Source-linked AI summary
Gated Condition Injection without Multimodal Attention: Towards Controllable Linear-Attention Transformers
Yuhe Liu, Zhenxiong Tan, Yujia Hu, Songhua Liu, Xinchao Wang
TL;DR
Controllable diffusion needs privacy-preserving, efficient deployment, but existing frameworks have limited heterogeneous-condition flexibility or converge slowly on linear-attention backbones. The paper proposes unified gated conditioning for models such as SANA, and reports robust control, faster convergence, and high-quality generation across aligned and subject-driven tasks.
Problem
Existing controllable generation frameworks lack flexibility for heterogeneous conditions or converge slowly on linear-attention diffusion models, limiting controllable on-device generation.
Method
The paper introduces a unified learnable gated conditioning framework that selectively preserves and fuses conditional information in linear-attention backbones.
Results
Experiments demonstrate strong controllability, faster convergence, and high-quality generation across spatially aligned and subject-driven tasks with 0.09M additional parameters.
Takeaways & Limitations
GateControl provides a flexible, efficient control framework for linear-attention diffusion models, supporting diverse conditions and on-device deployment.
Takeaways & Limitations
Spatially aligned tasks require significantly more training iterations than subject-driven tasks, reported as 50k versus 15k.
Abstract
from arXiv · showhide
Recent advances in diffusion-based controllable visual generation have led to remarkable improvements in image quality. However, these powerful models are typically deployed on cloud servers due to their large computational demands, raising serious concerns about user data privacy. To enable secure and efficient on-device generation, we explore in this paper controllable diffusion models built upon linear attention architectures, which offer superior scalability and efficiency, even on edge devices. Yet, our experiments reveal that existing controllable generation frameworks, such as ControlNet and OminiControl, either lack the flexibility to support multiple heterogeneous condition types or suffer from slow convergence on such linear-attention models. To address these limitations, we propose a novel controllable diffusion framework tailored for linear attention backbones like SANA. The core of our method lies in a unified gated conditioning module working in a dual-path pipeline, which effectively integrates multi-type conditional inputs, such as spatially aligned and non-aligned cues. Extensive experiments on multiple tasks and benchmarks demonstrate that our approach achieves state-of-the-art controllable generation performance based on linear-attention models, surpassing existing methods in terms of fidelity and controllability.
1. Introduction
The paper targets privacy-preserving, on-device controllable diffusion with linear-attention backbones, addressing limitations in flexibility and convergence. It introduces gated conditioning to support heterogeneous controls while improving efficiency and controllability.
- Cloud deployment of controllable diffusion raises privacy concerns, motivating lighter linear-attention models such as SANA for on-device generation.
- ControlNet assumes spatial alignment, limiting subject-driven control, while existing alternatives remain insufficiently flexible or efficient for linear-attention models.
- The proposed learnable gate selectively preserves conditional information compressed by linear attention and controls its influence across layers.
- Experiments report strong controllability and generalization across diverse tasks, with advantages over ControlNet and OminiControl.
- The framework supports both spatially aligned and non-aligned conditions, including edge, depth, and subject-driven generation.
- The method achieves faster convergence and comparable or better fidelity with only 0.09M additional parameters.
2. Related work
Related work has progressed from concatenated conditioning toward attention, adapters, dynamic fusion, and gated architectures. These approaches seek flexible control and efficient information flow, but controllable generation still faces computational and flexibility constraints.
- Diffusion and transformer models improve image synthesis but remain constrained by heavy computation and inflexible controllability.
- Conditional generation methods inject text, masks, poses, or layouts through concatenation, cross-attention, adapters, and dynamic fusion pathways.
- Modern conditioning techniques improve alignment and fine-grained consistency but often incur high training and inference costs.
- Gating mechanisms regulate information flow in recurrent, feed-forward, attention, and mixture-of-experts architectures.
3. Methodology
GateControl adapts controllable diffusion to linear-attention backbones through shared-module encoding, internal token interaction, and token-wise gated fusion. The design supports spatially aligned and non-spatial conditions while targeting faster convergence and low parameter overhead.
- 3.1. Preliminary: Linear attention reduces attention’s quadratic sequence-length cost, and SANA adopts ReLU-based linear attention for efficient larger-image generation.
- 3.2. Overview: GateControl targets flexible, minimal, and efficient controllable generation for spatially aligned and non-spatially aligned conditions.
- 3.2.1. Shared-module encoding: ControlNet introduces excessive architectural components and parameters while adapting poorly to non-spatially aligned inputs, especially subject-driven cases.
- 3.2.1. Shared-module encoding: Shared-module encoding processes image conditions and noisy latents through the same VAE encoder and model structure, reducing the need for separate alignment components.
- 3.2.1. Shared-module encoding: 18.9M trainable parameters, less than 1.18%, replace the 590M additional parameters required by ControlNet under the reported comparison.
- 3.2.2. Internal interaction: Internal bidirectional linear attention jointly processes noisy latents, text conditions, and image conditions, supporting depth, Canny, coloring, deblurring, and subject-driven generation.
- 3.2.2. Internal interaction: Without explicit spatial information injection or alignment, internal interaction converges much more slowly on spatial tasks, requiring 50k iterations versus 15k for subject-driven tasks.
- 3.2.3. Efficient gated control: Token-wise gating computes data-dependent sigmoid scores to softly preserve or suppress features before adaptively fusing latent and image-condition hidden states.
4. Experiments
Experiments evaluate the gated framework across spatially aligned and subject-driven tasks, showing stronger controllability, quality, convergence, and parameter efficiency. Ablations identify tokenwise gating and its placement before self-attention as effective design choices.
- Experimental setup: Experiments use SANA-1.0 with LoRA finetuning across five spatially aligned tasks and subject-driven generation.Spatially aligned tasks include Canny-to-image, depth-to-image, colorization, deblurring, and HED-to-image.
- Quantitative comparison: FID drops from 24.95 to 10.28 for colorization, while HED MSE falls from 2320 to 1168 against SANA-based baselines.The comparison reports gains across controllability, image quality, and alignment metrics.
- Qualitative comparison: The method preserves subject identity while supporting localized edits, including accessories, clothing changes, and shoe-color alterations.Additional visualizations use conditioning images sourced from the DreamBooth dataset.
- Convergence behavior: The method captures explicit positional relationships within the first 1k steps, whereas the baseline requires more than 10k steps to acquire spatial correspondence.Training loss declines more steeply, and CLIP-Image scores remain ahead from the earliest training stage.
- Parameter efficiency: The gated control adds 0.09M parameters, equal to 0.006% of the original SANA model size.This small overhead accompanies reported improvements in convergence speed and controllability.
- Ablation studies: Tokenwise gating enables dynamic token selection with minimal overhead, while element-wise gating uses 200M parameters and direct summation destabilizes training.Among insertion points, placement after cross-attention produces unstable loss, whereas placement after Mix-FFN causes only minor fluctuations; features before self-attention perform better.
5. Conclusion
The paper presents a unified gated framework for controllable diffusion on linear-attention backbones. Experiments across aligned and subject-driven tasks support stable, expressive control, faster convergence, and minimal parameter overhead for efficient on-device generation.
- Conclusion: The unified gated framework targets controllable diffusion on linear-attention backbones and addresses limitations of ControlNet and OminiControl.The conclusion frames the method as supporting diverse conditioning while retaining efficient deployment potential.
- Conclusion: Experiments demonstrate high-quality generation and robust control across spatially aligned and subject-driven conditions.The reported conclusion covers both controllability and convergence across diverse conditions.
Supplementary Material
Supplementary analyses examine robustness, convergence, and generality beyond the main linear-attention setting. The gated approach is reported to remain effective under varied sampling conditions and when added to OminiControl with softmax attention.
- Softmax-attention generality: The approach captures control signals earlier than original OminiControl, including when using softmax attention.This comparison is presented in Figure 1 as gated control on the original OminiControl.
- Sampling robustness: The model produces more stable outputs than OminiControl under low-step inference and varying guidance scales.Figure 2 evaluates robustness to both sampling steps and guidance scale.
- Softmax-attention generality: A single-block gated interaction on FLUX.1dev adds only 0.2M parameters while enabling earlier acquisition of control information.The variant uses image-condition tokens and latent tokens to test generality on softmax attention.
B. Visual analysis of ablations
Visual ablations show that gating is important for condition compliance and image quality, while gate-score inputs can influence global style and color. Comparisons also highlight more faithful condition following and natural colorization.
- Gated application: Removing gating substantially impairs conditional information use, causing condition violations and artifacts such as unnatural protrusions.The ablation also reports decreased overall image quality without gating.
- Subject-driven comparison: The proposed approach preserves object-specific features while adapting the environment naturally to editing prompts in subject-driven tasks.
- Gate-score inputs: Gate scores derived from different input sources can alter global image properties, including style and color.
- Spatially aligned comparison: Compared with SANA-based ControlNet, the method follows conditions more accurately and produces more natural, realistic colorization.
C. Sampling steps and guidance scale
The model remains robust under low-step inference and varying guidance scales, supporting low-latency use. It also adapts quickly to editing and can combine multiple conditions, although conflicts may slightly change subject shape.
- Sampling robustness: The model produces better and more stable outputs than OminiControl under low-step inference and varying guidance scales.These results support suitability for low-latency scenarios.
- Fast adaptation and multi-condition control: GateControl acquires basic editing capability within limited training steps and simultaneously incorporates subject and depth conditions.When conditions conflict, satisfying geometric constraints may slightly alter the original subject shape.
E.1. Evaluation for subject-driven generation
Subject-driven evaluation measures both identity fidelity and compliance with requested modifications across five criteria. Qualitative results show stronger object-feature preservation and natural environmental adaptation, alongside multi-condition capability.
- Evaluation protocol: Subject-driven evaluation uses Identity Preservation, Material Quality, Color Fidelity, Natural Appearance, and Modification Accuracy.Evaluations are conducted using the GPT-4o multimodal model.
- Editing and multi-condition control: The model demonstrates image editing and simultaneous incorporation of multiple conditions.
- Qualitative results: The method preserves object-specific features more faithfully while adapting the environment naturally and flexibly to editing prompts.
F. Comparisons on spatially aligned tasks
On spatially aligned tasks, the method more accurately follows conditions than SANA-based ControlNet and produces higher-quality, more natural colorization.
- Comparison with SANA-based ControlNet: Compared with SANA-based ControlNet, the method follows conditions more accurately and generates surrounding objects more faithfully.For coloring, it produces more natural and realistic colorization and higher-quality bus outputs.