用神经网络估算掩码序列模型中的变量依赖关系，加速生成
Neural Estimation of Pairwise Mutual Information in Masked Discrete Sequence Models
- 从预训练模型隐藏状态直接估计成对条件互信息
- 单次前向传播完成全部互信息计算，推理速度提升3-5倍
- 适合需要快速生成且关注变量依赖结构的研究者
理解变量间依赖关系对掩码扩散模型（MDMs）的可解释性和高效生成至关重要，但这类模型主要暴露边缘条件分布，未显式表示变量间依赖。我们提出一种神经框架，直接从预训练MDM的隐藏状态估计成对条件互信息（MI），利用模型自身条件分布计算的真实MI作为监督信号。所得估计器捕捉模型内部对依赖结构的认知，可在一次前向传播中预测完整的MI矩阵，实现基于互信息的并行解码，识别条件独立变量子集。我们在Sudoku和蛋白质序列生成（ESM-C）任务上评估该方法，发现MI图能恢复已知结构约束，在保持生成质量的同时，相比串行解码将推理时间减少3-5倍，优于基于熵的并行化方法。
原文摘要 · Abstract (English)
Understanding dependencies between variables is critical for interpretability and efficient generation in masked diffusion models (MDMs), yet these models primarily expose marginal conditional distributions and do not explicitly represent inter-variable dependence. We propose a neural framework for estimating pairwise conditional mutual information (MI) directly from the hidden states of a pretrained MDM, using ground-truth MI computed from the model's own conditional distributions for supervision. The resulting estimator captures the model's internal belief about dependency structure and predicts the full MI matrix in a single forward pass, enabling MI-guided parallel decoding by identifying conditionally independent subsets of variables. We evaluate our approach on Sudoku and protein sequence generation with ESM-C, where the MI maps recover known structural constraints and enable a 3-5x magnitude reduction in inference-time forward passes compared to sequential decoding, while preserving generative quality and outperforming entropy-based parallelization methods.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。