Truncate or Not? A Scaling Study of Looped Transformers with Truncated Backpropagation

Looped transformers reuse a shared block of layers to increase computational depth without increasing the number of trainable parameters. They are trained with backpropagation through time (BPTT). To save training compute, BPTT is often truncated to the last few loops, but truncation yields biased gradients that may weaken the benefit of looping. In this work, we study the value of truncation in looped transformer training in a systematic way: we jointly vary model size, training tokens, and forward and backward loop counts across nearly 600 training runs and fit a unified scaling law that captures the effects of looping and truncation through effective parameters. Despite the training-compute savings of truncation, our fitted law favors full BPTT when forward and backward loop counts are jointly optimized. In matched-compute runs, full BPTT through fewer loops consistently outperforms truncated settings of the same compute. The law further shows that at fixed model size, the compute-optimal loop count and training-token budget both grow with compute. When model size is also free, the law identifies two fully trained loops as the best looped setting. Experimentally, this setting achieves lower validation loss than the compute-optimal non-looped model with about 27% fewer non-embedding parameters and only 8% more training compute.

한 줄 요약

Looped Transformer에서 Truncated BPTT의 효과를 규명한 대규모 스케일링 연구.

핵심 기여도

**Unified Scaling Law**: Looping과 Truncation의 효과를 Effective Parameter로 통합 모델링.
**Full BPTT의 우수성**: 동일 컴퓨트에서 Full BPTT가 Truncated BPTT보다 성능 향상.
**2개 Full Loop 모델**: 27% 적은 non-embedding 파라미터로 8% 더 많은 컴퓨트만으로 non-looped 모델 성능 초과.

핵심 아이디어

Looped Transformer는 동일 블록을 반복 실행하여 파라미터 수를 유지하면서 계산 깊이를 늘리는 구조다. 그러나 BPTT를 Truncate하면 경사도 편향이 발생해 성능 저하가 우려된다. 본 연구는 600회 이상의 실험을 통해 Forward/Backward Loop Count, 모델 크기, 훈련 토큰 수를 **동시에 조절**하며, **Effective Parameter** 개념을 도입한 통합 스케일링 법칙을 제안한다. 이 법칙은 Truncation이 Loop Gain을 감소시키는 반면, Full BPTT는 Loop Count를 줄이고 Backward Loop를 늘리는 방식으로 더 높은 Effective Parameter를 제공함을 보여준다.

기술적 접근법

**모델**: Looped Transformer, 파라미터 수 58M–919M 범위.
**훈련**: Forward/Backward Loop Count를 독립적으로 조절, 최대 16개 Loop.
**알고리즘**: BPTT (Full/Truncated) 비교.
**데이터**: 토큰 수는 파라미터당 3자리 수 범위.
**하이퍼파라미터**: Loop Count, 모델 크기, 토큰 수를 독립적으로 조절하며 실험.
**법칙**: Effective Parameter 기반의 통합 스케일링 법칙을 적합.

주요 결과

**Full BPTT vs. Truncated BPTT**: 동일 컴퓨트에서 Full BPTT가 성능 우수.
**2개 Full Loop 모델**: Non-looped 모델 대비 27% 적은 non-embedding 파라미터로, 8% 더 많은 컴퓨트만으로 validation loss 개선.
**Compute-Optimal Loop Count**: 고정된 모델 크기에서 Loop Count는 컴퓨트에 따라 C^0.24 비율로 증가.
**모델 크기 자유일 때**: 2개 Full Loop가 최적, Non-looped 모델과 유사한 성능.

의의 및 한계

**의의**: Looping과 Truncation의 효과를 정량적으로 분석한 최초의 연구. Effective Parameter 개념을 통해 다양한 설정을 통합 모델링.
**한계**: 실험은 주로 Language Modeling에 초점, 다른 태스크 적용 가능성 미확. 또한, 실제 하드웨어 제약(메모리, 병렬성)은 고려되지 않음.

실용적 활용

**자원 효율적인 모델 설계**: 파라미터 수를 줄이면서도 성능 유지하는 모델 개발에 활용 가능.
**대규모 모델 훈련**: 컴퓨트 예산이 제한된 상황에서 Loop Count와 토큰 수를 최적화할 수 있는 기준 제공.
**AI 연구**: 스케일링 법칙 기반의 모델 설계 전략 수립에 기초 자료로 활용 가능.