Balance aware method speeds up parallel language model token generation
Balance-Aware Pivot Selection for Parallel Diffusion LLM Decoding
Computation and Language
Summary
Standard language models generate words one by one, which can be slow. Some new models generate many tokens at once but often pick tokens based on how confident the model is about each one. The authors noticed this can miss the best order for picking which words to generate next. They created a method that picks middle words first in groups of unknown words, making the overall process faster and more balanced. They tested this method on tasks like math problems and coding, showing it keeps quality while speeding up generation.
What this means in practice
- •For natural language processing engineers: Speed up language model generation in applications needing fast response times by using balanced pivot token decoding.
- •For software developers building code generation tools: Improve efficiency of code generation by using a pivot-based token decoding method that reduces dependencies and speeds output.
Authors
Yejin Lee, Hyundong Jin, Yo-Sub Han
Abstract
Diffusion language models enable parallel token generation, offering a promising alternative to sequential autoregressive decoding. Existing acceleration methods use predictive confidence to enable parallel token commitment while preserving generation quality. However, confidence-based selection prioritizes reliable token predictions without explicitly considering how token positions shape the remaining decoding schedule. Moreover, such ordering does not always reflect inter-token dependencies, suggesting that confidence need not be the sole criterion for token selection. Motivated by this limitation, we propose Balance-Aware Pivot Decoding (BPD), which favors central pivots within contiguous masked intervals and decodes the selected positions in parallel. Balanced pivots yield logarithmic decoding depth in the interval length. Our practical method accounts for model uncertainty by combining confidence-based prefilling with confidence-thresholded center selection. We evaluate the resulting quality-efficiency tradeoff on LLaDA and Dream across five mathematical reasoning, code generation, and instruction-following benchmarks. Our implementation is available at: https://anonymous.4open.science/r/DLM_random_pivot-EE62/