【Disffusion Large Language Model】Dream 7Bの推論を試す
はじめに
Dream 7B (Diffusion Reasonig Model)なるLLMが公開されました.
こちらのブログポストの内容に詳細がありますが,LLaDAなどの既存のDiffusion LMの性能を越え,Llama3 8Bなど既存の同程度の規模のAuto Autoregressive Modelと比較してgeneral, math, and codingのベンチマークで高い性能を示しているとのことです.
Mercury CoderやBlock Diffusionなど,のニュースも記憶に新しいですが,LLM x Diffusionの潮流が来ており,精度や速度の観点でも実用に近くなってきたと考えられます.
本記事では手元のGPUサーバでDream-v0-Instruct-7Bの推論をテストしてみます.
関連情報
デモもこちらで試せるようです:
開発者のアナウンスポスト:
Dream 7Bの基本情報
ほぼブログのまとめなので詳細はブログをご確認ください.
テキスト生成方式はDiscrete diffusion models (離散拡散モデル)を用いており,初期状態は全ての単語をノイズ(mask)で置換し,モデルが逐次的にテキスト全体を並行して洗練してくようなアプローチをとります.生成時に双方の文脈情報を統合することができ,長い文脈でも整合性の取れたテキストを生成しやすくなると言います.
下図はAuto Autoregressive (以下ARと略)とDIffusionの生成の違いを示しています.
- ARのTrasnformer DecoderではCausal Attention Maskを適用して文頭からnext token prediction
- DreamではTrasnformerが全トークンにAttention(Full Attention)して,各ステップで一部の
<mask>
トークンを予測,置換(remasking)しながら文全体を並列的に生成する
DreamはARモデルからアーキテクチャを継承をする設計になっており,Casual Attention MaskをFUll Attentionに変える学習をしています.Dream 7BはQwen2.5 7Bの重みを初期値としており,full scratchでの学習より効率的にtrainingをしているそうです.この時学習率の設定が難しく,値が大きすぎるとARモデルから引き継いだ知識が早期に消失する,小さすぎると拡散モデルの適用が進まない,などの現象が報告されています.
Dreamは生成順序を意図的に制御することもできるようで,通常の左から右への順次生成に近い振る舞いをさせれば文章をスムーズに続きを書かせることができ,ランダムな順序で生成を進める設定にすれば特定の制約を満たす文章構築(必要な要素を逐次充填していくような生成)ができるようです.文章の部分編集や特定フォーマットに向けた文生成など,きめ細かなテキスト操作を要する応用では期待できる仕組みだと思いました.
pretrain
公開データのAllen InstituteのDolma v1.7,OpenCoder,DCLM-Baselineなどを組み合わせた5,800億トークン規模のテキストを使っており,H800 GPU x 96枚, 256h(10日間)かけて学習されたそうです...
SFT
SFTについてはまだあまり情報が公開されていませんでした. Tulu3とSmolLM2の1.8Mのデータを用いて指示追従性高めているそうです.
実装コード
結論 transformers
で超簡単に推論が可能でした.(基本はQwenと同じデコーダ構造を持っているのでembeddingやfeedfowrd, unebmbeddingなどは共通しているため,mask部分を変更するだけでよく,推論実装時はそこまで考えることが少ないように思います.)オリジナルのGitHubにコードがありましたが,serve機能を実装したものを下記にpushしました.
コードは本当に雑に書いただけですが,下記で環境構築を行なったのち,
uv sync
モデルのserve
uv run bash scripts/serve.sh
対話は下記で可能です.
uv run bash scripts/client.sh
オリジナルのexampleコードは非常にシンプルで下記の通りです.
import time
import torch
from transformers import AutoModel, AutoTokenizer
model_path = "Dream-org/Dream-v0-Instruct-7B"
model = AutoModel.from_pretrained(model_path, torch_dtype=torch.bfloat16, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = model.to("cuda").eval()
messages = [
{"role": "user", "content": "Please write a Python class that implements a PyTorch trainer capable of training a model on a toy dataset."}
]
inputs = tokenizer.apply_chat_template(
messages, return_tensors="pt", return_dict=True, add_generation_prompt=True
)
input_ids = inputs.input_ids.to(device="cuda")
attention_mask = inputs.attention_mask.to(device="cuda")
output = model.diffusion_generate(
input_ids,
attention_mask=attention_mask,
max_new_tokens=512,
output_history=True,
return_dict_in_generate=True,
steps=512,
temperature=0.2,
top_p=0.95,
alg="entropy",
alg_temp=0.,
)
generations = [
tokenizer.decode(g[len(p) :].tolist())
for p, g in zip(input_ids, output.sequences)
]
print(generations[0].split(tokenizer.eos_token)[0])
model.diffusion_generate
というメソッドが見慣れないですが,model.generate()
に準拠しており,Diffusion LLM特有の推論パラメータを受け付けています.
特徴的なものは下記の通りです.
-
steps
: 拡散過程のタイプステップで,各ステップでmax_new_tokens
/steps
トークンが生成される.ステップ数は速度と品質のトレードオフ. -
alg
: diffusion samplingのmasking戦略でトークン生成の順序を制御する.下記の1つのランダム戦略と3つのconfidence-based戦略がある.-
origin
(default):https://arxiv.org/abs/2107.03006 の完全にランダムな順序でトークンが生成される. -
maskgit_plus
:https://arxiv.org/abs/2202.04200 のトップ1の信頼度に基づいてトークンが生成される -
topk_margin
:https://arxiv.org/abs/2502.06768 のtop1 - top2
のマージン信頼度に基づいてトークンが生成されます -
entropy
:各トークン分布のエントロピーに基づいてトークンが生成されます。
-
- alg_temp:信頼度ベースの戦略を使用する際に
alg
にランダム性を追加することができる
推論テスト
環境はRTX 4000Ada x 4枚の環境を用いてますが,1GPUで推論可能であることが確認できました.
math
まずは得意と謳っている数学のテストを試します.GSM8Kの適当な問題を与えましたが,下記のように正解していることがわかります.(ちょっと簡単な問題ですが)
備考:gold answer
Natalia sold 48/2 = <<48/2=24>>24 clips in May.
Natalia sold 48+24 = <<48+24=72>>72 clips altogether in April and May.
#### 72
coding
次は適当にコーディングさせてみます.最長交互増加部分列に関するシンプルな問題を作ってみました.言語はあえて指定していません.
与えられた数値配列から「交互増加部分列」の最大長を求めるプログラムを実装してください。「交互増加部分列」とは、隣接する要素が交互に増加または減少する部分列です。例えば配列[1,7,4,9,2,5]の場合、[1,7,4,9,2,5]全体が交互増加部分列となり、長さは6です。
うん.ダメでした.inputを色々試しましたが,空のresponseを返してきます.
japanese
次に日本語を話せるのかテストしてみました.例の如くELYZA-tasksk-100から適当な問題をサンプリングしてみましたが,意外にも正解していますね.
下記のような問題は不正解でした.知識を問う問題は回答できないことがわかりましたが,回答自体は日本語で返しています.
ネガティブな結果としては,例えば自由に何か生成させるタスクは回答してくれませんでした.下記はELYZA-tasksの自由文章生成タスクの問題(日英)ですが,日本語は無回答,英語はrefusalといった結果でした.
おわりに
今回は本当に雑に試しただけですが,あまり回答は安定しないという結果になりました.回答できている問いもベースとなったQwenの貢献度も大きそうです.ただ,Diffusion Language Model特有の推論の工夫とかも色々あると思うので,もう少し遊んでみようかなと思います.生成品質は一昔前のLlamaを思い出しました.これから日本語版のDiffusion Language Modelも出てきそうで楽しみです.Diffusion LMの入門として良い機会でした.
Discussion