Fine-tuning language models improves context use and prediction accuracy
Fine-Tuning Diffusion Language Models with Context Selection and Target Weighting
Artificial IntelligenceComputation and Language
Summary
Training language models to predict missing words depends on which words they see and which they guess. The authors propose GoldiMask, a method that smartly chooses which words to show as clues and weighs the remaining words to guess based on how helpful those clues are. This method helps models get better at reasoning and writing code by focusing training on more useful examples. Tests show GoldiMask improves accuracy and can reduce the time needed for generating answers while keeping quality high.
What this means in practice
- •For language model developers: Enhance training methods to improve model accuracy on reasoning and code tasks by selectively choosing context tokens during fine-tuning.
- •For software engineering teams: Reduce time to generate high-confidence code completions and mathematical reasoning responses in AI-assisted development tools.
Authors
Loay Mualem, Lluís Pastor-Pérez, Vinh Tong, Andrei Manolache, Tanja Bien, Steffen Staab, Mathias Niepert
Abstract
Supervised fine-tuning of discrete diffusion language models masks some response tokens and trains the model to recover their original values from the visible context. The masking pattern therefore determines both the context available to the model and the tokens it learns to predict. Uniform random masking does not explicitly account for the interaction between these choices. We introduce GoldiMask, which selects tokens to reveal as context by approximately maximizing a submodular objective. This objective uses model signals to balance the benefit of revealing tokens against their value as prediction targets. GoldiMask then weights the remaining targets according to how they benefit from the selected context and their remaining learning potential. Across three backbones and three training datasets, GoldiMask achieves the highest average accuracy in most evaluated settings, demonstrating gains on both reasoning and code generation. Component ablations show that both context selection and target weighting contribute to the gains. GoldiMask also reduces decoding iterations on GSM8K and MATH-500 under confidence-threshold parallel decoding, while maintaining comparable accuracy at higher confidence thresholds.