DeepLoop: ループ型Transformerの深度スケーリング

要旨

ループ型トランスフォーマーは、物理ブロックのコンパクトなスタックを複数ラウンドにわたって適用することにより、格納パラメータを増やすことなく展開深度を増加させ、逐次計算を拡張する。この再利用により、残差スケーリング問題が変化する。すなわち、非結合型トランスフォーマーでは各残差枝が自身のパラメータ更新を受け取り適用するのに対し、ループ型トランスフォーマーでは一つの共有更新が繰り返し訪問からの勾配を集約し、次の線形化順伝播において同じ訪問によって読み戻される。我々はこの結合深度効果を、訪問整列係数κ_Rによって制御される一次摂動境界を通じて形式化する。この境界は、訪問が無相関化する場合にはDeepNorm指数を回復するが、保守的な整列領域では、固定された物理深度でループ数が増加するにつれて指数が1/4から1/2に増加することを必要とする。得られた手法であるDeepLoopは、Post-LN DeepNormアーキテクチャを維持し、展開深度Nに対してα=(2N)^{1/2}およびβ=(8N)^{-1/2}と設定する。GPT-2 smallおよびGPT-2 mediumスケールのGPTスタイルのループ型言語モデルにおいて、DeepLoopは物理ブロックが再訪問されない場合には中立的であり、再帰深度が活性化されると検証損失と下流精度を改善する。これらの結果は、安定な再帰深度には、名目上の層数だけでなくパラメータ訪問を考慮した残差スケーリング則が必要であることを示している。

English

Looped Transformers scale sequential computation by applying a compact stack of physical blocks for multiple rounds, increasing unrolled depth without increasing stored parameters. This reuse changes the residual-scaling problem: in an untied Transformer, each residual branch receives and applies its own parameter update, whereas in a looped Transformer one shared update aggregates gradients from repeated visits and is read back by those same visits in the next linearized forward pass. We formalize this tied-depth effect through a first-order perturbation bound controlled by a visit-alignment coefficient κ_R. The bound recovers the DeepNorm exponent when visits decorrelate, but in the conservative aligned regime it requires the exponent to increase from 1/4 to 1/2 as loop count grows at fixed physical depth. The resulting method, DeepLoop, keeps the Post-LN DeepNorm architecture and sets α=(2N)^{1/2} and β=(8N)^{-1/2} for unrolled depth N. On GPT-style looped language models at GPT-2 small and GPT-2 medium scale, DeepLoop is neutral when no physical block is revisited and improves validation loss and downstream accuracy once recurrent depth is activated. These results show that stable recurrent depth requires residual scaling rules that account for parameter visits, not only nominal layer count.