Source-linked AI summary
FlashControl: One-Step Controllable Generator via Distillation-Friendly Single-Stream Teachers
Ngan Nguyen, Dung Nguyen, Quan Dao, Dimitris N. Metaxas, Minh Hoai, Cuong Pham, Anh Tran
TL;DR
One-step controllable generation is harder because a student must preserve both visual fidelity and spatial faithfulness, while existing dual-branch controllers are difficult to distill. FlashControl uses a single-stream controllable teacher and complementary student-training components, and the paper reports one-step generation on par with traditional two-branch methods.
Problem
One-step controllable generation must preserve visual fidelity and spatial faithfulness in a single network, and dual-branch control pathways are poorly matched to one-step distillation.
Source: §1 · p. 2, §4.1 · p. 4Method
The method adapts selected layers of a single-stream backbone as a controllable teacher, then distills it with DMD using AutoAdapter and Latent Variance Preservation.
Source: §1 · p. 2Results
FlashControl demonstrates one-step spatially controllable generation and performs on par with traditional two-branch methods.
Source: §1 · p. 3, §6 · p. 10Takeaways & Limitations
A streamlined single-branch design can preserve much of the teacher’s controllability while simplifying training and inference.
Source: §6 · p. 10Experiments focus on 512 × 512 images and six common spatial controls; the authors do not claim direct generalization to higher resolutions, video, or specialized domains.
Source: J Limitations · p. 19
Abstract
from arXiv · showhide
Diffusion models have become a leading approach in visual generative modeling, particularly effective in text-to-image and controllable generation. Their main drawback is slow inference as it requires many sampling steps. While recent distillation methods enable one-step text-to-image generation, extending these techniques to controllable generation is still challenging. Existing controllable models, such as ControlNet, rely on dual-branch design that make distillation difficult. In this work, we adopt a simpler strategy. Instead of adding extra conditional branches, we finetune a selected subset of layers in the original text-to-image diffusion model to directly enable controllable generation. This results in DeltaControl, a single-flow, multi-step diffusion model capable of supporting spatial-condition generation while matching the performance of dual-branch approaches. Building on this teacher model, we further apply step-distillation to obtain FlashControl, a one-step model capable of spatial controllable generation with a single forward pass.
1 Introduction
The introduction argues that dual-branch control pathways are poorly suited to one-step distillation, then presents a single-stream teacher and complementary training components for controllable generation.
1 Introduction
- DeltaControl adapts selected backbone layers for spatial control; feed-forward tuning provides a favorable balance across UNet and DiT backbones.The authors interpret feed-forward layers as a more stable control-injection subspace than attention layers.Source: §1 · pp. 2–3
- Vanilla DMD transfers the teacher to a one-step generator while isolating the teacher architecture’s role in the comparison.The main setting avoids extra adversarial components that could confound this comparison.Source: §1 · p. 2
- AutoAdapter aligns heterogeneous condition features to hidden-state statistics and gates tokens to support multiple control types and stabilize early training.The introduction names canny edges, depth maps, segmentation masks, and poses as examples of structurally different controls.Source: §1 · p. 2
- Latent Variance Preservation encourages the student to respond to input-noise perturbations, reducing its tendency toward deterministic collapse.The objective addresses late-stage loss of noise sensitivity during one-step distillation.Source: §1 · p. 2
- The framework assigns distinct roles to its components: DeltaControl addresses teacher architecture, AutoAdapter aligns controls, and Latent Variance Preservation maintains noise responsiveness.The authors describe FlashControl as a pipeline for one-step controllable generation rather than a loose collection of modules.Source: §1 · pp. 2–3
- Naive one-step distillation fails for dual-branch controllable models, while single-flow teachers provide a better foundation for one-step controllable generation.The proposed framework targets control mechanisms that survive compression into a one-step student.Source: §1 · p. 3
2 Related Work
Related work covers diffusion and flow-based generation, methods for spatially controllable synthesis, and approaches that reduce the number of sampling steps.
Diffusion and Flow-based Generative Models
Multi-step Controllable Diffusion Models
- T2I-Adapter [27] trains adapters on a frozen backbone, while ControlNet adds a parallel control branch; these and other methods still rely on multi-step sampling.Single-branch approaches such as UniCombine [45] and OminiControl [44] append conditions to the noisy latent input and use a LoRA-finetuned DiT.Source: §2 · p. 3
3 Preliminaries
The preliminaries describe diffusion and flow-based denoising, ControlNet’s use of a separate control branch, and DMD as the one-step distillation objective.
Diffusion and Flow-based Generative Models
- Diffusion and flow-based generation start from noisy samples and use a learned denoiser or flow field to recover clean data.The passage names the clean image or latent, timestep, and Gaussian noise; α_t and σ_t set signal and noise scales.Source: §3 · pp. 3–4
Preliminaries
- Figure 2 contrasts ControlNet and DeltaControl architectures and shows AutoAdapter’s architecture.The caption labels these as panels (a), (b), and (c), respectively.Source: Figure 2 · p. 4
Diffusion and Flow-based Generative Models
- At inference, repeated denoiser or flow-field evaluations produce samples, but the many evaluations make sampling slow.The model family may predict noise or a velocity/flow target.Source: §3 · p. 4
ControlNet
- ControlNet processes the noisy input, timestep, and spatial condition in an auxiliary branch, then injects condition-dependent residuals into the frozen base denoiser.The passage says zero-initialized layers are usually used for the residual injection.Source: §3 · p. 4
Distribution Matching Distillation
4 FlashControl
FlashControl is a two-stage pipeline for one-step controllable text-to-image generation: DeltaControl integrates spatial control into a single-stream teacher, which is then distilled into FlashControl. Its design addresses dual-branch distillation difficulties and adds mechanisms for adapting heterogeneous controls and preserving student sensitivity to noise.
4.1 Overview
- FlashControl distills a single-stream DeltaControl teacher into a one-step generator using DMD, retaining AutoAdapter and adding LSP to address late-stage collapse.AutoAdapter aligns control perturbations before injection, supporting faster and more stable distillation convergence.Source: §4.1 · p. 5, §4.4 · pp. 7–8
4.2 From Dual-branch to efficient single-stream teacher.
- Distilling dual-branch ControlNet and T2I-Adapter teachers with DMD produces low-quality, blurry results, attributed to the single-step student having to approximate both denoising and control corrections.Freezing either branch does not resolve the problem; optimization can collapse toward the base model or overfit noisy control residuals.Source: §4.2 · p. 5
- DeltaControl integrates control by selectively fine-tuning the base backbone rather than replicating components, analyzing layers across UNet and DiT and supporting diverse modalities.Unlike OminiControl and UniCombine, which target efficient multi-step control, DeltaControl is designed for teacher-to-student transfer with a normalized control interface.Source: §4.2 · pp. 5–6
4.3.1 Efficient Finetuning Strategy.
- FFN tuning gives the largest conditioning improvements versus attention-only tuning and consistently beats ControlNet across UNet and DiT training stages.The study matches ControlNet's trainable parameter count; the authors hypothesize that attention already aligns text and image features while FFNs refine their combined representations.Source: §4.3.1 · p. 6
4.3.2 AutoAdapter.
- AutoAdapter converts control tokens into hidden-state-compatible residuals and gates them token by token, addressing scale, sparsity, and reliability differences across conditions.Statistical alignment and bounded gating are intended to reduce erratic hidden-state updates and suppress noisy or incompatible tokens.Source: §4.3.2 · pp. 6–7
4.4 Stage 2: One-step Controllable Generative Model FlashControl
- Stage 2 uses DMD and retains AutoAdapter; aligning control before injection supports faster, more stable convergence of student and LoRA fake-score updates.The final objective combines baseline DMD with a sensitivity regularization term.Source: §4.4 · pp. 7–8
- LSP counters late-stage mode seeking that makes outputs nearly deterministic for a prompt-condition pair by preserving response to changes in Gaussian noise.It is a bounding regularization intended to preserve conditional faithfulness while preventing degenerate mapping collapse.Source: §4.4 · p. 8
5 Experiments
Across six control types, FlashControl performs competitively with multi-step controllable baselines in one step, while the experiments examine qualitative outputs, adapter and teacher-design ablations, and extended alignment metrics. Evaluation uses COCO2017-5k and reports FID, Precision, and Recall.
5.1 Experimental Setup
- Evaluation covers six control types on COCO2017-5k for multi-step and one-step models, using FID, Precision, and Recall.The conditions are Canny, Depth, Segmentation, Pose, Hough lines, and Inpainting; one-step models use 1 NFE.Source: §5.1 · p. 8
5.2 Quantitative comparison
- FID 13.64 for FlashControl (FLUX.2-Klein-4B, 1 NFE) on Canny approaches its 28-step teacher's 12.70 and beats OminiControl's 20.14.The paper reports that the pattern extends to other conditionings and that FlashControl also surpasses one-step baselines JDM and SDXS.Source: §5.2 · p. 9
5.3 Qualitative comparison
- With one step, FlashControl produces outputs described as on par with 28–50-step methods, with aligned images and high perceptual quality.The qualitative comparison notes consistent brushwork and fabric detail in portraits, coherent lighting and snow density outdoors, and color shifts around hole boundaries in several baselines.Source: §5.3 · p. 9, Figure 3 · p. 9
5.4 Ablation study
- Full AutoAdapter gives the lowest FID and fastest convergence across Canny, Depth, and Seg, while removing components degrades FID and removing the adapter yields the worst, less stable training.The ablation compares no adapter, no distribution matching, no token-wise gate, and full AutoAdapter; Table 2 reports FID at 30k iterations.Source: §5.4 · p. 9, Table 2 · p. 9
- Without LSP, FID worsens; single-branch DeltaControl reduces backward time by 50% and training memory by 40% versus dual-branch ControlNet, while FlashControl achieves 0.042s per 512 × 512 image on an NVIDIA A100 40GB.Source: §5.4 · p. 10
- Under the same one-step distillation recipe, branch-based teachers transfer poorly and DeltaControl yields the strongest one-step student among the compared single-stream teachers.Table 5 varies teacher architecture and reports FID on COCO-5k; the text says single-stream design alone is insufficient.Source: §5.4 · p. 10, Table 5 · p. 10
5.5 Extended Alignment Metrics for Quantitative Results
- Extended alignment evaluation uses RMSE for Canny and Depth structural consistency and mean Pixel Accuracy (mPA) for segmentation semantic alignment.Table 6 compares model performance on COCO-5K using these extended metrics; unavailable conditions are marked with a dash.Source: §5.5 · p. 10, Table 6 · p. 10
6 Conclusion
The paper concludes that one-step controllable generation is possible with a streamlined single-branch design, while its appendix analyzes adapter and LSP effects on distillation and feature alignment. It identifies simultaneous multi-condition composition as future work.
6 Conclusion
- FlashControl preserves much of the teacher's controllability and performs on par with traditional two-branch methods while simplifying training and inference.The conclusion identifies combining conditions such as depth and pose in one input as a future direction.Source: §6 · p. 10
Appendix Overview.
- The appendix covers extended alignment metrics, distillation behavior, user study, related work, dataset and training details, condition-count effects, and comparisons of distillation objectives.It also lists additional experiments, limitations, AutoAdapter details, and qualitative results.Source: Appendix Overview · p. 15
A AutoAdapter benefits distillation.
- AutoAdapter yields earlier recognizable structure, cleaner conditional outputs, and lower, smoother losses before late-stage instability, making Stage-2 updates easier and more stable to optimize.The comparison holds distillation conditions identical across the plain single-stream baseline, SS + AutoAdapter, and SS + AA + LSP.Source: A AutoAdapter benefits distillation · p. 15, Figure 6 · p. 16, Figure 7 · p. 16
B LSP benefits distillation.
- Without LSP, distillation shows a pronounced spike in student and LoRA losses and visual degradation; with LSP, losses remain stable and coherent structure is preserved.The analysis characterizes LSP as preserving sensitivity to input noise and preventing near-deterministic mapping when spatial control dominates.Source: B LSP benefits distillation · p. 15
C Effect of AutoAdapter on feature alignment.
- Full AutoAdapter reduces control-hidden statistic mismatch fastest and remains lowest, while removing distribution matching or the token-wise gate gives intermediate results.Removing AutoAdapter produces the highest mismatch; the experiments use SD2.1, DMD only, and the same hyperparameters.Source: C Effect of AutoAdapter on feature… · pp. 15–16
D User Study
In a study with 279 participants, FlashControl was preferred over three baselines on image quality, condition alignment, and textual alignment despite using one NFE rather than 50. The reported results show stronger perceptual quality at 50× faster generation.
D User Study
- FlashControl was preferred in 57.3% of image-quality cases, 43.0% of condition-alignment cases, and 47.0% of textual-alignment cases, using 1 NFE versus 50 for the baselines.The study included 279 participants, 30 prompts across Canny, Depth, and Segmentation, and comparisons against ControlNet, Ctrl-Adapter, and CtrLoRA.Source: D User Study · p. 16
E Dataset.
DeltaControl training uses a 2M-image subset of Re-LAION5B, with spatial conditions extracted from the same images using the ControlNet preprocessing pipeline. FlashControl distillation uses captions and conditions without requiring images.
E Dataset.
- DeltaControl is trained on a 2M subset of Re-LAION5B with conditions extracted from the same images using the ControlNet preprocessing pipeline.The conditions include Canny edges, depth maps, segmentation masks, M-LSD lines, Openpose, and inpainting; FlashControl distillation uses only captions and conditions.Source: E Dataset · p. 17
F Training Implementation Details
Table 7 covers training hyperparameters and compute settings for DeltaControl and FlashControl. The implementation applies LoRA across specified FFN and projection layers, while a separate training-difficulty result reports slightly worse FID and convergence with more control types under the same budget.
F Training Implementation Details
- For FLUX.2-klein-4B, LoRA covers all FFN layers in double-stream blocks and the first 14 fused projection layers in single-stream blocks, keeping trainable parameters comparable to ControlNet.Stable Diffusion 2.1 and PixArt-α apply LoRA to all FFN modules.Source: F Training Implementation Details · p. 17
- Table 7 summarizes training hyperparameters and compute settings for Stage 1 DeltaControl and Stage 2 FlashControl; conditioning projectors are randomly initialized and trained using LoRA.Source: Table 7 · p. 18
- With more control types and the same budget, FID and convergence rate are slightly worse.Source: Figure 10 · p. 18
G Number of Training Conditions
The experiments vary the number of training condition types and report worse image quality with more conditions at the same model size and budget. The supplied quantitative-results table describes FID comparisons across six tasks for different distillation methods.
G Number of Training Conditions
- With the same model size and training budget, training on more condition types hurts image quality.The experiment varies the number of condition types across 1, 3, and 6 and tracks FID during training.Source: G Number of Training Conditions · p. 18
- The authors suggest that broader condition coverage may force model parameters to explain a wider range of inputs.Source: G Number of Training Conditions · p. 18
- Table 8 reports FID results for different distillation methods across six tasks and identifies FlashControl as leveraging DMD.Source: Table 8 · p. 19
H Varying Distillation Methods
The comparison replaces DMD with DMD2 while retaining LSD loss in both configurations. With adversarial loss, DMD2 is reported to provide more robust performance than DMD.
H Varying Distillation Methods
- Both distillation configurations use LSD loss while the method choice changes from DMD to DMD2.Source: H Varying Distillation Methods · p. 19
- With adversarial loss, DMD2 provides more robust performance than DMD.Source: H Varying Distillation Methods · p. 19
I More Details About Our Experiments
The experiments compare parameter-matched DeltaControl variants with ControlNet on Stable Diffusion 2.1 and PixArt-α. All variants use a shared single-Canny conditioning pipeline.
I More Details About Our Experiments
- Figure 5 compares ControlNet with DeltaControl variants on a U-Net backbone, Stable Diffusion 2.1, and a DiT backbone, PixArt-α.Source: I More Details About Our Experiments · p. 19
- On PixArt-α, the variants tune attention layers, FFN layers, or a parameter subset drawn from both across all 28 transformer blocks.Source: I More Details About Our Experiments · p. 19
- On Stable Diffusion 2.1, the ResNet, FFN, and attention variants tune different combinations of convolutional, feedforward, and attention layers across the down, mid, and up blocks.Source: I More Details About Our Experiments · p. 19
- The comparison uses the same trainable-parameter percentage as ControlNet and a single Canny edge condition across all models.The condition is VAE-encoded, mapped by a small trainable embedding module, and added to the intermediate latent at the network head.Source: I More Details About Our Experiments · p. 19
J Limitations
The limitations section outlines FlashControl’s scope and unresolved issues, including costly two-stage training, 512 × 512 single-control generation, multi-condition composition, and misuse risks, alongside supplementary AutoAdapter, qualitative, and checklist details. It also notes gaps in user-study reporting and states that LLMs are not an important or non-standard component of the core method and experiments.
J Limitations
- FlashControl requires a strong pretrained backbone and costly two-stage training; experiments cover only 512 × 512 images and six spatial controls, not higher resolutions, video, or specialized domains.The authors also report that adding condition types under the same budget can slow convergence and slightly degrade quality, leaving robust multi-condition composition for future work.Source: J Limitations · p. 19
2. Limitations
- Fast one-step generation could lower the cost of creating misleading visual content, so responsible release should include usage restrictions, provenance tools, and safety filters.The checklist likewise identifies misleading-content risks and lists usage restrictions, provenance or watermarking tools when possible, and inherited content-safety filters as safeguards.Source: J Limitations · p. 19, §2 · p. 31
L AutoAdapter: Detailed Design
- AutoAdapter VAE-encodes and projects control inputs, aligns their token distribution to backbone states, then gates each aligned token before adding it to the corresponding hidden state.Its gate combines token similarity, a timestep factor favoring stronger control early in denoising, and predicted uncertainty; noisy or inconsistent tokens receive less influence.Source: L AutoAdapter: Detailed Design · p. 20
K More Qualitative Results
- Additional qualitative results show DeltaControl with 20 steps and FlashControl with one step for PixArt-α and FLUX.2-klein-4B.Figures 11–16 present these examples; the captions identify the models and step counts but do not state comparative outcomes.Source: Figure 11 · p. 21, Figure 12 · p. 22, Figure 13 · p. 23
NeurIPS Paper Checklist
- The checklist says the paper provides experimental details for reproduction but did not provide anonymized code or checkpoints at submission, with release planned subject to company policies.It also reports no error bars or multi-seed confidence intervals because training multiple large teachers and students is computationally expensive; comparisons instead use fixed evaluation sets and matched protocols.Source: §2 · pp. 28–29
2. Limitations
- The paper says it cites the datasets, pretrained models, and baselines used and follows their stated licenses and terms; it introduces no new dataset or benchmark.If code or checkpoints are released after publication, they will be documented with training details, intended use, limitations, and license information.Source: §2 · p. 32
- The paper reports a user study but omits participant instructions, screenshots, and compensation details, which it says should be added to supplemental material.The current version does not provide these materials or details.Source: §2 · p. 32
- The paper does not state whether the user study received IRB approval, an exemption, or equivalent review, and says this should be clarified while preserving anonymity.Source: §2 · p. 33
- The paper says its core method, experiments, and scientific claims do not use LLMs as an important, original, or non-standard component.Any writing or formatting assistance does not affect the methodology or results.Source: §2 · p. 33