要約
研究者らがマスク拡散言語モデルの推論を高速化する「モデルスケジューリング」手法を発表。 拡散ステップの両端に小型モデルを配置する「サンドイッチ型」が最も効果的で、FLOPs17%削減を実現。 生成品質の劣化をわずか3.4%に抑えつつ高速化でき、実用的なLLM推論最適化の選択肢として注目される。
マスク拡散LMの推論、17%高速化する新手法登場
Not All Denoising Steps Are Equal: Model Scheduling for Faster Masked Diffusion Language Models
原文を読む →
https://huggingface.co/papers/2604.02340
本論文はマスク拡散言語モデル（MDLM）の推論コスト削減を目的とし、拡散の各ステップで使用するモデルサイズを動的に切り替える「モデルスケジューリング」を提案する。MDLMは自己回帰LMと異なりKVキャッシュを活用できず、全シーケンスに対して何百回も大型Transformerを走らせる必要があるため推論が高コストという課題がある。提案手法では12ブロックの重いモデルと4〜10ブロックの軽いモデルを別々に学習し、推論時に特定のタイムステップで軽いモデルに置き換える。再学習・蒸留不要のアーキテクチャ非依存な手法である。OpenWebTextとLM1Bの2データセットで実験した結果、拡散軌跡の中間ステップが最も置換に敏感であり、両端（序盤・終盤）への小モデル配置が有効という一貫した知見が得られた。「サンドイッチスケジュール」（軽い→重い→軽い）が最良で、FLOPs16.7%削減・パープレキシティ劣化3.4%を達成。損失差分とKLダイバージェンスの分析によって中間ステップでの小大モデル間不一致が最大になることを機構的に説明し、連続画像拡散で見られる単調なトレンドとは異なる非単調な重要度構造を発見した。KVキャッシュ手法や反復削減手法と直交しており組み合わせ可能な点も実用上の強みである。