マスク離散系列モデルにおけるペアワイズ相互情報量のニューラル推定
研究者らは、事前学習済みマスク拡散モデルの隠れ状態から直接ペアワイズ条件付き相互情報量を推定するニューラルフレームワークを提案し、相互情報量に基づく並列復号を可能にすることで、数独とタンパク質配列生成タスクにおいて3~5倍の高速化を達成した。
機械学習や生成モデルの分野では、変数間の依存関係を理解することが解釈可能性や効率的な生成にとって重要です。しかし、現在主流のマスク拡散モデル(MDM)は周辺条件付き分布を提供するのみで、変数間の依存構造を明示的に表現できず、効率的な推論が求められるシナリオでの応用が制限されています。この問題に対処するため、研究チームの最新研究では、事前学習済みMDMの隠れ状態から直接ペアワイズ条件付き相互情報量(MI)を推定するニューラルフレームワークを提案し、モデル内部の依存構造を明らかにします。
このフレームワークの核心は、モデルの隠れ状態を入力とし、モデル自身の条件付き分布から計算される真の相互情報量を教師信号として、変数間のペアワイズ条件付きMIを予測するニューラルネットワーク推定器を訓練することです。訓練後、この推定器は単一の前向きパスで完全なMI行列を出力し、モデルの依存構造に対する内部信念を正確に捉えます。この行列に基づき、条件付き独立な変数の部分集合を特定することで、MIに導かれた並列復号手法を開発しました。これにより、生成過程で複数の変数を並列に生成でき、推論を大幅に高速化します。
提案手法の有効性を検証するため、数独とタンパク質配列生成の二つのタスクで評価を行いました。数独タスクでは、MIマップが数独ルールに内在する構造的制約(例えば、同じ行の数字は互いに異なる)を復元することに成功しました。ESM-Cモデルを用いたタンパク質生成では、従来の順次復号と比較して、推論時の前向きパス数を3~5倍削減し、生成品質を維持しました。さらに、エントロピーベースの並列化ベースライン手法と比較しても、生成品質で優れた結果を示しました。
本研究の成果は2026年1月にICML 2026に投稿され、関連コードとデータセットは公開予定です。この研究は、マスク拡散モデルの効率的な復号に新たな道を開くだけでなく、モデルの解釈可能性を高め、自然言語処理、バイオインフォマティクス、組合せ最適化など、より広範な離散系列生成タスクでの応用が期待されます。