アーキテクチャ

Mixture-of-Depthsとは

MoD

各トークンごとに処理へ使う層の数を動的に決め、計算量を必要最小限に絞るTransformerの設計

ひとことで言うと

簡単な単語は少ない処理で済ませ、難しい単語だけじっくり計算する、計算配分を賢くする仕組み。

概要

Mixture-of-Depths(MoD)は、Transformerの全トークンへ同じ数の層を一律適用するのではなく、各トークンごとにその層をスキップするか処理するかを動的に判断するルーティング機構を導入し、モデル全体の計算量を必要最小限に絞るアーキテクチャ。全てのトークンが同じ計算コストを必要とするわけではなく、単純なトークンは少ない層の処理で十分な一方、文脈的に重要なトークンにはより多くの層の処理を割り当てる方が効率的だという発想に基づく。各層の入り口でルーターが各トークンに対しその層を通すかスキップするかを判定し、あらかじめ定めた計算予算(1層あたり処理するトークンの割合)の範囲内で計算資源を動的に配分する。 Mixture-of-Experts(MoE)が「どの専門家を使うか」を動的に選ぶのに対し、Mixture-of-Depthsは「どれだけの深さ(層数)を使うか」を動的に選ぶ点が特徴で、両者を組み合わせる研究も進んでいる。

背景

標準的なTransformerは全トークンに同じ数の層を適用するため、単純な処理で済むトークンにも一律に高い計算コストがかかっていた。 Mixture-of-Depthsは、トークンごとに必要な計算量が異なるという直感に基づき、計算資源をより効率的に配分する目的で考案された。

歴史

2024年: Google DeepMindのRaposoらが論文「Mixture-of-Depths: Dynamically allocating compute in transformer-based language models」を発表。

ワークフロー

各層の入り口で、軽量なルーターが各トークンの表現からその層を通すべきかスキップすべきかを判定する。あらかじめ定めた計算予算の範囲内で、重要度の高いトークンだけをその層で処理し、それ以外のトークンは残差接続を通じてそのまま次の層へ渡す。これを各層で繰り返すことで、トークンごとに実質的な処理の深さが変化する。

コード例

トークン単位のルーティング(概念的な実装例)

def transformer_block(x, router, block, capacity):
    scores = router(x)
    keep_idx = scores.topk(capacity).indices
    x[keep_idx] = block(x[keep_idx])
    return x

利点

全トークンに一律の計算コストをかける必要がなく、モデル全体の計算量を削減できる
重要なトークンにより多くの計算資源を割り当てられ、同じ計算予算でも精度を高めやすい
既存のTransformer構造に軽量なルーティング機構を追加するだけで導入できる

欠点

ルーターの判定精度がモデル全体の性能に影響し、学習の難易度が上がる
動的なルーティングにより、ハードウェア上での並列処理の効率がやや複雑になる
比較的新しい技術であり、大規模モデルでの実運用の知見がまだ少ない

比較

MoE — MoEが使う専門家サブネットワークを動的に選ぶのに対し、Mixture-of-Depthsは処理に使う層の数を動的に選ぶ
Transformer — Mixture-of-Depthsは標準的なTransformerに計算量を動的配分するルーティング機構を追加した設計
残差接続 — スキップされたトークンは残差接続を通じてそのまま次の層へ渡される

よくある質問

Mixture-of-DepthsとMixture-of-Expertsの違いは？

MoEは各トークンに対しどの専門家サブネットワークを使うかを動的に選ぶのに対し、Mixture-of-Depthsはどれだけの層(深さ)を使って処理するかを動的に選ぶ。両者は組み合わせて使うこともできる。

スキップされたトークンはどうなる？

その層の処理を受けず、残差接続を通じてそのままの表現が次の層へ渡される。重要度が低いと判定されたトークンの計算コストを節約できる。

参考文献

Research PaperMixture-of-Depths: Dynamically allocating compute in transformer-based language models