Promises and Pitfalls of Generative Masked Language Modeling: Theoretical Framework and Practical Guidelines (2407.21046v1)

Abstract: Autoregressive LLMs are the currently dominant paradigm for text generation, but they have some fundamental limitations that cannot be remedied by scale-for example inherently sequential and unidirectional generation. While alternate classes of models have been explored, we have limited mathematical understanding of their fundamental power and limitations. In this paper we focus on Generative Masked LLMs (GMLMs), a non-autoregressive paradigm in which we train a model to fit conditional probabilities of the data distribution via masking, which are subsequently used as inputs to a Markov Chain to draw samples from the model, These models empirically strike a promising speed-quality trade-off as each step can be typically parallelized by decoding the entire sequence in parallel. We develop a mathematical framework for analyzing and improving such models which sheds light on questions of sample complexity and inference speed and quality. Empirically, we adapt the T5 model for iteratively-refined parallel decoding, achieving 2-3x speedup in machine translation with minimal sacrifice in quality compared with autoregressive models. We run careful ablation experiments to give recommendations on key design choices, and make fine-grained observations on the common error modes in connection with our theory. Our mathematical analyses and empirical observations characterize both potentials and limitations of this approach, and can be applied to future works on improving understanding and performance of GMLMs. Our codes are released at https://github.com/google-research/google-research/tree/master/padir

• The paper develops a theoretical framework that quantifies sample complexity and speed-quality trade-offs using functional inequalities.
• It demonstrates that adapting the T5 model for iteratively-refined parallel decoding achieves 2-3x speedups in machine translation with minimal quality loss.
• The paper provides practical guidelines, including high masking ratios and model distillation, to optimize GMLM training and mitigate common error modes like stuttering.

Overview of "Promises and Pitfalls of Generative Masked LLMing: Theoretical Framework and Practical Guidelines"

This paper explores Generative Masked LLMs (GMLMs) as an alternative to the prevalent autoregressive (AR) models used in text generation. While AR models predict tokens sequentially, GMLMs enable parallel decoding, offering empirical advantages in speed. The paper introduces a theoretical framework to analyze GMLMs, investigating sample complexity, inference efficiency, and quality.

Key Contributions

1. Theoretical Framework:
   • The paper develops a mathematical approach to assess GMLMs, focusing on sample complexity and speed-quality trade-offs.
   • It examines the asymptotic sample complexity using functional inequalities, such as Poincaré and entropy tensorization.
   • Larger masks in training improve statistical efficiency, with formal proofs supporting this claim.
2. Empirical Studies:
   • The authors adapt the T5 model for iteratively-refined parallel decoding, showing 2-3x speedups in machine translation tasks without significant quality losses.
   • Ablation experiments provide design recommendations for GMLMs and identify common error modes like "stuttering."
3. Practical Guidelines:
   • Suggestions include using high masking ratios, customized vocabularies, distillation from AR models, and positional attention.
   • GMLMs are found effective in translation tasks with low-entropy and less multi-modal outputs.

Theoretical Implications

• Conditional and Joint Distribution Learning: Theoretical analyses indicate scenarios where learning conditional probabilities also aids joint distribution learning.
• Design Space Exploration: The paper offers insights into optimizing losses and training procedures, helping navigate the design space for GMLMs.

Empirical Insights

• The experiments confirm that GMLMs with parallel decoding excel in machine translation.
• Challenges remain in tasks requiring strong target-side dependency modeling.

Future Directions

• The research suggests exploring Markov Chains that mix rapidly in scenarios with strong dependencies.
• Future work may leverage the findings to enhance training objectives, inference algorithms, and model architectures for improved parallel decoding performance.

In conclusion, the paper sheds light on the potential and limitations of GMLMs, balancing between theoretical underpinnings and empirical outcomes, with a prospect of inspiring further advancements in non-autoregressive LLMing.