SvMoE:MoEルータの教師あり学習

Mixture-of-Experts(MoE)は,エキスパートと呼ばれる複数のモジュールの出力をアンサンブルする技術である[1,2].アンサンブルには,ルータと呼ばれるモジュールの出力を重みとして使用する.深層学習[3]や言語モデリング[4]にも適用され,特にエキスパートの選択を疎にしたもの[4]はそのパラメタ数に対して計算コストが小さいという利点がある.しかし,MoEのルータによるエキスパートの選択ではその分布が偏る問題がある.近年の言語モデルにおける研究ではこれを避けるために,ルーティングを均一に近づけるためのLoad-Balancing損失(LB損失)を言語モデル損失に加えて訓練される[4,5,6].一方で,LB損失は言語モデルの性能に対して干渉するという問題もある[7].この問題を踏まえて本研究では,MoE全体の訓練の前にルータを教師あり学習で訓練する,SvMoE(SupervisedMixture-of-Experts)を提案する.各エキスパートの担当するドキュメントを事前に定義し,正規化されたTF-IDFを教師信号として使用することでエキスパートの学習した知識とルーティングを対応させる.ここでは訓練データのTF-IDFがそのトークンが得意なエキスパートに対応することを仮定している.つまり,(1)定義されたドキュメントでエキスパートをそれぞれ訓練し(2)それらをマージしてMoEモデルとして,(3)TF-IDFでルータを訓練したのちに(4)MoE全体を訓練することによりLB損失を使用せずに適切な学習を実現する.ケーススタディとして日本の法律へのドメイン特化を扱う.提案手法による日本の法律文を使用した継続事前学習を事前学習済みモデルに対して行い,司法試験のベンチマーク[8]で評価する.提案手法により訓練されたルータではTF-IDFに対応した選択が確認され,LB損失を用いない提案手法はLB損失を使用したベースラインとほぼ同等の結果を得た.