토큰마다 디코더 상태를 되먹이는 순환 루프 트랜스포머가 긴 시퀀스를 일반화한다

무엇인가

이 논문이 다루는 문제는 알고리즘적 상태 추적이다. 패리티는 비트를 하나씩 누적해야 하고, 순열 추적은 군 원소를 계속 합성해야 한다. 그런데 트랜스포머는 시퀀스 길이와 무관하게 토큰마다 고정된 층 수를 적용한다. 저자들은 TC0 = NC1이 아닌 한 고정 깊이 트랜스포머가 임의로 긴 시퀀스에서 S5 순열의 합성을 추적할 수 없다는 기존 결과(Merrill et al., 2024)를 근거로 든다. 순환 신경망은 각 업데이트의 결과를 다음으로 넘겨 계산 경로가 입력 길이에 따라 자라지만, 트랜스포머에는 그 연결이 없다는 것이 출발점이다.

어떻게 동작하나

제안 방법인 Recurrent Looped Transformer(RLT)는 8개 층을 병렬 인과 인코더와 순환 트랜스포머 디코더로 나눈다. 인코더는 알려진 토큰을 병렬로 처리해 토큰 표현 e_t와 전역 키-값 메모리 M≤t를 만든다. 디코더는 매 위치에서 게이트 병합으로 현재 인코더 표현과 직전 토큰의 최종 디코더 출력 s_{t-1}을 결합한다. 병합식은 u_t = e_t + α·g_t⊙W_s·RMSNorm(s_{t-1}), g_t = σ(W_g[e_t; r_{t-1}] + b_g) 형태이고, 실험에서는 피드백 스케일 α=0.1을 쓴다. 디코더 각 블록은 인과 슬라이딩 윈도 어텐션, 인코더 메모리에 대한 크로스 어텐션, FFN을 차례로 적용한다. 디코더 상태는 H_t = (s_t, C_t^D)로, C_t^D는 각 SWA 층이 보관하는 키/값 캐시다. 디코더 깊이가 L_D일 때 t개 토큰을 처리하면 t·L_D개의 디코더 블록을 지나는 순환 경로가 생기지만, 토큰 하나가 통과하는 층은 L_E + L_D개로 고정된다. 프롬프트와 응답 토큰은 같은 전이를 따른다.

무엇과 다른가

설계 변수는 두 가지다. 하나는 인코더/디코더 깊이 배분이고, 다른 하나는 피드백 간격 B다. RLT-1은 매 토큰마다(B=1) 피드백하고, RLT-2는 B개 토큰 청크마다 한 번 상태를 갱신하며, RLT-0은 피드백을 아예 제거한다(B=∞). 길이 T인 프롬프트에 대해 디코더가 필요한 순차 블록 단계는 RLT-1이 T·L_D, RLT-2가 ⌈T/B⌉·L_D, RLT-0이 L_D다. 세 변형 모두 T·L_D개의 디코더 블록을 평가하고 토큰 하나를 스텝마다 생성하므로, 큰 B는 피드백 빈도를 알려진 토큰의 병렬성과 맞바꾼다. 파라미터가 B에 의존하지 않으므로 청크 크기를 학습 중에 바꿀 수 있고, 저자들은 큰 B로 사전학습한 뒤 중간·사후 학습에서 B=1까지 줄이는 스케줄을 제안한다(다만 이 스케줄 자체는 실험하지 않았다). seed-42 CPU 타이밍에서 4+4 구성의 chunk4 학습 스텝은 RLT-1보다 2.27배, RLT-0은 4.17~4.30배 빨랐다.

어떻게 쓰나

실험은 8개 논리 층, 폭 512, FFN 폭 1,365, 어텐션 헤드 4개로 고정하고, 4+4·5+3·6+2·7+1·8+0 다섯 가지 RLT-1 분할을 8층 디코더 전용 트랜스포머와 비교한다. 시드 42·43·44 세 개, 전역 배치 512, AdamW(β1=0.9, β2=0.95, weight decay 0.1), 학습률 10^-4에서 200스텝 워밍업 후 5×10^-6까지 감쇠한다. 패리티는 3~40비트, S5는 32개 연산, 평면 mod-5는 홀수 길이 3~39, 괄호식은 길이 3~40으로 학습한다. 패리티 256비트에서 5+3과 7+1 분할이 모든 시드에서 100±0%를 기록한 반면 트랜스포머는 50.07±1.63%로 확률 수준에 머물렀다. 학습 중에도 500스텝에서 6+2가 검증 정확도 99.44±0.98%를 내고 트랜스포머는 48.48±0.53%였다. 별도의 16층 시리즈에서도 8+8, 9+7, 11+5, 16+0이 256비트에서 100%를 기록했고 트랜스포머 16은 49.41%였다. 학습 길이의 8배인 256개 연산의 swap 기반 S5 추적에서 4+4는 최종 상태 정확도 97.30±2.76%를 냈고 트랜스포머는 0.85±0.30%였다. 512개 연산에서도 4+4가 55.70±25.78%를 유지한 반면 트랜스포머는 0.85±0.30%에 그쳤다. 5,000스텝 후 모듈러 산술에서는 길이 63 평면식에서 6+2가 93.36±5.69%(트랜스포머 33.20±2.33%), 길이 64 괄호식에서 5+3이 67.97±2.91%(트랜스포머 46.71±1.21%)를 기록했다.

전제와 한계

깊이 배분은 과제에 따라 최적점이 다르다. 패리티는 디코더 층이 1개 또는 3개일 때 가장 잘 일반화되지만, swap 기반 S5는 256개 연산에서 디코더 깊이가 늘수록 정확도가 올라 1개 층에서는 확률 수준이던 것이 4개 층에서 97%가 된다. 피드백 간격 절제 실험도 분명한 결과를 준다. 피드백을 제거한 RLT-0은 모든 분할에서 64비트 패리티가 50.23±3.80%로 확률 수준이고 swap 기반 S5도 마찬가지다. 4토큰 청크를 쓰는 RLT-2 chunk4는 64비트 패리티를 98.99±1.66%로 유지했지만, 64개 연산 swap 기반 S5는 RLT-1의 100±0%에서 19.60±6.54%로 떨어졌다. 즉 순열 추적은 토큰 단위 피드백에 크게 의존한다. 청킹 효과는 분할에도 의존해서, chunk4 4+4는 256비트 패리티에서 69.34±24.17%를 유지한 반면 chunk4 7+1과 8+0은 64비트에서 이미 확률 수준이었다. swap 기반 S5에서는 chunk4 5+3이 48개 연산 79.92±10.89%, 64개 연산 50.16±18.77%로 chunk4 4+4보다 나았다.

범위 밖의 결과도 분명히 밝힌다. 덧셈은 모든 모델이 학습 폭 검증셋에서 100% 교사 강제 정확도를 내지만 8자리를 넘으면 떨어지고, 32자리에서 8층 모델 평균은 14.89~16.84%다. 표준 S5(swap이 아닌 전체 순열)는 여전히 어려워서 모든 피드백 변형이 1/120 균등 기준 근처에 머문다. 저자들은 이 실험이 지도 학습 기반 알고리즘 성능만 측정하며 강화학습 성능은 평가하지 않았다고 명시한다. 또한 긴 표현식에서 정확도가 떨어지고, 여러 평면 mod-5 분할이 시드에 따라 변동이 크다는 점, 청크 크기 스케줄을 실제로 검증하지 않고 각 청크 크기를 처음부터 학습했다는 점을 한계로 남긴다.

개발자 관점에서 이 논문의 실용적 신호는 두 가지다. 첫째, 긴 시퀀스에서 상태를 누적해야 하는 작업(상태 추적, 반복 합성, 긴 산술)에서 고정 깊이 트랜스포머가 학습 길이를 넘는 순간 무너지는 현상을 순환 피드백으로 완화할 수 있다는 구체적 수치를 제공한다. 둘째, 피드백 간격 B가 정확도와 병렬성 사이의 실질적 손잡이라는 점이다. 청크를 키우면 학습 스텝이 2~4배 빨라지지만 순열 추적처럼 토큰 단위 갱신이 중요한 과제에서는 정확도가 급락하므로, 과제 특성에 따라 B를 정하고 추론·RL 재생 시에는 최종 학습 단계의 청크 크기를 그대로 써야 한다. 파라미터 수가 동일 층수 기준으로 정확히 일치하지 않는다는 점(RLT-1 26.10~28.73M 대 트랜스포머 25.31M, RLT-0은 RLT-1보다 787,968개 적음)도 비교 해석 시 감안해야 한다.