Mixture-of-Depths: Transformerで計算資源を動的に割り当てる技術

 (arxiv.org)
2 ポイント 投稿者 GN⁺ 2024-04-08 | 1件のコメント | WhatsAppで共有
  • Google DeepMindのMixture-of-Depths(MoD) は、Transformer言語モデルがすべてのトークンに同じFLOPsを使うのではなく、各レイヤーで必要なトークンだけをattentionとMLP計算に参加させる方式
  • レイヤーごとのルーターが各トークンに対してスカラー重みを生成し、あらかじめ定めたtop-k容量内のトークンだけがブロックを通過し、残りはresidual connectionで迂回する
  • kを事前に固定することで計算グラフとテンソルサイズは静的に保ちながら、文脈に応じたトークンごとの計算配分は動的に変化する
  • 実験ではMoDはisoFLOP基準で基本のTransformerより、同じ学習時間でより低いlossを出すか、同じ性能でforward passあたりのFLOPsを減らして学習・サンプリングstepを高速化できた
  • 最良の設定は12.5%容量ブロックを1ブロックおきに適用する構成で、auto-regressive samplingではpredictorベースのルーティングによってtop-kの非因果性問題を軽減した

MoDが狙う計算の無駄

  • 一般的なTransformer言語モデルはforward passですべてのトークンに同一の計算量を使う
  • MoDは、すべてのトークンとシーケンスが予測に同じ時間と労力を必要とするわけではないという前提のもと、全体の計算予算を抑えつつ必要なトークンに計算を配分する
  • 既存の条件付き計算(conditional computation)は動的な計算グラフを作る可能性があり、現在のハードウェアが好む静的計算グラフと合わないことがある
  • ユーザーが学習前に総計算予算を定めると、モデルはその範囲内でトークン別・レイヤー別にどこへ計算を使うかを学習する
    • 計算量、メモリ使用量、forward passあたりのFLOPsを事前に予測できる
    • どのトークンをルーティングするかは入力文脈によって変わる

ルーティング構造と計算削減の仕組み

  • MoDはMoE Transformerのようにルーターを使うが、複数のexpertから1つを選ぶ代わりに、トークンを2つの経路のどちらかへ送る
    • attentionとMLPを含む標準ブロック計算
    • 値をそのまま渡すresidual connection
  • residual connectionは計算コストが低く、ブロック出力は入力値で決まる
  • ブロックのcapacityを全トークン数Tより小さく設定すると、基本のTransformerよりforward passあたりのFLOPsが減る
    • たとえばcapacityをT/2にするとattentionのquery-key行列積はT²ではなく(T/2)²となり、FLOPsは25%水準になる
  • capacityを0にすると非常に高速だがTransformerパラメータの大半を使わないため性能が悪化し、capacityをTにすると基本のTransformerと同じになる
  • 目標は、この2つの極端な設定の間で、基本のTransformerと同等以上の性能を出しながらstep速度を高められる点を見つけること

top-k expert-choiceルーティング

  • 各レイヤーlでルーターが各トークン埋め込みに対してスカラー重みを生成する
  • ユーザーが定めたcapacity Cに応じて、シーケンス中のtop-kに当たるトークンだけがattentionとMLP計算に参加する
  • 選ばれなかったトークンはブロック計算をスキップし、residual connectionで次のレイヤーへ渡される
  • この方式はexpert-choice routingに基づいている
    • token-choice routingではトークンが経路を選ぶため、load balancingの問題が生じうる
    • expert-choice routingでは各経路がtop-kトークンを選ぶため、capacityが正確に満たされる
  • MoDでは経路が実質的に「ブロック計算」と「迂回」の2種類なので、単一のtop-k演算でトークンを2つの集合に分けられる
  • ルーター重みはブロック出力に掛けられてgradient path上に置かれ、言語モデリング学習の過程で勾配降下法の影響を受ける

サンプリングで生じる非因果性の問題

  • expert-choice routingのtop-k演算は非因果的(non-causal)
    • あるトークンがtop-kに入るかどうかは、その後のトークンのルーター重みまで分からないと決められない
    • auto-regressive samplingでは未来のトークン情報を使えない
  • これを避けるため、2つの方法を実験した
    • ルーター出力のsigmoidがtop-k選択の可否を予測するようにするbinary cross-entropy補助損失
    • ルーターと同じ入力を受けるがstop-gradientを適用した小さな補助MLP predictor
  • 1つ目の方法は主たる言語モデリングobjectiveに約**0.2~0.3%**の影響を与えた
  • 2つ目のpredictorベース手法は言語モデリングobjectiveに影響せず、step speedにも有意な影響を与えなかった
  • 補助予測課題は実験的にすぐ99%精度へ到達した

実験結果と観察

  • 初期実験は総計6e18 FLOPsの予算でMoDハイパーパラメータを探索する形で進められた
  • 最適なMoD設定は1ブロックおきにルーティングを適用し、top-kを256にした場合だった
    • シーケンス長2048基準で256トークン、すなわち**12.5%**だけがattentionとMLPを通過する
    • 残り1792トークン、すなわち**87.5%**はブロックを迂回する
  • 220MパラメータのMoD変種はisoFLOP最適baselineよりわずかに良い性能を示しつつ、学習中のstep速度は60%以上速く、同一ハードウェアでの学習全体のwall-clock timeはほぼ同じだった
  • 6e18、2e19、1e20 FLOPsのisoFLOP分析でも、MoDのFLOP-optimalモデルはbaselineより多くのパラメータを持つ傾向が見られた
  • 同等の性能またはより低いlossを出しながら、forward passあたりのFLOPsが少ないMoD変種が存在した
  • step速度の向上は2つの要因による
    • 一部トークンがブロックを迂回するため、同じモデルサイズでもFLOPs per parameterが減る
    • baselineと同等の性能を出せる、より小さなMoD変種を選べる
  • stochastic routingを使ったMoDは基本のTransformerや通常のMoDより大幅に低い性能を示し、学習されたルーティングの重要性が確認された
  • auto-regressive評価では256,000シーケンス、5億トークンのheld-out dataを使い、predictorベースのルーティングに切り替えても性能低下は小さかった
  • MoDはMoEと組み合わせてMixture-of-Depths-and-Experts(MoDE) として構成することもできる
    • staged MoDEはattention前にトークンをブロックへ送るか迂回させるかを決める
    • integrated MoDEは既存のMLP expert群の中に“no-op” expertを入れてMoDルーティングを統合する
  • ルーティング分析では、一部のトークンは複数ブロックを頻繁に通過し、他のトークンは可能な限り迂回するパターンが観察された
  • ブロックを頻繁に通るトークンは出力予測のentropyが高い場合と相関するという予備分析結果があり、これは予測がより難しいトークンである可能性を示している

関連記事

1件のコメント

 
GN⁺ 2024-04-08
Hacker Newsのコメント

  • より複雑なルーティングは確実に一般的になっていくと思う。
    特に、いずれは専門家集合を再び通過させる再帰的ルーティングへ進むと見ている。今後は「思考の連鎖（chain-of-thought）」がモデル内部で再帰的に起こるようになりそう。

    • こうした仮想的なオブジェクトは Recursive Neural Networks と呼べばよさそう。
    • ここで描写されているものは、Universal Transformers 系の研究に少し似て見える。入力埋め込みを1つのTransformerブロックに何度も通し、別のモジュールが埋め込みが十分に「熟した」かどうかを判断して取り出す方式だ。
      「専門家」というアイデアにより近いものとしては、昨年の Sparse Universal Transformers 論文があり、Universal Transformerと疎な専門家混合を組み合わせて、ゲーティング機構がどのTransformerブロックをどの順序で使うかを決める。
      自分の専門分野ではないが、きちんと学習させるのが難しく、通常のTransformerと同程度の結果を出すには推論時に総計算量がより必要になる、と理解している。それでも興味深い方向であり、トークンあたりの計算ステップ数に上限がある点は、古典的なTransformer構造の大きな欠点の1つだと思う。
    • これがまだ実現していない理由は、学習時点で再帰を何回行うべきかを決める方法がないからだと思う。
      ランダムな回数を選んだり、複数の再帰深度を試したりすると、出力が「ぼやける」。つまり、ある層の出力が最終結果に重要な情報を出すべきなのか、それとも次の再帰ラウンドに最適な入力を出すべきなのかが分からなくなる。
    • アテンションは基本的にルーティングであり、こうした別のルーティング方式はモデルにより粗い選択肢を与え、潜在的に学習を容易にし得る。
    • 流れが確実に動的ルーティングの増加に向かっているのは確かだが、MoE/MoD/MoDEがより深い推論を可能にするというより、重みの中で重複を減らし、追加の事実を格納できるようにするものに近いと思う。
      より深い推論は、層単位の動的性よりもトークン単位の動的性から来る気がする。たとえば、モデルが捨てる根拠トークンを出力する最近のQuiet-STaR論文がある: https://arxiv.org/abs/2403.09629

  • 2024年で最も重要な論文かもしれない。
    すべてのトークンに同じ量の計算を使わなくてもよいモデルが欲しい、というアイデアは以前からあったが、そのための説得力あるメカニズムは初めて見た。

    Equipped with these new methods, we can sample autoregressively by choosing to route tokens to or around a block based on the router’s output, which does not depend on any information from future tokens. We provide empirical evidence that this is a relatively easy auxiliary task that quickly achieves 99% accuracy.
    これ、少し驚きではないだろうか？

    • Sparse Universal Transformer の方が古く、すでにルーティングベースの早期終了を行っていた。
    • 最重要とまで言う？ すべてのトークンが全体のコンテキストウィンドウを必要とするわけではない、というのは obvious optimization であるべきだ。

  • 簡単な入門版:
    言葉を理解して処理できる賢いアシスタントがいると想像すればよい。通常、このアシスタントは各単語が全体の意味にどれほど重要かに関係なく、すべての単語に同じ注意を払う。
    そこで、そのアシスタントに「脳のリソース」をより賢く使わせる方法を見つけたとしよう。すべての単語に同じ注意を向ける代わりに、意味理解に最も重要な単語へより集中し、文脈に応じてその集中をその場で調整する。
    アシスタントが過負荷にならないよう、特定の時点で使える総「脳のリソース」にも上限を設ける。予算を与えて「一度に一定数の単語にだけリソースを使える」と伝えるようなものだ。するとアシスタントは、どの単語が最も重要かを決めなければならない。
    この上限があっても、アシスタントはリソース使用を柔軟に調整する。ある単語には多く使い、別の単語には少なく使いながら、状況に合わせて優先順位を付ける。
    このように学習すると、アシスタントは知的に注意を向けつつ、効率的に動作する。すべての単語に同じように注意を向けるアシスタントと同じくらいよく理解しながら、全体の脳リソースは少なくて済むため、応答や新情報の処理がはるかに速くなる。

    • ELI5だというのは分かるが、説明されている仕組みはすでにアテンションがやっていることではないのか？ 以前のシーケンス内で文脈上重要な単語に特定して集中するのだから。

  • 自分の理解をもとに、ここに少しまとめてみた:
    https://lifeinthesingularity.com/p/googles-breakthroughs-in-...

    • 良い記事だ。New Scientist っぽいスタイルを思い出す。「進めながら説明しつつ、短く保つ」感じで、この分野の感覚をつかむのに良い。

  • 専門家混合と非常によく似ている。ただし、トークンを複数の専門家へルーティングする代わりに、「動的にスキップできる単一の専門家に配置する」ような形だ。

    • 2つを混ぜるとかなり面白そうだ。性能を維持しながらMoEの計算量をさらに減らせるかもしれない。

  • “This is more computationally efficient than performing a full content-based lookup across an entire memory buffer for each step in the future, and could be one step towards drastically increasing the context-length available for making a prediction.”
    これは1,000万トークンのコンテキストウィンドウを可能にする方法なのか？ それとも今後のより長いコンテキストウィンドウの話なのか？

  • RASPアルゴリズムをいくつか理解して実装しようとして得た結論は、特定の関数が動作するには一定数のTransformer層が必要だということだった。
    この論理に従うと、Transformerが学習した関数が複数のヘッドにまたがって分散し得ることが分かる。こうした関数を反復することは、問題を理解して解くうえで非常に価値があり得るが、現在の推論方式では連続したヘッド群を繰り返し実行できない。この論文は実際に有望な方向に見える。
    [1] https://arxiv.org/pdf/2106.06981.pdf
    [2] https://www.youtube.com/watch?v=t5LjgczaS80

  • LLMがあまりにも速く進歩することの唯一の欠点は、少なくともGoogleの外では、誰も改善点を学習してテストできる速度よりも速く論文が出てくることかもしれない。
    自分がディープラーニングを始めた頃は、ReLUとドロップアウトが盛んで、コンシューマ向けの1080でコードを1、2行変えて数時間以内に改善したかどうかをテストできた。今はMistralのようなところが試してくれるまで数週間待たなければならなさそうだ。

    • GPU貧困層へようこそ。
      私は量子化アプローチに集中しつつ、古い前世代GPUでテストしている。

  • hu-poがAI論文を詳細なライブ配信でレビューしている。
    強くおすすめする。ここでは議論中の mixture-of-depths 論文を扱っている: https://www.youtube.com/watch?v=Teru_qIdB8Y