PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving Few-Step Language Flows with Untied Self-Conditioning

Bocheng Li, Linli Xu

arXiv:2608.22244v1cs.CLcs.AIcs.LG

TL;DR

Few-step flow-matching language generation suffers from a train–inference mismatch in previous-prediction self-conditioning, which also distorts the solver update. The paper introduces Untied Self-Conditioning, a training-free sampler that corrects both effects from frozen model structure and prediction history, improving generation quality across tested settings. At eight evaluations, LangFlow OpenWebText GenPPL drops from 531 to 62, while ELF-B drops from 71 to 43.

  • Problem

    Previous-prediction self-conditioning couples the latent and explicit input pathways during sampling, although training computes them without an intervening solver step.

  • Method

    Untied Self-Conditioning dampens redundant self-conditioning directions and shifts solver predictions toward step averages using prediction history and offline trajectory statistics.

  • Results

    531 to 62: LangFlow OpenWebText GenPPL at eight evaluations; ELF-B GenPPL falls from 71 to 43, with improvements at every tested NFE.

  • Takeaways & Limitations

    The sampler provides quality improvements in few-step generation without retraining and with one model evaluation per step.

  • Takeaways & Limitations

    Comparisons use matched configurations with fixed pretrained models, solvers, time grids, tokenizers, decoding procedures, and shared overlap statistics.

Abstract

from arXiv · show

Flow-matching language models refine all token positions in parallel and can trade sampling steps for latency, yet generation quality still degrades sharply with few sampling steps. We trace a source of this degradation to a train--inference mismatch in previous-prediction self-conditioning: during training, the self-conditioning input is computed from the current noisy state with no intervening solver step; during sampling, the solver folds the previous prediction into the latent before that same prediction reappears as the explicit self-conditioning input. This coupling, absent during training, creates redundancy that grows with step width. We show that the mismatch degrades both the self-conditioning input and the solver update, and derive a correction for each from the model's own structure. From the frozen projection weights we identify directions along which the self-conditioning input is redundant with the latent and dampen them; from the solver's integration structure we derive that a step-average prediction is needed and approximate it from prediction history, with scale set by offline trajectory statistics. The resulting sampler, Untied Self-Conditioning, requires no retraining and uses one evaluation per step. At 8 sampling steps on LangFlow, it reduces OpenWebText generative perplexity from $531$ to~$62$ ($8.6\times$); under an adapted Arena-Hard-Auto~v2 protocol, its outputs are preferred in $96\%$ of pairwise comparisons. On ELF-B it reduces generative perplexity from $71$ to~$43$. Improvements hold from 8 to 256 sampling steps.

1 INTRODUCTION

Flow-matching language models enable parallel, variable-cost text generation, but few-step sampling remains difficult because self-conditioning creates a train–inference mismatch. Untied Self-Conditioning corrects the mismatch without retraining and improves quality across tested settings.

  • Motivation: Fewer iterative steps reduce latency, making each model evaluation more informative the central challenge for few-step generation.Flow-matching language models refine all token positions simultaneously in a continuous latent space.
  • Problem: Solver coupling creates redundancy because the previous prediction is folded into the latent and then reintroduced as explicit self-conditioning, unlike during training.The coupling grows with step width.
  • Method: Untied Self-Conditioning dampens redundant self-conditioning directions using frozen projection weights and shifts solver predictions toward step averages using prediction history.The sampler uses one model evaluation per step and stores the raw prediction for self-conditioning.
  • Results: 531 to 62: LangFlow GenPPL on OpenWebText at eight model evaluations, while ELF-B falls from 71 to 43.At 32 evaluations on ELF-B, GenPPL reaches 21 with 95% throughput retained.
  • Diagnosis: Strong coupling can make self-conditioning harmful even when individual prediction quality is unchanged.A controlled experiment varies coupling while preserving each step’s marginal distribution.
  • Results: Improvements occur at every tested NFE, with the largest gains in the few-step regime.The sampler requires no retraining and uses one model evaluation per step.

2 PRELIMINARIES

Language flow models embed discrete tokens continuously and generate by numerically integrating a learned transport field. Previous-prediction self-conditioning supplies each step with the preceding clean prediction, which is produced through a standard sampling recurrence.

  • Language flows: Language flows embed discrete tokens in a continuous space and learn a velocity field that transports noise to clean data.Sampling integrates the induced ODE over a discretized time grid.
  • Sampling recurrence: At each step, the model predicts a clean state from the latent, self-conditioning input, and time, then the solver updates the latent.The recurrence is pk = fθ(zk, mk, tk), zk+1 = Sk(zk, pk), and mk+1 = pk.
  • Self-conditioning: During sampling, the self-conditioning input is the previous clean prediction, while training computes it from the current noisy state without an intervening solver step.The intervening solver step creates cross-step correlation during inference.

3 SOLVER COUPLING CREATES A TRAIN–INFERENCE MISMATCH

The solver creates a train–inference mismatch by putting the previous prediction into both the latent and explicit self-conditioning pathway. This causes redundancy in self-conditioning and a step-average approximation gap in the solver, especially at low step counts.

  • Mismatch: During sampling, the solver folds the previous prediction into the next latent before passing that same prediction as explicit self-conditioning.Training has no intervening solver step, so this coupling is absent there and grows with step width.
  • Two consequences: The mismatch produces two consequences: redundant self-conditioning input and a solver prediction that differs from the step-average prediction required for accurate integration.The solver gap grows at low step counts.
  • Causal test: Strong coupling makes a fixed positive self-conditioning coefficient harmful after the optimal coefficient decreases monotonically with correlation.Varying correlation preserves both input marginals and changes only their joint dependence.
  • Causal test: The sampling recurrence exposes the previous prediction through both mk and zk, and frozen-language-model experiments confirm the predicted monotone decrease.The solver is the source of the cross-pathway dependence.
  • Projection geometry: Shared row-space directions of the latent and self-conditioning projections identify components where explicit self-conditioning is redundant.Large squared principal-angle cosines indicate strong shared response to the same input component.
  • Projection geometry: The overlap objective is reduced by applying (I−aK), which dampens high-overlap directions while leaving self-conditioning-specific directions less affected.Direction qi retains fraction 1−aρ2_i.
  • Solver integration: The exact solver update needs a weighted average prediction over each step, whereas one start-of-step evaluation supplies only pk.The approximation gap is first order in step width and its latent-update contribution is second order.
  • Solver integration: An exponential moving average of past predictions supplies a history-based direction that estimates recent prediction change without extra evaluations.The method combines this residual with solver-weighted history and normalizes using offline trajectory statistics.

4 UNTIED SELF-CONDITIONING

Untied Self-Conditioning implements separate corrections for the two mismatch consequences within one sampling loop. It dampens redundant self-conditioning directions and uses history-based solver corrections scaled by offline trajectory statistics.

  • Sampler design: A single sampling loop applies both mismatch corrections with one model evaluation per step.The self-conditioning pathway and solver pathway receive different corrected quantities.
  • Self-conditioning correction: The overlap transform Fk retains less of self-conditioning directions with greater projection overlap and leaves unique directions comparatively unchanged.Principal directions and squared cosines are computed once from frozen weights.
  • Self-conditioning correction: Dampening is strongest when the solver has folded more prior prediction into the latent and weakens when raw and corrected running averages disagree.At the first step, no prior prediction exists and the input passes unchanged.
  • Solver correction: Solver-weighted history tracks the predictions that entered the solver through the running average Ok.This distinguishes the corrected solver stream from the raw self-conditioning stream.
  • Solver correction: The correction direction combines the raw-prediction EMA residual with disagreement between the current prediction and the solver-running average.The disagreement also gates self-conditioning dampening strength.
  • Normalization: Offline trajectory statistics replace the current correction direction’s arbitrary mean and scale while preserving its centered direction.The statistics are per-step token means and centered RMS values from stored trajectories.
  • Pathway routing: The raw prediction updates self-conditioning and its EMA, while the corrected prediction advances the solver and its state.Routing the correction into self-conditioning would create a second mismatch.

5 EXPERIMENTS

Experiments evaluate Untied Self-Conditioning across matched model, solver, and decoding settings, alongside controlled tests of its coupling diagnosis and two corrections. The sampler improves quality across tested NFE, with especially large few-step gains and retained diversity or throughput.

  • Generation quality: Untied Self-Conditioning improves generation quality in every tested configuration across ELF and LangFlow.Comparisons keep the pretrained model, base solver, time grid, tokenizer, decoding procedure, initial latent, and random stream matched.
  • Generation quality: 531.2 to 61.6 (8.6×): LangFlow OpenWebText GenPPL at 8 NFE; ELF-B reaches 42.66 versus 70.63.At 32 NFE on ELF-B, GenPPL reaches 21.0 without additional training.
  • Efficiency and diversity: 95.5% geometric-mean throughput is retained across six workloads, while entropy moves closer to the corpus reference on LangFlow and is preserved on ELF.The sampler uses one model evaluation per step in its stated sampling loop.
  • Generation quality: 96.3% of non-tie comparisons favor Untied Self-Conditioning at 8 NFE under pairwise LLM judging.The evaluation uses 1,024 matched generations per NFE and reports an Arena score of 79.2 at 8 NFE.
  • Correlation study: 0.88 to 0.30: the MSE-optimal self-conditioning coefficient decreases as cross-step correlation rises at 8 NFE.At high correlation, coefficient one becomes harmful, with the MSE ratio rising to 1.07 at 8 NFE and 1.11 at 16 NFE.
  • Correction studies: 45.63 versus 67.27: applying the history-based correction in the solver improves GenPPL, while applying it only to self-conditioning raises GenPPL to 157.77.The solver correction provides the larger individual gain, and the complete sampler outperforms either component alone.

6 RELATED WORK

Related work has addressed self-conditioning, solver integration, and offline solver statistics through separate training-time or inference-time changes. This work differs by identifying solver-induced coupling and deriving corrections for both pathways.

  • Self-conditioned language flows: Prior self-conditioned language-flow methods refine all token positions in parallel using previous clean predictions as additional model input.These methods span continuous text-generation models using self-conditioning during sampling.
  • Training-time modifications: Training-time approaches address train–inference mismatch by modifying model training, exposing intermediate states, post-training, or distilling recurrent self-conditioning.These approaches require changes to model weights or training procedures.
  • Inference-time changes: Inference-time methods modify the self-conditioning input independently of the solver update, including moving averages, self-guidance, or repetition-direction subtraction.Some methods use two model evaluations or estimate directions from generated trajectories.
  • Solvers and multi-step methods: Multi-step solvers reuse past evaluations or aggregate past clean predictions to improve numerical integration or current solver predictions.In self-conditioned models, the same prediction can enter both the solver and self-conditioning pathways.
  • Position of this work: Untied Self-Conditioning jointly corrects solver-induced coupling using projection geometry and integration structure, rather than modifying only one sampling component.Its corrections are training-free and target both pathways implicated by the mismatch.

7 CONCLUSION

The paper diagnoses solver-induced coupling as a train–inference mismatch in self-conditioned flow sampling and corrects both resulting mechanisms without retraining. Across ELF and LangFlow, the sampler improves quality at every tested NFE, including a large eight-step LangFlow gain confirmed by pairwise preference.

  • Conclusion: Strong coupling can make a fixed self-conditioning coefficient harmful even when individual prediction quality is unchanged.The controlled experiment varies coupling while preserving marginal distributions.
  • Conclusion: Untied Self-Conditioning derives geometric and step-average corrections from frozen projection weights and solver integration structure.The sampler uses one model evaluation per step and requires no retraining.
  • Conclusion: 531 to 62: LangFlow OpenWebText GenPPL falls at eight evaluations, with gains also supported by pairwise LLM preference and external-evaluator likelihood.Improvements hold across ELF and LangFlow at every tested NFE.

A DERIVATIONS AND PROOFS

The derivations show that solver updates require a weighted step-average prediction rather than only the initial prediction, while the correction error has predictable step-size orders. History-based residuals provide a practical approximation to this average.

  • Solver coefficient structure: In the ELF solver, normalized update coefficients differ from full-latent coefficients, and re-noising changes the evaluation time while preserving that algebraic separation.The first normalized update has κ0 = 1 even though its full-latent coefficient is β0 = t1.
  • Step-size orders: The true correction to the initial prediction is first order in step width, while its contribution to the endpoint is second order.This makes the discrepancy small at many steps but increasingly relevant as steps become wider.
  • Step-size orders: For a candidate correction Bk, the endpoint-error difference is scaled by (1 −e−hk)^2, preserving the sign of the improvement criterion.The factor connects the correction objective to the squared endpoint error.
  • Weighted prediction over a solver step: The solver’s exact update uses a weighted average of predictions across each step, not merely the model evaluation at the step’s start.The weighting follows the affine flow’s variation-of-constants solution and specializes to the ELF and LangFlow solver coordinates.
  • History-based approximation: A history residual can recover the prediction increment exactly for an affine sequence, but converting it into solver-time derivatives also requires grid spacing and prediction curvature.The sampler therefore uses the raw residual as a history-based direction rather than treating it as an exact derivative in solver time.

B METHOD DETAILS

The method uses frozen projection geometry to damp redundant self-conditioning directions and offline trajectory statistics to scale a history-based solver correction. The resulting transform and correction preserve specific structural constraints while targeting complementary information and step-average behavior.

  • Principal-angle directions for self-conditioning: Frozen input-projection weights define principal directions where the latent and self-conditioning pathways overlap, identifying which self-conditioning components should be damped.ELF has a 128-dimensional bottleneck with squared principal-angle cosines ranging near zero to 0.92, whereas LangFlow’s full-rank projections yield isotropic dampening.
  • Principal-angle directions for self-conditioning: The self-conditioning transform preserves the orthogonal complement of the shared subspace, contracts principal directions, and uses a floor to avoid completely suppressing any direction.A separate uniform contraction can be isolated from the direction-dependent dampening.
  • History-based solver correction: The sampler forms a history residual from the current prediction and an exponential moving average, with startup bias decaying as (1−α)^k.A bias-corrected alternative exists, but the sampler uses the raw residual because the bias vanishes after a few steps.
  • Matching the empirical mean and RMS: Offline trajectory banks provide step-dependent mean and RMS statistics that set the correction’s strength, while the current trajectory supplies its centered direction.Inactive positions receive zero correction, and the centered component is chosen to maximize alignment among tensors with the same moments.
  • Matching the empirical mean and RMS: Whenever the proposed correction has positive alignment with the true correction, some positive correction strength strictly improves the endpoint over using pk alone.The endpoint criterion decomposes into alignment of the empirical mean and the angle between centered correction directions.

C.1 EVALUATION PROTOCOL

The evaluation protocol uses matched checkpoints, initial latents, solvers, grids, sequence lengths, random streams, and model-evaluation counts. ELF and LangFlow use their official grids and evaluator conventions, with independent seed runs supporting reported uncertainty.

  • Matched comparisons: Every matched comparison fixes the checkpoint, initial latent, base solver, time grid, sequence length, and number of model evaluations.Configurations also start from the same random stream, while generation and evaluation are run separately.
  • ELF protocol: ELF evaluates official sampling settings across 8, 16, 32, and 64 NFE, including specified SDE noise scales and a self-conditioning scale of 3.0.The protocol uses the official random logit-normal grid.
  • ELF protocol: ELF uses seeds 0–5, with 1,024 generations of length 1,024 per seed; each independent seed run is the unit for reported standard error.This defines the sampling volume and uncertainty unit for ELF results.
  • LangFlow protocol: LangFlow uses its official proposal-quantile time grid and seeds 0–5, generating 1,024 samples per setting with lengths 1,024 on OpenWebText and 128 on LM1B.The grid spans proposal quantiles between 1 −10−5 and 10−5.
  • Evaluation metrics: Official and Untied Self-Conditioning configurations share evaluator conventions for entropy and perplexity within each model comparison.ELF applies its official empty-sample filter before GPT-2 Large perplexity and unigram-entropy computation; LangFlow decodes before perplexity calculation.

C.2 PAIRWISE LLM-JUDGE EVALUATION PROTOCOL

The protocol compares official and Untied Self-Conditioning generations on matched latents through blind, reversed-order LLM judgments. It defines pairwise outcomes from weighted verdicts and reports associated evaluation context.

  • Evaluation protocol: 1,024 matched text pairs per NFE are presented twice in reversed order, yielding 8,192 judgments across the evaluation.The protocol covers NFE values in {8, 16, 32, 64} and uses the same latent and sample index for each pair.
  • Evaluation protocol: The judge sees neither method names nor NFE labels and scores fluency, coherence, semantic content, and degeneration artifacts.Verdicts use five levels ranging from strong preference for one system to strong preference for the other.
  • Outcome aggregation: Pair-level outcomes classify a candidate win when the weighted two-game score exceeds 0.5, an official win below 0.5, and a tie otherwise.Confidence intervals use 10,000 paired bootstrap resamples, with no judgments excluded or reassigned.

C.9 ADDITIONAL ABLATIONS

Additional ablations test correction direction, sampler comparisons, and throughput under matched computation. The reported gains depend on the history-based solver direction while retaining most baseline throughput.

  • Solver-correction controls: A random equal-RMS solver-correction direction nearly recovers official GenPPL, showing that the gain requires the specific history-based direction.Removing the sample-specific component raises GenPPL, while permuting it across samples reduces the gain further.
  • Matched sampler comparison: At 8 NFE on ELF-B, Untied Self-Conditioning reaches 45.63 GenPPL versus 56.90 for HiGS and 74.18 for the official sampler.HiGS combined with ACE reaches 65.84, while the solver correction alone reaches 50.04.
  • Matched sampler comparison: The paired improvement over HiGS with ACE is 0.36 GPT-2 NLL, with 95% interval [0.35, 0.38].The comparison uses matched conditions on frozen ELF-B at 8 NFE with 1,024 generations.
  • Throughput: Throughput ranges from 0.91 to 0.98 of the official sampler, with geometric mean 0.96 and over 95% baseline throughput retained.Both loops use the same compiled model, batch size, NFE, and generated outputs without additional model evaluations.

C.11 PARAMETER SENSITIVITY

The sensitivity and illustrative-comparison results examine robustness and decoded behavior across configurations. Correction performance is generally robust, but the solver coefficient is more sensitive than the self-conditioning coefficients.

  • Parameter sensitivity: All 24 coefficient settings improve GenPPL over the paired official sampler in the frozen ELF-B 8-NFE sensitivity study.The coefficient grid and evaluation protocol were fixed before generation, using 256 independent paired samples.
  • Parameter sensitivity: The ηR sweep varies GenPPL by less than 0.8 across the [×1, ×2] range, while ηO varies by 3.0.The reported ranges are 39.4–40.2 for ηR and 38.4–41.3 for ηO.
  • Parameter sensitivity: Doubling ηT raises GenPPL to 64.3, still below the official sampler’s 65.8, while reducing ηT to ×1⁄2 gives 56.7 at α = 1.The solver coefficient is therefore more sensitive than the ηR and ηO sweeps in this study.
  • Paired likelihoods: At 8 and 16 NFE, all 6,144 paired LangFlow OpenWebText samples have lower NLL under Untied Self-Conditioning.The paired-sample comparison is summarized in Table 15.
  • Illustrative text comparisons: The 64-NFE LangFlow example removes a repeated phrase and sustains article-like segments longer, while the 8-NFE ELF-B example recovers from token-level collapse into connected sentences.These examples use matched initial latents and sample indices.
Loading 2608.22244v1…