アシストされた生成：低遅延テキスト生成への新たな方向

'Assisted Generation A New Direction Towards Low-Latency Text Generation'

大規模な言語モデルは最近注目を集めており、多くの企業がそれらを拡大し、新たな機能を開放するために多大なリソースを投資しています。しかし、私たち人間は注意力が減少しているため、彼らの遅い応答時間も嫌いです。レイテンシは良好なユーザーエクスペリエンスにおいて重要であり、低品質なものである場合でも（たとえば、コード補完において）小さいモデルが使用されることがよくあります。

なぜテキスト生成は遅いのでしょうか？破産せずに低レイテンシな大規模な言語モデルを展開する障害物は何でしょうか？このブログ記事では、自己回帰的なテキスト生成のボトルネックを再検討し、レイテンシの問題に対処するための新しいデコーディング手法を紹介します。私たちの新しい手法であるアシスト付き生成を使用することで、コモディティハードウェアでレイテンシを最大10倍まで削減できることがわかります！

テキスト生成のレイテンシの理解

現代のテキスト生成の核心は理解しやすいです。中心となる部分であるMLモデルを見てみましょう。その入力には、これまでに生成されたテキストや、モデル固有のコンポーネント（Whisperのようなオーディオ入力もあります）など、テキストシーケンスが含まれます。モデルは入力を受け取り、順番に各層を通って次のトークンの非正規化された対数確率（ログット）を予測します。トークンは、モデルによって単語全体、サブワード、または個々の文字で構成されることがあります。テキスト生成のこの部分について詳しく知りたい場合は、イラスト付きのGPT-2を参照してください。

モデルの順方向パスによって次のトークンのログットが得られますが、これらのログットを自由に操作することができます（たとえば、望ましくない単語やシーケンスの確率を0に設定する）。テキスト生成の次のステップは、これらのログットから次のトークンを選択することです。一般的な戦略は、最も可能性の高いトークンを選ぶことで、これはグリーディデコーディングとして知られています。また、トークンの分布からサンプリングすることも行います。モデルの順方向パスと次のトークンの選択を反復的に連結することで、テキスト生成が行われます。デコーディング手法に関しては、この説明はアイスバーグの一角に過ぎません。詳細な探求のために、テキスト生成に関する当社のブログ記事を参照してください。

• RWKVとは、トランスフォーマーの利点を持つRNNの紹介です
• 単一のGPUでChatgptのようなチャットボットをROCmで実行する
• より小さいほうが良いです：Xeon上で効率的な生成AI体験、Q8-Chat

上記の説明から、テキスト生成のレイテンシのボトルネックは明確です。大規模なモデルの順方向パスを実行することは遅いため、連続して何百回も実行する必要があります。しかし、さらに詳しく見てみましょう。なぜ順方向パスが遅いのでしょうか？順方向パスは通常、行列の乗算によって支配されます。対応するウィキペディアのセクションを素早く訪れると、この操作における制限はメモリ帯域幅であることがわかります（たとえば、GPU RAMからGPUコンピュートコアへ）。つまり、順方向パスのボトルネックは、デバイスの計算コアにモデルのレイヤーの重みを読み込むことから来ており、計算自体ではありません。

現時点では、テキスト生成の最大の効果を得るために探求できる3つの主な方法がありますが、すべてがモデルの順方向パスのパフォーマンスに対処しています。まず第一に、ハードウェア固有のモデルの最適化があります。たとえば、デバイスがFlash Attentionに対応している場合、操作の並べ替えによってアテンションレイヤーの処理を高速化することができます。また、モデルのウェイトのサイズを削減するINT8量子化もあります。

第二に、同時にテキスト生成のリクエストがあることがわかっている場合、入力をバッチ化し、小さなレイテンシのペナルティを支払いながらスループットを大幅に増加させることができます。デバイスに読み込まれたモデルのレイヤーのウェイトは、複数の入力行で並列に使用されるため、ほぼ同じメモリ帯域幅の負荷でより多くのトークンを出力することができます。バッチ処理の問題は、追加のデバイスメモリが必要であることです（またはメモリをどこかにオフロードする必要があります）-このスペクトルの終端では、レイテンシを犠牲にしてスループットを最適化するFlexGenなどのプロジェクトが見られます。

# バッチ化された生成の影響を示す例。計測デバイス: RTX3090
from transformers import AutoModelForCausalLM, AutoTokenizer
import time

tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2").to("cuda")
inputs = tokenizer(["Hello world"], return_tensors="pt").to("cuda")

def print_tokens_per_second(batch_size):
    new_tokens = 100
    cumulative_time = 0

    # ウォームアップ
    model.generate(
        **inputs, do_sample=True, max_new_tokens=new_tokens, num_return_sequences=batch_size
    )

    for _ in range(10):
        start = time.time()
        model.generate(
            **inputs, do_sample=True, max_new_tokens=new_tokens, num_return_sequences=batch_size
        )
        cumulative_time += time.time() - start
    print(f"1秒あたりのトークン数: {new_tokens * batch_size * 10 / cumulative_time:.1f}")

print_tokens_per_second(1)   # 1秒あたりのトークン数: 418.3
print_tokens_per_second(64)  # 1秒あたりのトークン数: 16266.2（1秒あたりのトークン数が約39倍増加）

最後に、複数のデバイスが利用可能な場合、テンソル並列処理を使用してワークロードを分散し、レイテンシを低くすることができます。テンソル並列処理では、メモリ帯域幅の負荷を複数のデバイスに分散させますが、複数のデバイスを実行するための金銭的なコストに加えて、デバイス間の通信のボトルネックも考慮する必要があります。利点は主にモデルのサイズに依存します。シングルの一般的なデバイスに簡単に収まるモデルでは、利点は非常に限定的です。この DeepSpeed ブログ記事の結果を見ると、17B パラメータモデルを 4 つの GPU に分散させることで、レイテンシを 1.5 倍短縮できることがわかります（図 7）。

これらの 3 つの改善策は併用することができ、高いスループットのソリューションを実現することができます。ただし、ハードウェア固有の最適化を適用した後、レイテンシを低減するためのオプションは限られており、既存のオプションは高価です。それを解決しましょう！

言語デコーダのフォワードパス、再考

前述の通り、各モデルのフォワードパスは次のトークンのログを生成しますが、これは実際には不完全な説明です。テキスト生成中、典型的なイテレーションでは、モデルは最新の生成されたトークンを入力として受け取り、すべての他の前の入力のキャッシュされた内部計算を返し、次のトークンのログを生成します。キャッシュは冗長な計算を避けるために使用され、より高速なフォワードパスを実現しますが、必須ではありません（一部のみ使用することもできます）。キャッシュが無効になっている場合、入力にはこれまでに生成されたトークンの完全なシーケンスが含まれ、出力にはシーケンス内のすべての位置の次のトークンに対応するログが含まれます！位置 N のログは、入力が最初の N トークンで構成され、シーケンス内のすべての後続トークンを無視した場合の次のトークンの分布に対応します。貪欲デコーディングの特定の場合では、生成されたシーケンスを入力として渡し、結果のログに argmax 演算子を適用すると、生成されたシーケンスが得られます。

from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")

inputs = tok(["The"], return_tensors="pt")
generated = model.generate(**inputs, do_sample=False, max_new_tokens=10)
forward_confirmation = model(generated).logits.argmax(-1)

# We exclude the opposing tips from each sequence: the forward pass returns
# the logits for the next token, so it is shifted by one position.
print(generated[0, 1:].tolist() == forward_confirmation[0, :-1].tolist())  # True

つまり、モデルのフォワードパスを別の目的に使用することができます。次のトークンを予測するためにいくつかのトークンを供給するだけでなく、モデルにシーケンスを渡して同じシーケンス（またはその一部）を生成するかどうかを確認することもできます。

一瞬、魔法のレイテンシフリーなオラクルモデルにアクセスできると仮定してみましょう。このオラクルモデルは、任意の入力に対してモデルと同じシーケンスを生成します。議論のために、直接使用することはできず、ジェネレーション手順のアシスタントに限定されています。上記で説明した特性を使用すると、このアシスタントモデルを使用して、候補の出力トークンを取得し、モデルを使用してそれらが本当に正しいかどうかを確認することができます。この理想的なシナリオでは、テキスト生成のレイテンシは O(n) から O(1) に短縮されます（ここで n は生成されたトークンの数です）。長い生成においては、数桁のレベルでの短縮です。

現実に近づく一歩として、アシスタントモデルはオラクルの特性を失っていると仮定しましょう。これは、あるモデルによれば、いくつかの候補トークンが誤っているということです。タスクの自己回帰的な性質により、アシスタントがトークンを誤ると、その後のすべての候補は無効になる必要があります。ただし、モデルで誤ったトークンを修正した後、アシスタントに再度クエリを行い、このプロセスを反復することは可能です。アシスタントがいくつかのトークンに失敗しても、テキスト生成は元の形式よりも桁違いにレイテンシが少なくなります。

もちろん、レイテンシフリーなアシスタントモデルは存在しません。ただし、他のモデルのテキスト生成の出力を近似するモデルを比較的簡単に見つけることができます。同じアーキテクチャのより小さなバージョンは、この特性を備えていることがあります。さらに、モデルサイズの差が大きくなると、いくつかのフォワードパスをスキップする利点を考慮した後に、より小さなモデルをアシスタントとして使用するコストは考慮に入れる必要があります。これで、アシストされたジェネレーションの核心を理解しました。

アシストされたジェネレーションにおける貪欲デコーディング

アシストされたジェネレーションはバランスの取れた行為です。アシスタントが候補のシーケンスを素早く生成する一方で、可能な限り正確にすることが望ましいです。アシスタントの品質が低い場合、アシスタントモデルを使用するコストがかかるだけでなく、ほとんど利益を得ることができません。一方、候補のシーケンスの品質を最適化するには、遅いアシスタントの使用が必要となり、結果としてネットの遅延が発生します。アシスタントモデルの選択を自動化することはできませんが、アシスタントとの時間を制限するための追加の要件とヒューリスティックを含めました。

まず、要件 – アシスタントはモデルとまったく同じトークナイザを持っている必要があります。この要件がなければ、高価なトークンのデコードと再エンコードの手順を追加する必要があります。さらに、これらの追加手順はCPU上で行われる必要があり、それによって遅いデバイス間のデータ転送が必要になるかもしれません。アシスタントの高速な使用は、支援生成の利点が現れるために重要です。

最後に、ヒューリスティック。この時点で、映画「インセプション」と支援生成の類似点に気付いたかもしれません – 結局のところ、テキスト生成をテキスト生成の内部で実行しています。候補トークンごとにアシスタントモデルの前方パスが行われ、前方パスが高価であることはわかっています。アシスタントモデルが正しく取得するトークンの数を事前に知ることはできませんが、この情報を追跡してアシスタントに要求する候補トークンの数を制限することができます – 出力の一部は他の部分よりも簡単に予測できるからです。

すべてをまとめると、支援生成ループの元の実装は次のようになります（コード）：

1. アシスタントモデルを使用して貪欲デコードを行い、candidatesという一定数の候補トークンを生成します。最初の支援生成呼び出しのとき、生成される候補トークンの数は5と初期化されます。
2. モデルを使用して、candidatesを使用して前方パスを行い、logitsを取得します。
3. トークンの選択方法（貪欲探索の場合は.argmax()、サンプリングの場合は.multinomial()）を使用して、logitsからnext_tokensを取得します。
4. next_tokensをcandidatesと比較し、一致するトークンの数を取得します。この比較は左から右への因果関係で行われる必要があることを覚えておいてください：最初の不一致後、すべての候補は無効になります。
5. 一致するトークンと（モデルが有効な候補の部分系列から生成する）最初の分岐トークンを含むように、トークンをスライスし、未確認の候補トークンに関連する変数を破棄します。
6. 次の反復で生成する候補トークンの数を調整します – 元のヒューリスティックでは、すべてのトークンが一致する場合に2増やし、それ以外の場合に1減らします。

🤗 Transformersでは、このプロセスが手間いらずになるようにAPIを設計しました。新しいassistant_modelキーワード引数でアシスタントモデルを渡すだけで、遅延時間の短縮を享受できます！このブログ投稿のリリース時点では、支援生成はバッチサイズ1に制限されています。

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

prompt = "Alice and Bob"
checkpoint = "EleutherAI/pythia-1.4b-deduped"
assistant_checkpoint = "EleutherAI/pythia-160m-deduped"
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer(prompt, return_tensors="pt").to(device)

model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
assistant_model = AutoModelForCausalLM.from_pretrained(assistant_checkpoint).to(device)
outputs = model.generate(**inputs, assistant_model=assistant_model)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
# ['Alice and Bob are sitting in a bar. Alice is drinking a beer and Bob is drinking a']

追加の内部的な複雑さはそれに値するでしょうか？ 貪欲デコードの場合の待ち時間の数値を見てみましょう（サンプリングの結果は次のセクションにあります）、バッチサイズ1を考慮に入れます。これらの結果は、追加の最適化なしで直接🤗 Transformersから取得されたものであり、設定で再現できるはずです。

収集した数値を一目見ると、支援生成はさまざまな環境で著しい待ち時間の削減をもたらすことがわかりますが、それは万能な解決策ではありません – ユースケースに適用する前にベンチマークを行う必要があります。支援生成は次のように結論付けることができます：

1. 🤏 より大きな差があるほど、モデルよりも少なくとも1桁小さいアシスタントモデルにアクセスする必要があります（差が大きいほど良い）；
2. 🚀 GPUメモリにモデルが収まる場合、INT8の存在下で最大3倍の高速化、それ以外の場合は最大2倍の高速化が得られます；
3. 🤯 GPUに収まらないモデルでメモリオフロードを利用している場合、最大10倍の高速化が見られることがあります；
4. 📄 自動音声認識や要約などの入力に基づくタスクで輝きます。

アシストされた生成のサンプル

グリーディデコーディングは、入力に基づいたタスク（自動音声認識、翻訳、要約など）や事実の知識を求めるタスクに適しています。言語モデルをチャットボットとして使用するような大きな創造性を必要とするタスクなど、大量の創造性を必要とするオープンエンドのタスクでは、サンプリングを使用する必要があります。アシストされた生成は、本来的にグリーディデコーディング向けに設計されていますが、マルチノミアルサンプリングと組み合わせてアシストされた生成を使用することも可能です！

次のトークンの確率分布からサンプルを抽出すると、グリーディアシスタントはより頻繁に失敗し、レイテンシの利点が減少します。ただし、多くのサンプリングベースのアプリケーションで存在する温度係数を使用して、次のトークンの確率分布の鋭さを制御することができます。温度が0に近い場合、サンプリングはグリーディデコーディングに近づき、最も確率の高いトークンを優先します。一方、温度を1よりも大きな値に設定すると、サンプリングはカオスになり、一様分布から抽出されます。したがって、低い温度はアシスタントモデルにとって有利であり、アシストされた生成のレイテンシの利点の大部分を保持することができます。

自分自身で確認して、アシストされた生成の感触を掴んでみませんか？

今後の展望

アシストされた生成は、現代のテキスト生成戦略が最適化の余地があることを示しています。現在の問題は、計算によるものではなく、メモリに制約があることを理解することで、利用可能なメモリ帯域幅を最大限に活用するためのシンプルなヒューリスティックスを適用できます。さらに、アシスタントモデルの使用方法をさらに洗練させることで、レイテンシの削減がさらに進むと考えています。たとえば、アシスタントに複数の候補の継続を生成するように要求することで、さらにいくつかの順方向のパスをスキップできるかもしれません。自然なアシスタントモデルをリリースして使用できるようにすることは、利点を実現し増幅するために重要です。

初めに、🤗 Transformersライブラリで公開され、.generate()関数とともに使用することを想定していますが、将来的にはHugging Faceのユニバース全体で提供する予定です。また、その実装は完全にオープンソースですので、テキスト生成に取り組んでおり、当社のツールを使用していない場合は、参考として自由に使用してください。

最後に、アシストされた生成はテキスト生成における重要な問題を再浮上させます。このフィールドは、すべての新しいトークンが固定された計算量によって生成されるという制約の下で進化してきました。純粋な自己回帰的な方法で、一度の同質な順方向のパスで1つのトークンずつ生成されます。このブログ記事は、それが必ずしもそうであるべきではないという考えを強化しています。生成された出力の大部分も、サイズの一部で生成されるモデルによって同様に生成できます。そのためには、新しいモデルアーキテクチャとデコーディング手法が必要です。私たちは将来の展望に期待しています！

関連する研究

このブログ記事の初版リリース後、同じ基本原理（より長い継続を検証するための順方向のパスを使用する）を探究した他の研究も存在することがわかりました。特に、以下の研究にご覧ください：

• Google Brainによるブロックごとの並列デコーディング
• DeepMindによる仮説的なサンプリング

引用

@misc {gante2023assisted,
    author       = { {Joao Gante} },
    title        = { Assisted Generation: a new direction toward low-latency text generation },
    year         = 2023,
    url          = { https://huggingface.co/blog/assisted-generation },
    doi          = { 10.57967/hf/0638 },
    publisher    = { Hugging Face Blog }
}

謝辞

このブログ記事を改善するために多くの有益な提案を共有してくれたSylvain Gugger、Nicolas Patry、Lewis Tunstallに感謝します。また、ウェブページで見ることができる美しいカバーをデザインしてくれたChunte Leeにも感謝します。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!