반복 트랜스포머는 역전파를 자르기보다 루프 수를 줄이는 편이 낫다

무엇인가

반복 트랜스포머(looped transformer)는 하나의 공유 블록을 여러 번 재사용해 학습 파라미터를 늘리지 않고 계산 깊이를 늘린다. 문제는 학습 방식이다. 이 모델은 BPTT(backpropagation through time)로 학습하는데, 계산을 아끼려고 마지막 몇 번의 루프에만 역전파를 흘리는 절단 BPTT가 흔히 쓰인다. 절단은 역방향 계산과 활성값 메모리를 아끼지만 편향된 그래디언트를 만들어 루프의 이득을 깎는다. 그런데 기존 스케일링 연구들은 절단 비율을 미리 정해둔 레시피를 비교할 뿐, 순전파 루프 수 T와 역전파 루프 수 B를 동시에 변수로 두고 손실을 모델링하지 않았다. 이 논문은 "고정된 학습 연산 예산에서 우리는 잘라야 하는가"를 정면으로 묻는다.

어떻게 동작하나

방법은 대규모 실험과 통합 스케일링 법칙이다. 아키텍처는 prelude–recurrent–coda 구조로, prelude가 토큰을 은닉 상태로 만들고, 공유 recurrent 블록이 [s0; s_{t-1}]을 입력 주입(input injection)으로 이어받아 T번 루프를 돌며, coda가 다음 토큰 분포를 낸다. 각 루프 끝에는 RMSNorm이 붙는다. 저자들은 58M~919M 비임베딩 파라미터의 5개 모델 크기, FineWeb-Edu 100B 토큰, T ∈ {1,2,4,8,16}에 걸친 15개 (T,B) 조합으로 총 592회 학습을 돌렸다(스케일링 스위트만 347개 구성, 파라미터당 1.1~1,291 토큰). 이 중 230개 구성으로 법칙을 적합하고 117개는 별도 검증용으로 남겼다. 법칙 자체는 기존의 결합형 크기–데이터 스케일링 법칙에 루프를 얹는 방식이다. 비임베딩 파라미터 N을 한 번만 도는 N_once와 매 루프 재사용되는 N_rec로 쪼갠 뒤, N_rec에 루프 계수 g(T,B)를 곱한 유효 파라미터 수로 N을 대체한다. g는 루프를 늘릴수록 체감 수익을 보이며 포화하고, 절단(B<T)은 이 이득에 할인을 건다.

무엇과 다른가

절단에 대한 결론은 조건부다. 적합된 지수는 순전파 루프 p≈0.237, 역전파 루프 q≈1.022로 q > 2p를 크게 만족한다. 역방향 루프가 순방향의 두 배 FLOPs를 쓰므로 제약이 T+2B=S 형태가 되는데, 이때 손실을 최소화하는 해는 T=B=S/3, 즉 전부 학습하는 full BPTT다. 반면 T가 이미 고정되어 있다면 이야기가 달라진다. 절단으로 아낀 역방향 연산을 더 큰 모델이나 더 많은 토큰에 쓰는 것이 할인된 루프 이득을 상쇄하고도 남는다는 기존 연구 결과를 법칙이 그대로 재현한다. 별도로 적합한 두 법칙을 비교해도 같은 방향이다. T+2B=12로 묶었을 때 (4,4)는 (8,2)와 같은 손실에 도달하는 데 8~14% 적은 연산을 쓴다. D=16·D0의 동일 연산 실험에서도 (4,4)가 모든 크기에서 (8,2)보다 낮은 손실을 냈다.

어떻게 쓰나

연산 배분 규칙도 도출한다. 모델 크기가 메모리 때문에 고정된 경우, 최적 루프 수는 C^0.24, 토큰 수는 C^0.76으로 늘어난다. 연산이 10배가 되면 루프는 약 1.7배, 토큰은 약 5.8배 늘어난다는 뜻이다. 참고로 Parcae가 보고한 경험적 지수는 루프 0.40, 토큰 0.78인데, 저자들은 Parcae가 절단 BPTT와 샘플링 기반 루프를 쓴 결과라 조건부 최적값이라고 지적한다. 모델 크기까지 자유롭게 둘 수 있으면 법칙은 루프를 쓰지 않는 쪽을 고른다. 루프 T개를 안 묶인 복사본 T개로 바꾸면 같은 연산·같은 토큰에서 유효 크기가 더 커지기 때문이다(g(T,T) < T). 이때 최적 배분은 N* ∝ C^0.56, D* ∝ C^0.44, T*=B*=1이다. 루프를 쓰는 설정 중에서는 (2,2)가 가장 좋은데, g(2,2)≈2라서 두 번의 공유 루프가 안 묶인 복사본 두 개에 거의 맞먹는 유효 크기를 주기 때문이다.

전제와 한계

이 (2,2) 설정은 실험으로도 확인된다. 세 가지 학습 예산에서 각 예산에 맞춰 손실을 최소화한 (2,2) 모델은 루프 없는 최적 모델 대비 0.008 낫 이내 차이를 유지하면서 비임베딩 파라미터가 27~32% 적었다. 2.1×10^19 FLOPs 지점의 쌍을 골라 (2,2) 모델을 8% 더 많은 토큰, 즉 8% 더 많은 연산으로 학습시키면 손실이 루프 없는 모델보다 0.007 낫 낮아지고, 파라미터는 여전히 27% 적다. 이 차이는 시드 간 편차보다 크다. 요약하면 같은 연산에서 (2,2)는 거의 동등하고, 연산을 조금 더 쓰면 앞선다.

개발자 관점의 실무 규칙은 이렇다. 절단 BPTT를 기본값으로 두고 있다면 먼저 순전파 루프 수가 정말 고정인지 확인해야 한다. 아키텍처 제약으로 T를 줄일 수 없다면 절단은 합리적이고, 아낀 연산을 모델 크기나 토큰에 재투자하는 편이 낫다. 반대로 T와 B를 함께 조정할 수 있다면 루프 수를 줄여서라도 모든 루프를 역전파하는 쪽이 같은 연산에서 이긴다. 이때 비교는 반드시 역방향 루프가 순방향의 두 배 FLOPs라는 점을 반영한 등연산 기준으로 해야 한다. 디바이스 메모리 때문에 파라미터를 늘릴 수 없는 상황이라면 추가 연산의 대부분은 토큰에 넣고 루프는 천천히 늘리는 것이 법칙이 말하는 배분이다. 파라미터 예산이 빠듯한데 품질을 크게 잃고 싶지 않다면 두 번의 완전 학습 루프가 현실적인 타협점이다.

저자들이 밝힌 한계는 분명하다. 첫째, dense 모델만 다뤘고 MoE에는 검증하지 않았으므로 절단과 연산 배분에 대한 결론이 루프 MoE로 확장되는지는 열린 문제다. 둘째, 검증 손실만 측정했고 다운스트림 태스크 성능이나 추론 능력으로 이어진다는 증거는 제시하지 않았다. 셋째, 실험은 T ≤ 16에 묶여 있어 이득의 상한과 이후의 더 깊은 루프로의 전환은 외삽이다. 넷째, prelude·coda 파라미터와 recurrent 블록 파라미터의 비율을 모델 크기마다 고정했으며, 이 비율이 루프 이득과 최적 루프 수에 어떤 영향을 주는지는 향후 과제로 남겼다. 깊은 재귀가 잔차 폭발이나 손실 스파이크를 일으키는 문제도 이 논문의 범위 밖이다.