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

Abstract: Autoregressive (AR) language modeling is the dominant paradigm for text generation, yet its sequential token-by-token decoding makes inference memory-bound and inefficient. Existing acceleration approaches, such as speculative decoding and diffusion LLMs, can yield speedups under certain conditions but do not directly address high-load batch serving--the scenario most critical for industrial-scale deployment. We introduce K-Forcing, a push-forward language modeling paradigm for joint next-k-token decoding. K-Forcing distills an existing AR model into a conditional push-forward mapping--one that transforms independent uniform noise variables into a joint sample of multiple future tokens in a single forward pass. This design preserves fixed-length outputs, reuses the AR teacher backbone, and remains compatible with standard AR serving infrastructure. We train this mapping via progressive self-forcing distillation, which gradually expands the prediction window while enabling the student to closely match the sequence distribution of the AR teacher. We evaluate K-Forcing on LM1B and OpenWebText using a standard causal Transformer backbone. When aggressively configured to generate k = 4 tokens per forward pass, K-Forcing delivers approximately 2.4-3.5x speedup across different batch sizes, while incurring modest quality degradation relative to its AR teacher. As inference increasingly dominates the lifetime compute cost of modern LLMs, K-Forcing offers a promising route toward accelerating AR generation under real-world high-load deployment.

• The paper introduces a push-forward technique that jointly decodes multiple tokens in a single forward pass, drastically reducing inference time.
• It employs progressive self-forcing distillation to overcome noise inversion issues and accurately model the AR teacher’s joint token distribution.
• Empirical evaluations on LM1B and OpenWebText demonstrate up to 3.5x speedup while maintaining batch regularity and with modest quality degradation.

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

Motivation and Background

Autoregressive (AR) LLMs have achieved state-of-the-art scores in natural language generation and understanding but remain bottlenecked by sequential, single-token generation. Each output token requires a separate forward pass, making AR decoding memory-bound, particularly in high-load, large-batch inference settings on hardware accelerators. Prior attempts to accelerate AR inference—speculative decoding and various forms of parallel, masked, or diffusion-based decoding—either break batching regularity, degrade sample quality, or fail to capture the joint conditional distribution of multiple tokens, especially in the high-throughput batch-serving regime critical for industrial deployment.

The K-Forcing framework addresses the statistical modeling bottleneck by enabling joint next-k-token decoding in a single forward pass, thus amortizing the cost across multiple generated outputs and preserving batch regularity.

Figure 1: Comparison of multiple LLM inference paradigms highlighting how K-Forcing achieves multi-token joint decoding per single forward evaluation, as opposed to sequential AR or speculative/MDLM approaches.

The Push-Forward Language Modeling Paradigm

The K-Forcing proposal consists of a push-forward mapping that generates $k$ tokens jointly given a context and $k$ i.i.d. uniform noise variables. Formally, for a context $x_{\leq t}$, the model implements a deterministic map $G_\theta(x_{\leq t}, \mathbf{z})$ producing a length-$k$ block of tokens. This construction is operationally equivalent to sampling from the joint future-token distribution that would be generated by $k$ sequential AR sampling steps, but is collapsed into a single forward pass.

Significantly, K-Forcing does not factor the output token block into conditionally independent per-position predictions, as is typical in diffusion or masked models. Instead, it learns an implicit mapping from noise to tokens so that the composite distribution matches the teacher AR model’s joint, offering a statistically sound multi-token sampler.

Limitations of Existing Multi-Token Decoding Approaches

Draft-then-verify (e.g., speculative decoding) techniques are lossless in isolation but have inherent drawbacks for batch serving: their variable output lengths break synchronized KV-caching and batch alignment, reducing effective throughput under high-load. Diffusion-based approaches (e.g., MDLMs) and standard multi-token predictors (MTP) independently sample per-position marginals in parallel, which provably cannot recover the full joint distribution except under restrictive data assumptions. Theoretical analysis establishes an NFE (number of forward evaluations) lower bound: lossless decoding from marginal predictors cannot require fewer steps than standard AR autoregression in the worst case.

K-Forcing directly targets both issues: it (1) produces fixed-length outputs per step for batching regularity, and (2) models the true joint conditional distribution for blocks of future tokens, not just marginals.

Progressive Self-Forcing Distillation

Training the push-forward model is nontrivial. Simple noise inversion—associating tokens back to noise values via AR model inverse CDF—introduces distributional mismatches and severe numerical instabilities due to narrow CDF bins in the tail of the teacher’s distribution. K-Forcing instead employs progressive self-forcing distillation:

1. Bootstrap: Distill the AR teacher into a PFLM ($k=1$) using forward sampling—teacher generates tokens conditioned on sampled noise, supervising the student.
2. Progressive Window Expansion: Sequentially double the window. A PFLM($k$) generates $k$ tokens; the extended model PFLM($2k$) is trained by concatenating two such generations, supervising on larger noise blocks.
3. Self-Rollout: At each stage, the teacher for supervision is itself a previously trained PFLM, which sidesteps the train-inference mismatch and avoids the instability of noise inversion.

This approach ensures the model always learns to map genuinely uniform noise—reflective of inference-time distribution—to token blocks, while maintaining numerically stable supervision.

Architectural Choices

K-Forcing can be built atop standard AR Transformers and supports two variants:

• Standard MTP: Concatenate $k$0 noise variables, decode with $k$1 independent heads from their joint hidden state.
• Fully Causal: Represent each noise variable as a dedicated token under causal attention, decoded by a shared head, mimicking the recursive structure of AR block sampling and facilitating parameter sharing for multi-token inference.

  Figure 2: Comparison of the standard MTP architecture (left) vs. the fully causal K-Forcing architecture (right), both extending AR backbones with noise injection.

The fully causal design provides a more direct inductive bias for the push-forward mapping and is empirically superior in matching teacher distributions.

Experimental Results

Quality—Throughput Trade-off

Evaluations on LM1B and OpenWebText with a 12-layer, 100M-parameter Transformer backbone demonstrate that K-Forcing achieves $k$2–$k$3 speedups over AR inference, generating up to 4 tokens per pass, with modest quality degradation (as measured by LLM-based win rates and generative perplexity). Varying $k$4 provides a smooth control trade-off between speed and sample coherence: $k$5 nearly matches the AR teacher in human and LLM-based evaluations, while $k$6 achieves the highest throughput with a tolerable decrease in win rate.

Supervision and Architectural Ablations

Ablations reveal that both the supervision strategy and architecture are critical. Only the progressive self-forcing + fully causal configuration reliably achieves low per-position negative log-likelihood, often outperforming both simple noise inversion schemes and non-causal architectures. The per-position NLL remains above zero, indicating that some distributional gap to the AR teacher persists, attributable to numerical noise, occasional mismatches, or capacity limitations.

Figure 3: Training dynamics (validation NLL) of K-Forcing on LM1B, comparing supervision/architecture strategies, highlighting faster convergence and lower loss with progressive self-forcing.

Comparative Analysis with Other Paradigms

Against Medusa (speculative), PTP (parallel token prediction), and MDLM, K-Forcing consistently achieves the lowest number of forward evaluations (NFE) for comparable or better output qualities. Speculative approaches require two forward passes (draft + verification); MDLMs collapse in quality when unmasking multiple tokens per pass, consistent with theoretical analysis. In contrast, K-Forcing's implicit push-forward mapping enables a strictly lower NFE bound while preserving batch-serving regularity.

Implementation and Practical Considerations

K-Forcing preserves standard KV-cache and batching paradigms, making it suitable for integration into production LLM serving stacks. The training cost is dominated by the need for two sequential teacher passes and one student pass per context; while the theoretical overhead is only $k$7, the current implementation incurs $k$8 due to dense mask usage. Custom block-sparse kernels, as well as batch-invariant floating-point primitives (as in DeepSeek-V4), are crucial for scaling to larger windows and minimizing teacher-student distributional drift due to kernel non-determinism.

Implications, Limitations, and Future Directions

K-Forcing establishes that statistical modeling-level changes—as opposed to purely system or kernel-level optimizations—can produce substantial batch-serving speedups without incurring lossless sample degradation, provided the full joint conditional is learned robustly. The main remaining challenge is to close the gap to AR quality for larger $k$9, which may require improved consistency in training signals, larger-capacity backbones, and more efficient custom kernels.

Practically, K-Forcing offers a promising route to improving inference economics in high-throughput, latency-constrained environments where AR decoding is currently rate-limiting. Theoretically, it affirms the feasibility of amortized, blockwise sequence modeling using implicit push-forward maps.

Directions for Further Research

• Efficient Block-Sparse Kernels: Engineering custom attention and training kernels to reduce runtime to $x_{\leq t}$0 cost for large windows is a priority for scalability.
• Deterministic Supervision: Further development of batch-invariant, numerically stable GPU primitives for logit computation and attention.
• Teacher-Free Training: Investigating the training of push-forward LMs ab initio, without AR teacher supervision, or via joint objectives, remains an open and challenging research question with implications for both comprehension and generation tractability.
• Temperature Conditioning and Diversity Control: As shown, K-Forcing supports continuous temperature control through explicit conditioning; systematic investigation of entropy modulations and sample calibration is warranted.

Conclusion

K-Forcing introduces a statistically principled framework to perform joint blockwise decoding for LLMs via implicit push-forward mappings, leveraging progressive self-forcing distillation to bridge teacher-student gaps. The result is a practical, throughput-optimized inference engine compatible with standard AR infrastructures but producing multi-token outputs per forward pass with controllable trade-offs. Closing the residual sample quality gap and generalizing the approach to larger models and more diverse decoding regimes constitute essential future frontiers.

(2606.10820)