On the Identifiability of Masked Prediction: Mode Blindness and Mask Schedules

Technical Summary: On the Identifiability of Masked Prediction

Problem Statement
Masked prediction methods (e.g., BERT, masked autoencoders, discrete diffusion) learn representations by optimizing a schedule-weighted family of conditional distributions rather than the joint data law directly. While rapid-mixing regimes allow for the recovery of the joint law from conditionals via approximate tensorization, it remains unclear when near-optimal conditional prediction pins down the underlying joint law in the presence of multimodal data. Specifically, the paper investigates whether a model can achieve near-optimal masked prediction loss while assigning substantially different weights to well-separated global modes (e.g., code vs. prose, or different languages) in the data distribution. This phenomenon is termed mode blindness.

Methodology
The authors analyze the problem at the population level using a novel theoretical framework:

1. $\epsilon$-Identifiability Modulus: They introduce a modulus $A^\mu_d(p, \epsilon)$, defined as the maximum distributional error (e.g., Total Variation distance) between the true data law $p$ and any admissible model law $q$ that incurs a masked prediction risk of at most $\epsilon$. This unifies the study of exact identifiability ($\epsilon=0$) and stability under small excess risk.
2. Mode-Reweighting Witness Family: To probe identifiability, they construct a family of laws $\{q_\lambda\}$ that share the exact same within-mode conditional distributions as the true data law $p$ but differ only in the global mixing proportions (mode weights). If the masked objective cannot distinguish between $p$ and $q_\lambda$, the joint law is not identifiable.
3. Information Decomposition: The paper derives an exact decomposition of the masked log discrepancy between $p$ and a reweighted law $q_\lambda$. The discrepancy is shown to be proportional to the Kullback-Leibler divergence between the true and reweighted mode weights, scaled by the schedule-averaged residual mode uncertainty (measured by the conditional Minimum Mean Squared Error, MMSE, of the mode given the visible context).
4. Assumptions on Data Geometry: The analysis relies on two structural assumptions about the data law:
   • Large-Context Mode Pinning (Assumption 5.1): For sufficiently large visible contexts, the mode is identified with exponentially high probability (residual uncertainty decays exponentially).
   • Low-Visibility Residual Uncertainty (Assumption 5.2): For small visible contexts, the mode remains ambiguous, and the residual uncertainty is bounded below.

Key Contributions

1. Proof of Mode Blindness: Under the assumption of large-context mode pinning, the paper proves that if a mask schedule is dominated by large visible contexts (i.e., the minimum number of visible coordinates is large), the masked objective becomes exponentially insensitive to global mode weights. A model can assign arbitrary weights to the modes (within a constant range) while incurring an exponentially small masked loss. Consequently, the $\epsilon$-identifiability modulus remains macroscopic even at exponentially small excess risk.
2. Schedule-Dependent Identifiability: The paper demonstrates that the mask schedule acts as the "observation model" for the objective.
   • Blindness: Schedules with no mass on low-visibility contexts (e.g., fixed-ratio masking with high visibility) fail to recover mode weights.
   • Recovery: Schedules that assign positive mass to low-visibility contexts (small visible sets) restore sensitivity to mode weights. The strength of recovery is governed by the mass assigned to these low-visibility contexts and the residual uncertainty at those scales.
   • Full-Mask Coercivity: The presence of any positive mass on the full-mask (zero visible context) provides a uniform guarantee that bounds the identifiability modulus over all admissible models, not just the reweighting family.
3. Exact Information Decomposition: The authors provide an exact decomposition showing that the masked discrepancy for reweighted laws is the product of the mode-weight information and the expected residual mode uncertainty. This explains why large contexts suppress sensitivity (pinning reduces uncertainty to near zero) and why low-visibility contexts are necessary for recovery.
4. Lower Bounds on Approximate Tensorization: The paper shows that the failure of joint-law recovery in the bimodal regime forces the approximate-tensorization constants (which bound joint KL by conditional KL) to grow exponentially with the sequence length $N$ for large-context schedules. This indicates the failure is intrinsic to the objective, not a limitation of existing analysis techniques.

Results

• Theoretical Bounds: The paper establishes that for schedules with $v_{\min}(\mu) \geq s_{\text{pin}}$, the discrepancy between the true law and a reweighted law is bounded by $O(e^{-c \cdot v_{\min}(\mu)})$, while the Total Variation distance remains constant. Conversely, for schedules with low-visibility mass $\pi_s(\mu) > 0$, the mode-weight error is bounded by $O(\sqrt{\epsilon / \pi_s(\mu)})$.
• Empirical Verification:
  • Exact Enumeration: On computable bimodal laws (constructed to satisfy the pinning assumptions), the authors verify the predicted exponential decay rates of the discrepancy and the exact sensitivity ratios.
  • Gradient Training: Training a model on these laws confirms that the "blindness–recovery" dichotomy survives stochastic optimization. Models trained with blind schedules fail to converge to the true mode weight, while those with low-visibility mass recover it, with error scaling as $\pi_0(\mu)^{-1/2}$.
  • Real Corpora: Measurements on code/prose and German/English corpora show that natural text exhibits the required geometric structure: residual mode uncertainty is bounded below at short contexts and suppressed at long contexts. However, the decay follows a power law rather than a strict exponential, suggesting that while mode blindness exists, it may be milder (polynomial vs. exponential) in natural language than in idealized thermodynamic systems.

Significance and Claims
The paper claims to resolve a foundational question in masked prediction: when does near-optimal masked prediction identify the underlying joint law? The answer is shown to be schedule-dependent.

• The joint law's global structure (specifically mode weights) is recoverable or invisible depending solely on the mask schedule's mass on low-visibility and full-mask contexts.
• The failure to identify global structure in standard large-context schedules is not a model capacity issue but an intrinsic property of the objective when applied to multimodal data with metastable regimes.
• The work provides a theoretical basis for understanding why conditionals extracted from masked models can be mutually inconsistent (failing to define a single joint law) and suggests that low-visibility mass (or full-mask mass) is a necessary design component for recovering global data structure.
• The authors emphasize that their results are population-level guarantees; while gradient training experiments reproduce the phenomena, the paper does not claim to solve finite-sample or architectural limitations of current large-scale pretraining.

In summary, the paper establishes that the mask schedule determines not just how masked prediction learns, but what it can identify, with global mode weights being invisible to schedules that lack low-visibility or full-mask components in the presence of well-separated data modes.