IB-Flow: Information Bottleneck-Guided CFG Distillation for Few-Step Text-to-Image Generation

Abstract: While large-scale text-to-image generative models have achieved unprecedented visual performance, their inherent reliance on multi-step iterative solvers incurs severe inference latency. Few-step distillation targeting the Classifier-Free Guidance (CFG) trajectory has emerged as the prevalent dual-dimensional compression paradigm. However, existing frameworks remain subjugated by a coarse-grained blind injection paradigm that perpetually enforces a globally static guidance strength while indiscriminately sampling the supervisor timestep. This state-agnostic design completely disregards the intrinsic nature of image generation as a dynamic evolutionary process characterized by progressive entropy reduction, which not only restricts the performance boundary of few-step compression but also precipitates severe CFG over-conditioning artifacts. To transcend these limitations, we re-examine the distillation procedure through the theoretical lens of Information Theory, formally modeling it as a dynamic mutual information game constrained by the Information Bottleneck (IB) principle. Specifically, we dismantle traditional blind assumptions via a dual-track adaptive framework. To determine the injection target, we propose an instance-aware selection mechanism that transmutes the intractable KL divergence constraint into a zero-overhead closed-form solution predicated on the local vector field norm. To regulate the injection strength, we introduce an entropy-aware schedule that dynamically decays alongside the SNR, applying maximal thrust for initial structural anchoring before smoothly reverting to the natural manifold to refine micro-details. Extensive empirical evaluations corroborate that our framework fundamentally eradicates over-conditioning artifacts, shattering the performance ceiling to achieve SOTA generative fidelity under extremely stringent 2-step configurations.

• The paper introduces two zero-overhead schedules that adapt the CFG supervisor timestep to the local residual norm and reduce guidance strength as signal-to-noise ratio rises.
• IB-Flow achieves state-of-the-art 2-NFE results across FLUX.1-dev, OpenUniL-L-512, and Qwen-Image-20B, with IBFlow-Qwen reaching 88.67 DPG-Bench versus 88.32 for its 50-step teacher.
• The method improves global layout, texture, and text alignment, but its theoretical schedules remain sensitive to the Fisher approximation, hyperparameters, and the limitations of the teacher distribution.

IB-Flow addresses a specific inefficiency in few-step text-to-image distillation: existing CFG-distillation frameworks inject teacher guidance with a globally static guidance scale $\omega$ and a randomly sampled supervisor timestep $\tau_{CA}$, regardless of the student's current position along the entropy-reduction trajectory. The authors recast this injection decision as an Information Bottleneck (IB) problem and derive two closed-form, zero-overhead schedules — one for where to inject (the supervisor timestep) and one for how strongly to inject (the guidance scale). Applied to FLUX.1-dev, OpenUni-L-512, and Qwen-Image-20B, the resulting 2-NFE students surpass prior state-of-the-art distillations on GenEval, DPG-Bench, and OneIG-Bench.

Motivation: the blind injection paradigm

The paper builds on Decoupled DMD's decomposition of the distillation gradient into a Distribution Matching regularizer ("shield") and a CFG Augmentation engine ("spear"), where the CA term is governed by two control knobs: the guidance strength $\omega$ and the supervisor timestep $\tau_{CA}$. Prior implementations fix $\omega$ across all timesteps and sample $\tau_{CA}$ from an empirical prior. The authors argue this is doubly misaligned with generation as a progressive entropy-reduction process: early high-entropy states require conservative structural anchoring (a distant target creates a semantic gap that destabilizes distillation), while late low-entropy states require retracted guidance (rigid conditioning produces over-conditioning artifacts such as color oversaturation and texture sharpening). The contribution is to make both knobs adaptive, per-instance and per-timestep, without adding inference cost.

IB formulation of the injection target

Selecting $\tau_{CA}$ is formalized as maximizing the conditional mutual information $I(X_{\tau_{CA}}; X_1 \mid X_t)$ subject to a KL divergence budget $D_{KL}(p_t \| p_{\tau_{CA}}) \le \delta$, where $\delta$ acts as an instantaneous channel capacity. Because predictive information about $\tau_{CA}$0 increases monotonically with $\tau_{CA}$1 under flow matching, the optimum lies at the largest admissible $\tau_{CA}$2. The key tractability result is a Fisher-information approximation showing that the local KL divergence satisfies $\tau_{CA}$3, justified by the identity that the CFG residual is proportional to the spatial gradient of pointwise mutual information between $\tau_{CA}$4 and the condition. This yields the closed-form solution:

$\tau_{CA}$5

The mechanism is instance-aware: samples with large CFG residuals (sharp conditional gradients) receive nearby supervision targets, while smoother states permit longer injection strides. Since both vector fields are already computed during CFG distillation, the schedule adds no overhead. It should be noted that the derivation relies on a second-order Taylor/Fisher expansion whose validity degrades for large injection spans; the ablation on $\tau_{CA}$6 confirms this empirically, with performance collapsing at $\tau_{CA}$7 (DPG drops to 86.13 versus 88.67 at the chosen $\tau_{CA}$8).

Entropy-aware guidance strength

For the strength knob, the authors minimize an IB objective trading conditioning cost $\tau_{CA}$9 against generative fidelity $\omega$0. Via the chain rule of mutual information and the I-MMSE relation, they show the achievable fidelity gain from raising $\omega$1 is bounded by the residual entropy $\omega$2, which under the rectified-flow path scales inversely with the instantaneous SNR $\omega$3. Setting the supplementary drive $\omega$4 proportional to this residual uncertainty gives:

$\omega$5

This guarantees the boundary condition $\omega$6 as $\omega$7: late-stage guidance smoothly retracts to the unconditional manifold, eliminating over-conditioning artifacts at their source rather than post hoc. The proof assumes a locally Gaussian channel model and proportionality up to logarithmic factors, so the schedule is analytically motivated rather than exact; nevertheless it outperforms constant, linear-decay, and cosine-decay alternatives in ablation (DPG 88.67 vs. 88.37 / 87.82 / 88.16).

Empirical results

Training uses 32 A100 GPUs for roughly 1.5 days per run on a 2.3M-prompt corpus, with $\omega$8, $\omega$9, and $\tau_{CA}$0 inherited from each teacher. Headline results at 2 NFE:

Model	NFE	GenEval	DPG-Bench
Qwen-Image-20B (teacher)	50×2	0.87	88.32
ArcFlow (Qwen)	2	0.84	87.94
TwinFlow (Qwen)	2	0.82	87.01
pi-Flow (GM-Qwen)	2	0.83	86.45
IBFlow-Qwen	2	0.86	88.67
FLUX.1-dev (teacher)	50	0.66	84.16
IBFlow-FLUX	2	0.68	85.13

Two claims stand out. First, IBFlow-Qwen exceeds its own multi-step teacher on DPG-Bench (88.67 vs. 88.32), indicating the dynamic schedules do more than preserve teacher quality under compression. Second, on OneIG-Bench the method improves the Text axis substantially over ArcFlow (0.465 vs. 0.368 on FLUX; 0.863 vs. 0.789 on Qwen), though it concedes some Style score relative to ArcFlow on FLUX (0.312 vs. 0.350). A fine-grained DPG breakdown attributes the largest gains to Global (+1.22) and Texture (+1.09) axes, matching the intended roles of the two schedules: $\tau_{CA}$1 anchors global layout early, while the SNR decay removes late-stage over-sharpening. Component ablations show each schedule contributes independently ($\tau_{CA}$2 alone: 88.37 DPG; $\tau_{CA}$3 alone: 88.50; combined: 88.67), and qualitative comparisons confirm that removing the dynamic target damages structure and alignment whereas removing the dynamic strength damages color and texture.

Limitations and open questions

The authors are explicit about three constraints. As a distillation framework, generative capacity remains upper-bounded by the teacher's distribution and training corpus; the method cannot exceed what the teacher has mastered. The closed-form derivations are specific to Flow Matching and Rectified Flow paths; extending them to SDE-based diffusion with non-linear variance schedules would require additional approximations. Methodologically, the Fisher-information bound underlying $\tau_{CA}$4 is only locally valid, making the divergence budget $\tau_{CA}$5 a sensitive hyperparameter, and the SNR-based strength schedule rests on a Gaussian-channel idealization. An open question the results raise but do not resolve is whether instance-adaptive supervision generalizes to autoregressive or video-generation settings where the entropy trajectory is less well characterized by a scalar SNR.

Conclusion

IB-Flow replaces the static, blind injection paradigm of CFG distillation with two theoretically derived, inference-free schedules: an IB-constrained supervisor timestep selected via the local CFG-residual norm, and an SNR-driven guidance decay that provably retracts to unconditional sampling as $\tau_{CA}$6. Across three backbones at 2 NFE, the approach achieves state-of-the-art fidelity and text alignment, including surpassing its multi-step teacher on DPG-Bench, with ablations confirming the complementary contributions of the target and strength schedules.