Consistency LLM: LLMを並列デコーダに変えて推論速度を最大3.5倍向上

 (hao-ai-lab.github.io)
2 ポイント 投稿者 GN⁺ 2024-05-09 | 1件のコメント | WhatsAppで共有
  • 長い応答を生成する際、トークンを1つずつ生成する自己回帰(AR)デコーディングが遅延の主要なボトルネックであり、CLLMはこれをn個のトークン単位の並列デコーディングに減らそうとするアプローチ
  • Consistency Large Language Models(CLLMs) は事前学習済みLLMをファインチューニングし、任意のnトークン状態をAR greedyの結果と同じ固定点へ高速に写像するよう学習する
  • Jacobi decodingは理論上AR greedy生成と同じ結果に収束するが、既存のLLMでは平均1.05倍程度の限定的な速度向上にとどまり、実用性は低かった
  • CLLMはSpider、Human-Eval、GSM8k、MT-benchの実験で2.4倍〜3.4倍の生成速度向上を示し、Medusa2、Eagleのような高速推論手法と同等以上と評価された
  • 推論時に補助モデルコンポーネントや対象モデルのアーキテクチャ変更が不要なため、速度改善とメモリ効率を両立できる

CLLMが狙うデコーディングのボトルネック

  • LLMはプログラミング、法務、健康アドバイスなど多様な分野で使われているが、推論段階では通常自己回帰(AR)デコーディングでトークンを1つずつ生成する
  • 応答が長くなるほどトークン単位の生成方式のため遅延が大きくなり、ユーザーが体感する待ち時間も増える
  • 既存の高速推論手法は、複数トークンを一度に生成するためにアーキテクチャの改変、補助コンポーネント、ドラフトモデルを必要とする場合が多い

Jacobi decodingの動作方式と限界

  • Jacobi decodingはJacobiおよびGauss-Seidelの固定点反復に着想を得た方式で、greedyデコーディングではAR生成と同一であることが証明されている
  • 逐次生成をn個の変数とn個の非線形方程式からなるシステムに置き換え、Jacobi反復で並列計算できるようにする
  • 具体的な流れは以下の通り
    • 入力プロンプトから次のn個のトークンをランダムに推測する
    • プロンプトとnトークン列をLLMに入れて反復的に更新する
    • 列がそれ以上変化しなくなれば固定点(fixed point) に到達する
    • 最終的なnトークン列は、greedy戦略においてARデコーディングが生成した出力へ収束する
  • 初期のランダム推測から最終的なAR生成結果に至るまでの経路をJacobi軌道(Jacobi trajectory) と呼ぶ
  • 基本的なJacobi decodingは実際のLLMでARデコーディング比平均1.05倍程度の限定的な速度向上しか示さない
    • ARで学習されたLLMは、前方のトークンに誤りがあると後方トークンを正しく予測しにくい
    • ほとんどのJacobi反復ではnトークン列のうち1トークンしか修正されず、軌道が長くなる
  • Lookahead decodingとspeculative decodingはJacobi decodingおよび従来のARデコーディングの非効率を減らそうとするが、推論時に追加のメモリコストが発生する

Consistency LLMの学習目標

  • CLLMは、事前学習済みLLMがJacobi軌道上の任意の地点から固定点へ一貫して移動するよう適応させる方式
  • この目標は、拡散モデルの高速化手法であるconsistency modelsの目標に近い
  • 対象モデルから収集したJacobi軌道を用い、Jacobi反復中の単一ステップ収束を促す損失でモデルを学習する
  • 各対象モデルをCLLMへ変換する学習は2つの部分から構成される
    • Jacobi軌道の準備
      • 各プロンプトについて、nトークン単位のすべての切り出し区間ごとにJacobi decodingを順次実行する
      • 応答全体の列は、連続する固定点をつなぎ合わせた形になる
      • 軌道中に生成された各列が1つの学習データ項目として扱われる
      • 長い応答で入力全体を評価することによる速度低下を避けるため、nトークンの切り出しを用いる
    • Consistency lossとAR lossの共同最適化
      • consistency lossは複数トークンを一度に予測するよう促す
      • AR lossはCLLMが対象LLMの分布から逸脱しないようにし、生成品質の維持を助ける

損失関数の構成

  • 対象LLMをp、CLLMをqθとし、qθはpのパラメータで初期化される
  • Global consistency(GC) lossは、Jacobi軌道の任意状態yを入力したときCLLMが固定点y*を出力するよう促す
  • Local consistency(LC) lossは、Jacobi軌道上で隣接する状態y(j)とy(j+1)が同じ出力を返すよう合わせる
  • 分布間距離D(·||·)はGKD手法で議論された選択肢に従い、本実験では主にforward KLを用いる
  • AR lossは対象LLM pの生成結果lに基づいて従来の自己回帰損失を適用し、対象LLMの分布から外れることを防ぐ
  • 学習全体の損失はconsistency lossと、重みwを掛けたAR lossの和で構成される

実験設定と結果

  • 実験は3つの特化ドメインタスクと1つのオープンドメイン対話ベンチマークを含む
    • Spider: text-to-SQL
    • Human-Eval: Pythonコード補完
    • GSM8k: 数学
    • MT-bench: オープンドメイン対話
  • タスクに応じて対象モデルとしてfine-tuned coder LLM、Deepseek-coder-7B-instruct、LLaMA-2-7B、ABEL-7B-001を使用
  • 学習と評価はいずれもNVIDIA A100 40GBサーバー上で実施された
  • 特化ドメインでは、元の対象モデル、Medusa2、speculative decodingなどのベースラインと比較してCLLMが最大の速度向上を示した
  • MT-benchでは、ShareGPTデータセットでLLaMA2-7Bから学習したCLLMをlookahead decodingと組み合わせた場合、Medusa2と同程度の速度向上を達成した
    • MT-benchのスコアも比較可能な水準
    • CLLMは対象モデルの元のアーキテクチャ変更を必要としない
    • 補助コンポーネントがないためメモリ効率が高い

学習コスト

  • CLLMのファインチューニングコストは中程度とされる
  • LLaMA-7BはSpiderデータセットで約100万トークンだけを通し、3.4倍の速度向上を達成した
  • CodeSearchNet-Pythonのようにデータセットが大きい場合、データセットの10%だけをJacobi軌道生成に使って約2.5倍の速度向上を得た
  • 総トークン数は次式で見積もられる
    • プロンプトあたりの平均軌道数 × 平均軌道長 × プロンプト数
  • データセット別の推定学習コストは以下の通り
    • Spider: 200万トークン、事前学習コストの< 0.01%
    • CodeSearchNet-Python: 1億トークン、事前学習コストの~0.1%
    • GSM8K: 1,000万トークン、事前学習コストの~0.01%
    • ShareGPT: 2億トークン、事前学習コストの~0.2%

Fast forwardingとstationary tokens

  • 対象LLMは通常、1回のJacobi反復で正しいトークンを1つしか生成できない
  • CLLMでは、1回のJacobi反復で複数の連続トークンを正しく予測するfast forwarding現象が現れる
  • 対象LLMでは、それ以前に正しく生成されたトークンが後続の反復で誤ったトークンに変わることがある
  • CLLMは、前方トークンに誤りがあっても正しいトークンを先回りして予測し、そのトークンを以後も維持する能力を示す
    • このようなトークンをstationary tokensと呼ぶ
  • fast forwardingとstationary tokensは、CLLMのJacobi decodingをより速く収束させ、生成速度向上に寄与する

言語パターンの学習

  • CLLMは学習を通じてコロケーション(collocations) という言語概念を獲得することが観察されている
  • コロケーションとは、ランダムな偶然よりも高頻度で共起する単語や用語の連なりを指す
  • 自然言語とコードの両方にコロケーションの例がある
    • 自然言語: talk to, remind … of …
    • 動詞+名詞構造: make a decision, catch a cold
    • ドメイン別の文法構造: SELECT … FROM …, if … else
  • consistency generationの目標は、CLLMがJacobi軌道のどの地点でもこうした構造を推論できるようにし、反復回数を減らすため複数語を同時予測するよう促すことにある

資料とコード

関連記事

1件のコメント

 
GN⁺ 2024-05-09
Hacker News の意見

  • 「自由ドローイング」（講義なし）の授業を受けたときの経験に似ている
    子どもの頃から絵がうまいと言われていたが、実際には以前描いたような似た細密画を繰り返したり、かなり時間をかけて描いたりしていた記憶が残っている。時間と忍耐さえあれば、誰でも場面をかなりそれらしく描けると思う
    授業にはルールも講義もなく、各自が好きな画材を持ってきた。インクの人もいれば鉛筆の人もいて、私は木炭を持っていった。決まっていたのはモデルのポーズ間の時間だけだった。最初の数ポーズは1分ほどと非常に短く、だんだん長くなって5分ポーズまで進み、いつでも絵を破って同じポーズを描き直すことができた
    短いウォームアップは、実際には比率と輪郭を最初の試みで合わせることを強制していた。急ぐと失敗するという通念とは違い、技術を学んだり磨いたりするときには、焦りが注意と学習を促すストレス要因として働くように思う
    授業の前でも同じくらいの品質の絵は描けただろうが、5〜10倍は時間がかかったはずだと確信している。遠回りできないよう強制され、性急なミスの代償を感じるやり方が効いた
    ただ、この手法を Consistency と呼ぶのは少し残念だ。性能改善には合うが、推論速度の改善にはあまり合わない名前のように思うし、「1トークンずつ生成したときに最終的に出てきたであろう結果との一貫性」という意味は理解している。むしろ同じ出力を期待しつつ、同じ結論へたどたどしく到達する抑制のない Proficiency LLM と呼びたい

    • CLLM の著者として、経験と洞察を共有してくれてありがとう。絵の腕を磨く過程は CLLM の学習プロセスに似ているように見える。ただし、現在の CLLM 学習におけるストレス要因は、だんだん厳しくなる形ではない
      絵では各試行ごとに許容時間を決め、それを徐々に短くできる。CLLM では、Jacobi 軌跡上で最終状態からますます遠い状態を最終状態へマッピングさせることで、学習プロセスをより難しくできそうだ
      「consistency」という用語は、拡散画像生成における consistency model と consistency LLM の類似性から持ってきたもので、学習プロセスが互いに似ているためだ
    • ある夏の 無脊椎動物学の実験授業で興味深い経験をした
      学生たちは実験室に入り、標本を受け取り、指示は「これを30分以内に描け。始め」だけだった
      「絵はこう描く」「これはして、あれはするな」といった言葉はなく、実質的には「不安だろうが、自分は描けないと思っていようが関係ない。言い訳せずにただ描け。今すぐ」に近かった
      私たちはみな描いた。そして夏の間ずっと、さらに多くの動物が持ち込まれ、同じ訓練を繰り返すうちに、全員がものすごく上達した
      その授業が教えてくれたのは、誰でも、本当に誰でも描けるということだ。集団としての態度が「これが可能なのかも分からない」から「当然できる。簡単で、日常的で、大したことではない」へと変わった
      強く勧められるアプローチで、大学で受けた授業の中で最も解放感があり、驚きのある授業だった
    • システムは通常、ストレスを受けるとより効率的になる。同時に局所最適解へ押し込まれることもあるので、何事にも長所と短所がある

  • 著者たちは Jacobi デコーディングが貪欲な自己回帰デコーディングと同じだと言っているが、実際には反復や過度に一般的な回答を避けるために、サンプリング温度を0より高くしたい場合が多いのではないかと思う
    このデコーディング戦略をまったく知らないので、それを反映する簡単な方法を見落としているのかもしれない

    • 良い質問だ。貪欲サンプリング以外の サンプリング戦略もサポートしようと積極的に取り組んでいる
      CLLM の学習文脈では、Jacobi デコーディングで得られた静的な固定点を学習目標へマッピングする代わりに、これを動的固定点と呼んでいる。新しい進捗は GitHub リポジトリを見ていればよい
    • 同意する。あるトークンが最大値の選択だったかどうかは確認しやすいが、あるトークンが望む確率で現れたかどうかを確認するのは難しそうだ
      望ましい統計を持つ n トークン補完へ軌跡が近づくように学習するファインチューニング段階はまだ可能だろうが、固定点の確認段階をどう置き換えられるのかはよく分からない。おそらく「尤度について、この固定しきい値以上だったかを確認する」といった方法かもしれない

  • LLM の内部で何が起きているのかを本当に理解する前にやるには、かなり危険な最適化のように感じる。たとえば幾何学的解釈を信じる側には言い分があるだろうし、「穴埋め」トークンを使う場合には有害になり得るようにも見える
    また「頭の中で完全な文を作ってから単語ごとに話す」という仮定は普遍的な事実ではなく、あくまで仮定にすぎず、私たちの心で起きている活動を過度に単純化しているように思う。本当に話したりタイプしたりする前に完全な計画があるのだろうか。仏教徒としては、それは幻想に近いと思う。さらに、同時並行的な思考はどうなるのか。私たちは文単位で線形に考えているのか
    いずれにせよ、数学はかなり見事だ

    • この最適化は LLM の結果に影響せず、直接デコードしたものと同等の結果を生成することが保証されている
      LLM を私たちの心に似た魔法のようなものとして扱うのはやめよう。単に、もっともらしい文を生成するもう一つのプログラムにすぎない
    • その仮定はこの文脈では有用かもしれないが、事実ではないことはかなり明らかに見える
      複雑な過去の出来事を複数の筋道で説明してみろと言うと、人々は出来事全体の範囲を収めるために、文の途中で断片、補足、脇道を素早く差し込む。文単位の粒度仮説を真面目な科学的文脈で見たことはないように思う
    • 全員に当てはまるとは言えないが、少なくとも私は話す前に完全な文を頭の中で作らない
      ときどき文の途中で文法的な行き止まりに追い込まれ、不自然な単語や句で考えを締めくくらなければならなかったり、単に止まって最初から言い直したりする
    • 単語は複数の言語で意味を持つより小さな単位に分解できるにもかかわらず、私たちはその下位構成要素から順番に単語を作っているようには見えない
      この現象が文レベルで突然崩れる明確な理由もなさそうだ
    • 幾何学的解釈とは何なのか気になる

  • もっと大きな注目を集めていないのが驚くほど。これは推論性能に明確な利点があるように見える
    このファインチューニングのコストは妥当で、元の事前学習コストの約0.01%程度。性能上の利得もかなり一貫しているように見える

    • LLMの性能にとって非常に大きな結果に見える
      この程度までLLMの推論性能を高められると提案した他の論文はあまり知らない。以前にあっただろうか？
      少なくとも、出力品質を維持し、クエリのレイテンシだけでなく全体のスループットも改善し、追加の計算を必要とせず、実装も比較的実用的で大きな複雑さを加えない、という条件ではなおさらだ
      並列/Jacobiデコーディングで行われてきた研究の上に積み上げたものなので、洞察自体は漸進的とも言える。以前の成果も必要で重要だったが、並列デコーディングの可能性から現実世界での価値を引き出したのはこの結果かもしれない
    • 同程度、あるいはそれ以上の推論上の利得は、すでに広く使われている投機的デコーディングでも得られる
      なのでこの研究は本当に興味深く、私の知る限り以前にもあまり成功しない形で試みられたことはあるが、実際の影響がどれほど大きいかはまだ明確ではない
    • 私たちの研究に関心を持ってくれてありがとう。consistency loss + AR lossでデータセットの一部だけを学習しても、かなりの速度向上が得られ、コストは事前学習の0.01%程度だった
      より多くのデータで学習すれば速度はさらに上がる。モデルがより頻繁に出現するコロケーションや構文から学べるためだ
      詳細は論文を見ればよく、学習データのサイズが大きくなるにつれて速度向上が飽和することも確認できる

  • 最初は、後続トークンを推測するためにunembed headを追加で使うMedusa系の論文だと思ったが、まったく違った
    本当にすごい。追加パラメータも使わず、単に補助的な学習損失を追加しているだけだ

    • MedusaとCLLMの唯一の共通点は、どちらも高速な推論のためにLLMを学習・適応させるという点だけだ
      学習手法もデコーディング手法も完全に異なり、指摘の通りCLLMは追加パラメータも、ツリーベース検証のためのアテンションマスク設定も必要ない

  • そのうち、モデルを必ずしも学習する必要はないと気づくことになりそうだ
    必要なのは良いインデックス化とサンプリングだ
    本質的には、あるレベルではすべてのLLMはデータセットのデータベースのようなもので、その上に優れた自然言語インターフェースが載っていると見なせる
    どちらも保存されたデータを探索する異なる方法にすぎない

    • LLMは学習データセットにないデータを容易に作り出せる
      LLMは保存されたデータを探索しているわけではない。LLMは学習データのデータベースではない
    • しかしインデックス化も学習だ。ただしエンドツーエンドの勾配降下法を使わないだけだ
    • モデルは学習データの圧縮版よりも何桁も小さいので、そのデータベースと同等ではあり得ない
    • それならInfinigramの論文が気に入るかもしれない。最近議論されていた
      https://news.ycombinator.com/item?id=40266791

  • 私のように詳しくない人が「AI専門家に聞く」ことができる場所はあるだろうか？
    例えば、LLMが同じプロンプトを受け取っても、なぜ同じ決定論的な方式で応答しないのかを尋ねたい
    これを学びたいし、YouTubeの「1時間でLLMを作る」のような動画を追うべきなのかもしれない

    • ソフトウェアには、モデルが出す重み付きの次トークン候補の中から1つを選ぶ乱数生成器が文字通り入っている
      選択プロセスには、応答を操作するためのさまざまなつまみがあり得る。決定論的にしたくてソフトウェアに直接アクセスできるなら、使っているソフトウェアに応じて top-k = 1 または temperature = 0.0 に設定すれば決定論的になる
      通常、デフォルト設定は決定論的ではない。完全に決定論的にすると、結果の品質があまり良くない傾向があるためだ
    • その答えは3blue1brownの動画を参考にするとよい
      LLMモデルはトークンに対する確率ベクトルを出力し、LLMのユーザーは乱数を使って可能性の高いリストからトークンを選ぶ
    • LLMは本質的に確率行列だからだ
      プロンプトを入力すると次の単語が出る確率を計算し、そのプロセスを繰り返して最終的に文を作る。学習された確率は学習データに基づく
      このような基盤となる確率モデルのため、100%決定論的ではない。さらにChatGPTのようなモデルは、意図的に temperature パラメータを設けてプロセス全体にランダム性を加えている
      さらに読みたいなら、この回答は次の論文に基づいている: The Matrix: A Bayesian learning model for LLMs, https://arxiv.org/abs/2402.03175
    • ほとんどのシステムでは、temperatureという推論設定パラメータでこれを制御できる
      ただし温度を可能な限り低く設定すると、回答品質が非常に低くなる傾向がある。システムがある局所最適解から抜け出せず、繰り返し続けるようになる。そうした回答は「決定論的」ではあり得るが、良いものではない
    • この記事はかなり体系的に説明しつつ、大局を見失わない良い出発点だった
      https://writings.stephenwolfram.com/2023/02/what-is-chatgpt-...

  • ただ飯はないので、私の見るところここにも何らかの経路損失がある
    例えば一部のJacobi軌跡は、定義上、より高い温度の経路を除外する。データ検索の観点ではむしろプラスかもしれないが、創造性を最大化したいならマイナスになり得る

    • より良いアルゴリズムと、より悪いアルゴリズムはある
      「ただ飯はない」が常に特別に意味のある形で当てはまるのかはよく分からない。あるものはパレートフロンティア上にない

  • 「投機的デコーディング方式は推論時に追加のメモリコストがかかる」という部分について、詳しい説明が気になります
    投機的デコーディングでは、より小さいモデルが高速だが不正確な可能性のある「分岐」を生成し、その後、大きいモデルでこれらの分岐を検証します。しかし投機的デコーディングに必要なのは単一トークンに相当するメモリだけで、他の分岐のトークンは推論中に単にマスクされます。コンテキストサイズが1000で、5トークンの分岐が約30個なら、メモリオーバーヘッドは3%で無視できる程度です。コンテキストサイズが分岐数に比べてはるかに小さいなら、コンテキストウィンドウが50トークンしかない生成LLMのユーザーが生成速度を気にするのかと思います
    また、投機的デコーディング手法は貪欲サンプリングに限定されません。元のモデルとまったく同じように動作し、期待確率でサンプリングする必要があります。投機的デコーディング関連の文献の多くは、すでに2.6〜3.5倍の高速化を報告しています。このブログ記事は2.4〜3.4倍の生成速度を報告していますが、それほど大きなアップグレードなのかは分かりません
    上で投機的デコーディングに触れましたし、著者が比較した技術はMedusa2とEagleのようですが、核心的な問題はそのままです。トークンを事前に予測するどんな方式を使っても、次のトークンを予測する前に前のトークンが絶対に必要になる特定の地点があります。これはモデルや手法の問題ではなく、数学的に何が可能かという問題です。5番目に続くトークンの確率分布が先行する4トークンに大きく依存するなら、どうやって一度に5トークンを予測できるのでしょうか？ 投機的デコーディングでも、Jacobiデコーディングでも、マルチトークン並列デコーディングでも同じです
    この方式が貪欲サンプリングだけをサポートするなら、利点が何なのか疑問です。他の手法がすでに期待された高速化を達成している点まで考えると、なおさらです。貪欲サンプリングの高速化とランダムサンプリングの高速化を比較するのは、リンゴとオレンジを比べるようなものですし、この方式をランダムサンプリング向けに変えた後でも、上で述べた核心的な問題のために同じ高速化が残るのかは疑わしいです

    • 「次のトークンを予測する前に前のトークンが絶対に必要だ」という部分が、この論文の核心的な貢献なのかもしれません
      consistency trainingによって、前のトークンに誤った推測があってもLLMが次のn個のトークンを予測できることを示したのかもしれません
      一方、数学的には p(x_t|x_1,...,x_t-1) が x_1 から x_t-1 までのすべてに依存するのは正しいですが、実際には x_t の予測に必要なのは x_1 から x_t-2 までで、x_t-1 へのアテンションはごく小さいこともあります。したがって、x_1 から x_t-2 と不正確な x_t-1 でも x_t を予測できます
    • 投機的デコーディングでは、より小さいモデルをメモリに載せ、そのモデルで推論を実行する必要があります

  • 興味深い内容です。多くの人が思いつきそうなアイデアですが、記事と発表はよく整理されていました

    • その通りです。ルームメイトと私は1年前にこういう話をしていました。LLMステアリングにも似たことができます