Mixture-of-Depths（動的計算配分）

なぜ「トークンごとに深さを変える」のか

標準的な Transformer は、すべてのトークンをすべての層に一律で通します。文の途中の「、」も、意味の要となる名詞も、まったく同じ回数だけ attention と MLP を潜るわけです。しかし直感的には、必要な計算量はトークンごとに違うはずです。冗長な機能語に深い層をフルに割くのは、計算予算の無駄遣いに見えます。

Mixture-of-Depths（MoD） は、この一律処理をやめる仕組みです。各ブロックに小さなルータを置き、そのブロックで実際に処理する少数のトークンだけを選抜します。選ばれなかったトークンは、層の計算を丸ごとスキップして残差接続だけで次へ流れます。結果として、トークンごとに「実際に通った層の数」＝実効的な深さが変わります。これが「動的計算配分（dynamic compute allocation）」と呼ばれる所以です。

容量制約：総 FLOPs を静的に固定する

MoD の設計で最も重要なのは、1ブロックが処理できるトークン数に固定の上限（容量）を設ける点です。系列長を T、容量比を c（例 0.125）とすると、そのブロックが処理するのは上位 c × T 個のトークンだけと決め打ちします。

容量 C = c × T          # 例: c=0.125 なら系列の 12.5% のトークンだけ処理

そのブロックで:
  選抜された C 個   ──▶ [self-attention + MLP] を通る
  残りの T − C 個   ──▶ 残差接続だけで素通り（層をスキップ）

ここが MoD の肝です。容量 C は入力内容に依存しない固定値なので、モデル全体の総 FLOPs は系列長さえ決まれば静的に確定します。「動的」なのはあくまでどのトークンを通すかという中身であって、計算量そのものは事前に分かる。これにより計算グラフのサイズが一定になり、行列積の形も固定できるため、ハードウェア上で効率よく実行できます。

ルータの動作：スカラースコアと Top-k 選抜

ルータの実体は、MoE のルータ と同じく小さな線形層1枚です。ただし出力が違います。MoE のルータが「どのエキスパートか」を選ぶために E 個のロジットを出すのに対し、MoD のルータは各トークンに対して**スカラー値1つ（ルータ重み）**を出すだけです。

r_t = x_t · w_router        # w_router は [隠れ次元 × 1]、トークン t のスカラースコア

このブロック内の全 T トークンについて r_t を並べ、
  top_idx = argtopk(r, C)   # 上位 C 個のトークン番号を選抜

選ばれた C 個のトークンは attention と MLP を通り、その出力にルータ重み r_t を掛けてから残差に加算します。

選ばれた t:   y_t = x_t + r_t · Block(x_t)      # r_t 経由で勾配が流れる
選ばれない t: y_t = x_t                          # ブロックは素通り

MoE と同じ論点がここにも現れます。「上位 C 個を選ぶ」argtopk は微分不可能な離散操作です。それでもルータが学習できるのは、選ばれたトークンの出力に掛かる r_t を経由して勾配が流れるからです。あるトークンをブロックに通すことが損失を下げるなら、その r_t を大きくする（＝スコアを上げる）向きに w_router が更新されます。選抜は離散だが、選抜後の連続的な重み付けが学習信号になる――MoE と共通の設計思想です。

因果性の問題：推論時にどう Top-k を作るか

Top-k 選抜には見落としやすい欠陥があります。上位 C 個を決めるには系列内の全トークンのスコアを見て順位付けする必要があるため、これは本質的に**非因果的（non-causal）**な操作です。学習時は系列全体が手元にあるので問題ありませんが、自己回帰生成では未来のトークンがまだ存在せず、あるトークンが上位に入るかを今の時点で判定できません。

原論文はこれを2通りで解決します。

どちらも狙いは同じで、「系列全体を見ないと決まらない Top-k」を、各トークン単体で因果的に近似できるようにすることです。これで推論時も、過去だけを見てそのトークンをブロックに通すか判断できます。学習時の理想的な Top-k と、推論時の因果的な近似とのズレが MoD 固有の注意点になります。

Mixture-of-Experts との違い：疎化する「向き」

MoD と MoE は「疎な活性化で計算を節約する」点で似ていますが、疎にする軸が直交しています。ここを取り違えると設計を誤ります。

要点は、MoE では全トークンが必ずどこかの FFN を通る――疎なのは「どの重みを使うか」であって、層を飛ばすわけではない、ということです。対して MoD は層そのものをスキップします。MoE が「同じ計算量で容量（総パラメータ）を増やす」道具なのに対し、MoD は「同じ深さ・同じ品質を、より少ない計算量で実現する」道具、と役割が対照的です。

実装上の勘所と効果

MoD ブロックは通常、全層ではなく1層おきに配置します。連続でスキップを許すと情報が伝わらなくなりうるため、密なブロックと MoD ブロックを交互に挟む構成が安定します。容量比 c は品質と計算量の直接のつまみで、小さくするほど FLOPs は減りますが、飛ばされるトークンが増えて品質が落ちるリスクが上がります。

効果の出方は独特です。MoD は同じ FLOPs 予算のもとで層数を増やせる（各層が一部トークンしか処理しないぶん、浮いた計算で深さを稼げる）ため、同一計算量で密モデルを上回る、あるいは同等品質を少ない FLOPs・速いステップで達成する、という形で効きます。とくに1回の前方計算の FLOPs を大きく削れるので、KV キャッシュと推論最適化 で扱う推論コストの観点でも利点があります。

まとめ：計算予算を「均等」から「必要な所へ」

MoD を一言でいえば、固定した計算予算を、トークンごとの必要性に応じて配分し直す構造です。全トークンを全層に通す一律処理をやめ、各層で本当に処理すべき少数トークンだけを選ぶ。総計算量は容量制約で据え置いたまま、その使い道を賢くする点が本質です。

この見立てがあると、「パラメータも深さも同じなのに速い」「一部トークンだけ深く処理される」といった MoD 特有の挙動が、容量制約・深さ方向の疎性・因果近似の言葉で説明できます。LLM と Transformer の各ブロックを「全トークンではなく選ばれたトークンだけ通す」と置き換えるだけで、動的計算配分の効きどころが線でつながります。