[논문리뷰] Truncate or Not? A Scaling Study of Looped Transformers with Truncated Backpropagation

링크: 논문 PDF로 바로 열기

저자: Xuyuan Xiong, Siyu Chen, Zihan Shao, Yue Su, Weiyang Jin, Ping Luo, Zhuoran Yang, Xihui Liu

1. Key Terms & Definitions

본 논문에서 다루는 핵심 용어 및 개념은 다음과 같다.

• Looped Transformers: 공유된 레이어 블록을 반복적으로 사용하여 학습 가능한 파라미터 수를 늘리지 않고 computational depth를 증가시키는 Transformer 모델 아키텍처이다.
• Truncated Backpropagation Through Time (TBT): Recurrent neural network 학습 시 gradient가 전파되는 recurrent steps의 수를 제한하는 기법으로, training compute와 activation memory를 절약하지만 biased gradients를 유발할 수 있다.
• Full BPTT: Recurrent neural network 학습 시 모든 recurrent steps (loops)에 걸쳐 gradient를 전파하는 방식으로, TBT와 대비되는 개념이다.
• Effective Parameters: Looped transformer가 동일한 token budget에서 동일한 loss를 달성하는 데 필요한 non-looped model의 파라미터 수를 의미하는 지표로, looping의 이점과 truncation의 손실을 반영하여 모델의 "실질적인" 크기를 정량화한다.
• Looping Coefficient g(T,B): Recurrent block의 parameter count에 곱해져 effective parameters를 결정하는 계수로, T (forward loops) 및 B (backward loops)에 따라 looping gain과 truncation discount를 포착한다.

2. Motivation & Problem Statement

Looped transformers는 뛰어난 성능과 파라미터 효율성을 동시에 제공하여 주목받고 있지만, 학습 시 Truncated Backpropagation Through Time (TBT) 사용 여부에 대한 최적의 전략이 명확하지 않다. TBT는 backward computation과 activation memory를 절약하지만, biased gradients를 생성하여 looping의 이점을 약화시킬 수 있는 tradeoff를 수반한다. 기존 scaling studies는 prescribed backpropagation recipes만을 검토하거나 forward loop count (T)와 backward loop count (B)를 jointly modeling하지 않아, fixed training compute budget 하에서 forward loops, backward loops, model size, 그리고 training tokens 사이에 compute를 어떻게 분배해야 할지에 대한 통합적인 지침을 제공하지 못했다. 따라서 본 연구는 model size, training tokens, forward loops, backward loops가 loss에 미치는 영향을 체계적으로 분석하고, fixed training compute budget에서 backpropagation을 truncate할지 여부를 결정하는 unified scaling law를 제시하고자 한다.

3. Method & Key Results

본 논문은 looped transformers의 loss와 model size, training tokens, 그리고 forward 및 backward loop counts 간의 관계를 설명하는 unified scaling law를 제안한다. 이 scaling law는 effective parameter count (N_eff) 개념을 도입하여 looping의 benefit과 truncation의 discount를 통합적으로 모델링한다. N_eff는 N_once + g(T,B) * N_rec로 정의되며, 여기서 g(T,B)는 exp[c * (1 - T^(-p) * B^(-q))] 형태의 looping coefficient이다.

주요 실험 결과 및 핵심 내용은 다음과 같다.

• Looping Coefficient의 특성: full BPTT (B=T)에서 looping coefficient g(T,T)는 T가 증가함에 따라 finite ceiling으로 수렴하는 diminishing returns를 보였으며 [cite: 1, Figure 9a], Truncated BPTT는 이 gain을 discount했다 [cite: 1, Figure 9b]. g(T,B)는 B/T 비율이 클수록 높게 나타나, trained fraction이 중요함을 시사한다 [cite: 1, Figure 9b]. 특히, effective parameters 증가는 forward loop count (T)보다 backward loop count (B)에 의해 더 크게 좌우되는 것으로 분석되었으며 [cite: 1, Figure 9c], fitted exponents인 q (~1.022)가 p (~0.237)보다 현저히 큰 값으로 이를 정량적으로 뒷받침한다.
• Truncation의 Optimal 여부: forward loop count (T)가 prescribed된 경우, truncation은 backward computation을 절약하여 더 큰 model size와 더 많은 training tokens를 허용함으로써 loss를 낮출 수 있다. 예를 들어, T=8에서 B=2로 truncation할 경우, full BPTT 대비 equal optimized loss에서 34%의 compute saving을 예측한다 [cite: 1, Table 5]. 그러나 forward 및 backward loop counts가 모두 jointly optimized될 수 있는 상황에서는 Full BPTT가 compute-optimal한 전략으로 나타났다. T+2B가 동일하게 fixed된 compute budget 하에서 Full BPTT (T=4, B=4)가 Truncated BPTT (T=8, B=2)보다 더 낮은 validation loss를 달성했다 [cite: 1, Table 6, Table 7].
• Compute Allocation 규칙: model size가 fixed된 경우, compute-optimal한 allocation은 대부분의 추가 compute를 training tokens에 할당하고 loop count는 C^0.24로 slowly 증가시키는 것을 제안한다. 반면 model size가 free인 경우, no looping ((1,1) 설정)이 computationally optimal하다고 예측되었다. 하지만, 두 개의 fully trained loops ((2,2) 설정)는 parameter-efficient한 대안으로, compute-optimal non-looped model과 유사한 loss를 달성하면서도 약 **27-32% 더 적은 non-embedding parameters**를 사용했다 [cite: 1, Table 12]. 또한, (2,2) 모델은 **8% 더 많은 training compute**를 사용하여 non-looped model보다 더 낮은 validation loss를 달성했다 [cite: 1, Figure 14].

4. Conclusion & Impact

본 연구는 looped transformers의 truncation 전략과 compute allocation에 대한 심층적인 scaling study를 수행했다. effective parameters를 기반으로 한 unified scaling law를 통해 looping gain과 truncation discount를 정량화했으며, forward loop count가 prescribed된 경우 truncation이 compute saving에 도움이 될 수 있지만, forward 및 backward loop counts가 모두 tunable할 때는 Full BPTT가 optimal하다는 중요한 결론을 도출했다. Compute allocation 측면에서는 fixed model size일 때 optimal loop count는 slowly 증가하고 training tokens가 compute의 대부분을 가져가며, free model size일 때는 no looping이 optimal하지만, **two fully trained loops**는 약 **27-32% 더 적은 non-embedding parameters**로 compute-optimal non-looped model에 필적하는 성능을 제공하며, 약간의 추가 compute로 더 낮은 loss를 달성하는 parameter-efficient training의 가능성을 입증했다 [cite: 1, Figure 14]. 이 연구는 Looped Transformers의 design 및 training에 대한 실질적인 guidelines를 제공하여 parameter-efficient하고 scalable한 language models 개발에 기여할 것으로 예상된다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.