Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Saber: An Efficient Sampling with Adaptive Acceleration and Backtracking Enhanced Remasking for Diffusion Language Model

Yihong Dong, Zhaoyu Ma, Xue Jiang, Zhiyuan Fan, Jiaru Qian, Yongmin Li, Jianha Xiao, Zhi Jin, Rongyu Cao, Binhua Li, Fei Huang, Yongbin Li, Ge Li

TL;DR

Diffusion language models offer parallel, bidirectional generation but struggle with the speed-accuracy trade-off in code generation. The paper proposes Saber, a training-free sampling algorithm that combines adaptive acceleration via dynamic unmasking with backtracking-enhanced remasking. Saber speeds up inference by selectively unmasking more tokens as context grows and reverses selected past tokens when new context reveals errors. On five code-generation benchmarks across multiple diffusion language models, Saber achieves an average Pass@1 improvement of 1.9 percentage points and an average speedup of 251.4%, significantly narrowing the gap with autoregressive approaches and demonstrating strong general applicability across DLM implementations.

Abstract

Diffusion language models (DLMs) are emerging as a powerful and promising alternative to the dominant autoregressive paradigm, offering inherent advantages in parallel generation and bidirectional context modeling. However, the performance of DLMs on code generation tasks, which have stronger structural constraints, is significantly hampered by the critical trade-off between inference speed and output quality. We observed that accelerating the code generation process by reducing the number of sampling steps usually leads to a catastrophic collapse in performance. In this paper, we introduce efficient Sampling with Adaptive acceleration and Backtracking Enhanced Remasking (i.e., Saber), a novel training-free sampling algorithm for DLMs to achieve better inference speed and output quality in code generation. Specifically, Saber is motivated by two key insights in the DLM generation process: 1) it can be adaptively accelerated as more of the code context is established; 2) it requires a backtracking mechanism to reverse the generated tokens. Extensive experiments on multiple mainstream code generation benchmarks show that Saber boosts Pass@1 accuracy by an average improvement of 1.9% over mainstream DLM sampling methods, meanwhile achieving an average 251.4% inference speedup. By leveraging the inherent advantages of DLMs, our work significantly narrows the performance gap with autoregressive models in code generation.

Saber: An Efficient Sampling with Adaptive Acceleration and Backtracking Enhanced Remasking for Diffusion Language Model

TL;DR

Diffusion language models offer parallel, bidirectional generation but struggle with the speed-accuracy trade-off in code generation. The paper proposes Saber, a training-free sampling algorithm that combines adaptive acceleration via dynamic unmasking with backtracking-enhanced remasking. Saber speeds up inference by selectively unmasking more tokens as context grows and reverses selected past tokens when new context reveals errors. On five code-generation benchmarks across multiple diffusion language models, Saber achieves an average Pass@1 improvement of 1.9 percentage points and an average speedup of 251.4%, significantly narrowing the gap with autoregressive approaches and demonstrating strong general applicability across DLM implementations.

Abstract

Diffusion language models (DLMs) are emerging as a powerful and promising alternative to the dominant autoregressive paradigm, offering inherent advantages in parallel generation and bidirectional context modeling. However, the performance of DLMs on code generation tasks, which have stronger structural constraints, is significantly hampered by the critical trade-off between inference speed and output quality. We observed that accelerating the code generation process by reducing the number of sampling steps usually leads to a catastrophic collapse in performance. In this paper, we introduce efficient Sampling with Adaptive acceleration and Backtracking Enhanced Remasking (i.e., Saber), a novel training-free sampling algorithm for DLMs to achieve better inference speed and output quality in code generation. Specifically, Saber is motivated by two key insights in the DLM generation process: 1) it can be adaptively accelerated as more of the code context is established; 2) it requires a backtracking mechanism to reverse the generated tokens. Extensive experiments on multiple mainstream code generation benchmarks show that Saber boosts Pass@1 accuracy by an average improvement of 1.9% over mainstream DLM sampling methods, meanwhile achieving an average 251.4% inference speedup. By leveraging the inherent advantages of DLMs, our work significantly narrows the performance gap with autoregressive models in code generation.
Paper Structure (34 sections, 8 equations, 4 figures, 3 tables, 1 algorithm)

This paper contains 34 sections, 8 equations, 4 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Motivation
  3. Insight 1: Difficulty Decreases over DLM Generation Process.
  4. Insight 2: Dynamic Context of DLM Generated Tokens.
  5. Summary.
  6. Related Work
  7. Diffusion Language Models for Code
  8. Efficient DLM Sampling Methods
  9. Saber
  10. Preliminaries
  11. Adaptive Acceleration via Dynamic Unmasking
  12. Backtracking-Enhanced Remasking Mechanism
  13. Overall
  14. Experiment Setup
  15. Datasets
  16. ...and 19 more sections

Figures (4)

  • Figure 1: Left: Illustration of DLM Sampling. Right: The trade-off of DLM Sampling between inference speed and output quality on HumanEval benchmark.
  • Figure 2: Motivation Example. Left: (a) Average confidence per step of DLM sampling. RightL (b) An example of the confidence drop for incorrect tokens of DLM sampling, where gray token means the token generated in the future step.
  • Figure 3: An Overview of Saber in DLM sampling, which consists of two key components, i.e., Adaptive Acceleration via Dynamic Unmasking (AADU) and Backtracking-Enhanced Remasking Mechanism (BERM), during each iterative sampling process.
  • Figure 4: Case Study.