単一GPUによる大規模言語モデルの効率的事前学習

重み共有アーキテクチャと通常Transformerの比較実験

抄録

大規模言語モデル（LLM）の学習には通常、数十〜数百基のGPUクラスタが用いられるが、本研究では単一
GPU（NVIDIA RTX 3090, 24GB VRAM）のみを用いて、大規模言語モデルの効率的な事前学習が可能であること
を実証した。対象モデルは、我々が開発したBathysRDT（Bathys Recursive Deep Transformer）であり、重
み共有による再帰構造を持つ。学習対象となるユニークパラメータ数は3.29Bであるが、推論時には再帰ブロッ
クが12回適用されることで展開時7.02B規模の計算深度で動作する。同程度のユニークパラメータ数（3.19B）
を持つ通常のTransformerを同一条件で比較した結果、同一評価コードによるCross Entropy（CE）比較におい
て、BathysRDTがn=12（全深度）でCE=8.44を達成したのに対し、通常TransformerはCE=12.01となり、
BathysRDTはCEを約30%低減した。さらに、再帰深度を1回に制限した場合でも、通常Transformerとほぼ同
一の計算量条件（FLOPs/token）下でCE=10.02を達成し、Transformer（CE=12.01）を上回った。再帰深度12
ではCE=8.44まで改善し、計算量に応じた性能制御が可能であることを示した。学習データの反復に対する耐
性にも顕著な差が観測された。BathysRDTが200Mトークン時点で最良性能（CE=8.44, 蒸留検証PPL=4,612）
を記録したのに対し、Transformerは100Mトークン時点（CE=9.44, 蒸留検証PPL=12,604）をピークにその後
急速に劣化した。本研究は、VRAM制約下での大規模モデル学習における重み共有アーキテクチャの有効性を
示すとともに、同等計算量（FLOPs/token）での比較においても再帰的重み共有構造が通常Transformerに対
して優位であることを確認した。