Chengyue Wu Hao Zhang Shuchen Xue Shizhe Diao Yonggan Fu Zhijian Liu Pavlo Molchanov Ping Luo Song Han Enze Xie

要約

自己回帰型（AR）大規模言語モデル（LLM）は、幅広い自然言語処理タスクにおいて優れた性能を達成しているが、その固有の逐次デコード構造により、推論効率に限界がある。本研究では、事前学習済みARモデルを効率的に拡散型言語モデル（dLLM）に変換するための、精密に設計されたブロック拡散型言語モデル「Fast-dLLM v2」を提案する。このモデルは、並列テキスト生成を実現するため、約10億トークン程度のファインチューニングで十分であり、Dream（5800億トークン）のようなフルアテンション拡散LLMと比較して、学習データ量を500倍削減しつつも、元のモデルの性能を維持する。我々のアプローチは、ブロック拡散機構と補完的アテンションマスクを組み合わせた新しい学習手法を導入し、ARモデルの学習目標を損なうことなく、ブロック単位での双方向的文脈モデリングを可能にする。さらにデコード速度を向上させるために、階層的なキャッシュ機構を設計した。具体的には、ブロック単位のキャッシュにより、複数ブロック間の歴史的文脈表現を保存し、サブブロックキャッシュにより、部分的にデコードされたブロック内での効率的な並列生成を実現する。このキャッシュ機構と並列デコードパイプラインを組み合わせることで、Fast-dLLM v2は生成品質を損なうことなく、標準的なARデコードと比較して最大2.5倍の高速化を達成する。多様なベンチマーク上で実施した広範な実験により、Fast-dLLM v2がARベースラインと同等あるいはそれを上回る精度を達成するとともに、dLLMの中でも最先端の効率性を実現することを示した。これは、高速かつ高精度なLLMの実用的導入に向けた重要な一歩である。コードおよびモデルは公開予定である。