Neural Estimation of Pairwise Mutual Information in Masked Discrete Sequence Models

2026-05-22 · Source: cs.LG updates on arXiv.org · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Life Sciences & Biology · Depth: Expert, long

Summary

A neural framework is proposed for estimating pairwise conditional mutual information (MI) directly from the hidden states of pretrained masked diffusion models (MDMs). This framework uses ground-truth MI, computed from the MDM'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. This enables MI-guided parallel decoding, identifying conditionally independent variable subsets. The approach was evaluated on Sudoku and protein sequence generation with ESM-C, demonstrating a 3-5x magnitude reduction in inference-time forward passes compared to sequential decoding. It preserves generative quality and outperforms entropy-based parallelization methods. For Sudoku, a model with 4,158,346 parameters and a 99,969-parameter MI predictor head was trained over 10 epochs. For proteins, ESM-Cambrian (300M parameters) was used with an ~810K-parameter predictor head trained over 10000 proteins and 5 epochs.

Key takeaway

For machine learning engineers optimizing discrete sequence generation with masked diffusion models, you should consider integrating MI-guided parallel sampling. This method significantly reduces inference-time forward passes by 3-5x while maintaining generative quality, outperforming simpler entropy-based approaches. Implement a neural MI estimator to dynamically identify conditionally independent tokens, enabling more efficient and consistent parallel decoding in applications like protein design or structured data generation.

Key insights

A neural estimator can predict pairwise mutual information from MDM hidden states for efficient, interpretable parallel decoding.

Principles

• MDMs implicitly learn relational structures.
• MI-guided decoding maintains consistency.
• Ground truth MI can supervise neural estimators.

Method

Train a lightweight neural estimator on MDM hidden states to predict pairwise MI, using brute-force conditional probing to generate ground truth MI matrices for supervision.

In practice

• Apply MI maps to interpret model dependencies.
• Use MI-guided sampling for parallel decoding.
• Reduce inference passes in discrete sequence generation.

Topics

• Mutual Information Estimation
• Masked Diffusion Models
• Parallel Decoding
• Protein Sequence Generation
• Sudoku Solving
• Generative AI

Best for: AI Scientist, Machine Learning Engineer, Research Scientist

Related on AIssential

• DepCap: Adaptive Block-Wise Parallel Decoding for Efficient Diffusion LM Inference — DepCap optimizes DLM inference by adaptively determining block boundaries and enabling conflict-free parallel decoding.
• Who Should Lead Decoding Now? Tracking Reliable Trajectories for Ensembling Masked Diffusion Language Models — TIE enhances MDLM sequence generation by iteratively ensembling models based on reliable decoding trajectory confidence.
• Re-evaluating Confidence Remasking in Masked Diffusion Language Models — Post-hoc confidence-based remasking benefits in masked diffusion language models are highly setting-dependent.
• Saber: An Efficient Sampling with Adaptive Acceleration and Backtracking Enhanced Remasking for Diffusion Language Model — Saber improves Diffusion Language Models' code generation by adaptively accelerating and backtracking to correct errors.
• Forward-Learned Discrete Diffusion: Learning how to noise to denoise faster — Learning the forward noising process in discrete diffusion enables few-step, high-quality generation.
• Beyond Fully Random Masking: Attention-Guided Denoising and Optimization for Diffusion Language Models — Attention-guided denoising and optimization (AGDO) significantly enhances dLLM reasoning by leveraging intrinsic token dependencies.

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by cs.LG updates on arXiv.org.