Recurrent Looped Transformer, 시퀀스 길이에 따라 계산 경로 확장

배경

Transformer는 시퀀스를 병렬로 처리할 수 있지만, 각 토큰에 적용되는 층 수는 일반적으로 고정되어 있습니다. 패리티 계산, 상태 머신 시뮬레이션, 순열 추적처럼 입력을 읽을 때마다 내부 상태를 갱신해야 하는 과제에서는 이런 고정 깊이가 긴 시퀀스로의 일반화를 제한할 수 있습니다.

이 논문은 고정된 층 수를 단순히 늘리는 대신, 시퀀스를 따라 상태를 반복적으로 전달하는 Recurrent Looped Transformer(RLT) 를 제안합니다.

작동 방식

RLT는 8개 층을 병렬 인과 인코더와 순환 디코더로 나눕니다. 인코더는 현재 토큰과 이전 문맥을 이용해 표현을 만들고, 디코더는 이 출력과 직전 토큰에서 얻은 최종 디코더 상태를 결합합니다.

그 결과 인코더의 병렬 처리 장점을 유지하면서 상태는 토큰 순서에 따라 다음 단계로 전달됩니다. 시퀀스가 길어질수록 상태가 갱신되는 횟수는 늘어나지만, 토큰 하나에 필요한 국소적 계산 비용은 비교적 일정하게 유지됩니다.

연구진은 8개 층을 인코더와 디코더에 어떻게 배분하는지 바꾸며, 병렬 표현 능력과 상태 추적 능력 사이의 균형도 분석했습니다.

주요 결과

6개 알고리즘 과제에서 5가지 층 분할을 평가하고, 8층 표준 Transformer와 비교했습니다. 모든 비교는 세 개의 랜덤 시드에서 진행됐습니다.

• 패리티: 최대 40비트로 학습했지만, 두 RLT 구성은 256비트까지 일반화해 모든 시드에서 100% 정확도를 기록했습니다. 표준 Transformer는 거의 무작위 수준에 머물렀습니다.
• S_5 순열 추적: 학습 길이의 8배인 평가에서 RLT는 최종 상태 정확도 97%를 달성했지만, 표준 Transformer는 1% 미만이었습니다. 디코더 층이 많을수록 성능도 높아졌습니다.
• 모듈러 산술: 학습 길이를 넘어선 평가에서 RLT는 최대 93%, 표준 Transformer는 33%를 기록했습니다.
• 피드백 제거 실험: 순환 피드백을 제거하면 패리티와 swap-based S_5 성능이 모든 층 분할에서 무작위 수준으로 떨어졌습니다.

병렬성과 상태 정확도의 절충

피드백을 네 토큰마다 한 번씩 갱신하는 청크 방식도 실험했습니다. 청크 안의 토큰을 병렬로 처리할 수 있어 64비트 패리티에서는 99% 정확도를 유지했습니다.

그러나 순열 추적에서는 성능 저하가 컸습니다. 길이 64의 swap-based S_5에서 정확도는 100%에서 20%로 하락했습니다. 즉 패리티처럼 집계 중심인 계산은 지연된 업데이트를 견딜 수 있지만, 순서에 민감한 상태 전이에는 토큰별 피드백이 필요합니다.

의미와 한계

RLT는 Transformer의 층을 무작정 늘리는 대신, 깊이의 일부를 시퀀스 방향의 순환 시간으로 바꾸는 방법을 보여줍니다. 길이 외삽에는 모델 규모뿐 아니라 입력 진행에 맞춰 상태를 적절한 단위로 업데이트하는 구조도 중요하다는 점을 시사합니다.

다만 이번 근거는 알고리즘 벤치마크에 집중되어 있어 일반적인 언어 모델링 성능 향상을 직접 입증하지는 않습니다. 병렬 실행과 정밀한 상태 추적 중 무엇을 우선할지는 향후 시스템 설계에서도 중요한 문제로 남습니다.

출처: Hugging Face Daily Papers