効率的なクロススケール・ハイパーパラメータ転送による確率的トランスフォーマーのスケーリング

Abstract

文脈における単語表現のためのホワイトボックス確率モデルである確率的トランスフォーマー（Probabilistic Transformer; PT）は、小規模モデルおよび小〜中規模データセットにおいて、計算構造と下流タスク性能の両面で標準的なトランスフォーマーと大きな類似性を示してきました。とはいえ、PTは標準的なトランスフォーマーよりもハイパーパラメータの選択に対する頑健性が低く、効率的にスケールすることが難しくなります。本研究では、最大更新パラメータ化（Maximal Update Parametrization; muP）に従ってPTのパラメータを再スケーリングし、小規模モデルで最適化されたハイパーパラメータを追加のチューニングなしに大規模モデルへ転移できるようにします。このアプローチにより、PTを最大0.4B（4,000万）パラメータのモデルまで正常にスケールすることができました。実験の結果、Masked Language Modeling（MLM）タスクにおいて、同一のパラメータ予算のもとでPTは標準的なトランスフォーマーを一貫して上回ることが示されました。本研究が、将来に向けて、より大規模における確率モデルの実用的な導入に貢献することを期待しています。