K-Forcing: Joint Next-K-Token Decoding via Push-Forward Language Modeling

Problem
The paper addresses the inefficiencies of autoregressive (AR) language modeling during inference, particularly in high-load batch serving scenarios critical for industrial applications. Existing acceleration methods, such as speculative decoding and diffusion models, do not adequately resolve the memory-bound and sequential nature of AR decoding. This work is a preprint and has not undergone peer review.

Method
K-Forcing proposes a push-forward language modeling paradigm that enables joint next-k-token decoding. The method distills an existing AR model into a conditional push-forward mapping, which transforms independent uniform noise variables into a joint sample of multiple future tokens in a single forward pass. This approach maintains fixed-length outputs and leverages the AR teacher model, ensuring compatibility with standard AR serving infrastructure. The training process employs progressive self-forcing distillation, gradually expanding the prediction window to allow the student model to closely approximate the sequence distribution of the AR teacher. The architecture utilized for evaluation is a standard causal Transformer backbone.

Results
K-Forcing was evaluated on the LM1B and OpenWebText datasets, achieving a speedup of approximately 2.4-3.5x when configured to generate k = 4 tokens per forward pass across various batch sizes. While this acceleration comes with a modest degradation in quality compared to the AR teacher, the trade-off is deemed acceptable given the significant performance improvements in inference speed.

Limitations
The authors acknowledge that K-Forcing incurs some quality degradation relative to the AR teacher model, although they do not quantify this degradation in detail. Additionally, the method’s performance may vary depending on the specific characteristics of the dataset and the chosen batch size. The paper does not explore the implications of using K-Forcing in low-load scenarios or its adaptability to different model architectures beyond the causal Transformer.

Why it matters
K-Forcing presents a significant advancement in the efficiency of autoregressive language models, particularly for applications requiring high-throughput text generation. By enabling joint next-k-token decoding, it addresses a critical bottleneck in inference speed, which is increasingly important as the compute cost of large language models (LLMs) continues to rise. This work lays the groundwork for future research into optimizing AR models for real-world deployment, potentially influencing the design of next-generation language generation systems. For further details, see the full paper available on arXiv.