複数生成パラダイムを統合するオープンソースライブラリ DANTINOX──JAX/Flaxによる制御されたクロスパラダイム比較を実現

論文概要

近年の言語生成研究において、自己回帰デコーディング（Autoregressive）、離散マスク拡散（Discrete Masked Diffusion）、連続フローマッチング（Continuous Flow-Matching）という3つの主要なパラダイムが混在している。しかし、これらは個別のコードベースで実装されることが多く、報告される性能差が純粋なパラダイムの違いによるものか、実装の詳細やコードの差異に起因するのかを切り分けることが困難であった。本論文では、単一のモジュール式Transformerバックボーンを用いてこれら3つのパラダイムを統一的に扱うオープンソースのJAX/Flaxライブラリ「DANTINOX」を提案する。

関連研究

従来の大規模言語モデル（LLM）の研究は自己回帰モデルが主流であったが、非自己回帰的な生成手法として拡散モデルやフローマッチングの適用が進んでいる。これまでの比較研究はコードベースが完全に独立していたため、トークナイザや初期化戦略の不一致がボトルネックとなり、公平なクロスパラダイム比較を行うことは困難であった。

新規性と貢献

本研究の最大の貢献は、バックボーンアーキテクチャ、トークナイザ、初期化戦略、学習インフラストラクチャを完全に共通化しつつ、設定の変更のみで生成パラダイムを切り替えられる統一フレームワークを提供した点にある。これにより、コードの差異によるノイズを排除した厳密なクロスパラダイム比較が単一のAPI上で実現された。

提案手法の詳細

DANTINOXはJAX/Flaxベースで実装されており、バックボーンとなるTransformerの構造を変更することなく、設定ファイル（Configuration）の変更だけで生成パラダイム、アテンション機構、ハードウェアトポロジを動的に切り替えることができる。初期化戦略やデータパイプラインを共通化することで、異なるパラダイム間での公平なアブレーションスタディやスケーリング則の検証が可能になる。

評価・考察

論文では、異なる生成パラダイムにおける推論効率の比較や、ハードウェア利用率（MFU: Model FLOPs Utilization）の屋根線分析（Roofline analysis）が行われている。生成長やバッチサイズを変えた場合のレイテンシ、スループット、メモリ使用量の変化が同一のバックボーン上で定量的に測定され、各パラダイムの実用上の特性が明らかにされている。

応用例と今後の展望

自然言語処理（NLP）分野における基盤モデルの研究開発において、非自己回帰生成モデルの導入を検討している研究チームやテックリードにとって、モデルのアーキテクチャ起因の性能差を排除した検証が可能になる。特に大規模な分散学習やストリーミング推論の最適化を行う際の実務的なベンチマークツールとして活用できる。

結論

DANTINOXは、バラバラに発展してきた複数の言語生成パラダイムを単一のコードベースに統合し、公平な比較と効率的な検証を可能にする実用的なフレームワークである。

注釈

• 自己回帰（Autoregressive）: トークンを1つずつ順番に生成していく従来のLLMの主流方式。
• フローマッチング（Flow-Matching）: 連続的な確率変数の軌跡を学習することで高速な生成を目指す生成モデルの枠組み。