← 返回 2026-04-03

无需多模态注意力的门控条件注入：迈向可控线性注意力Transformer Gated Condition Injection without Multimodal Attention: Towards Controllable Linear-Attention Transformers

Yuhe Liu, Zhenxiong Tan, Yujia Hu, Songhua Liu, Xinchao Wang 📅 2026-03-29 👍 18 2026-07-13 08:36

可控生成扩散模型端侧部署线性注意力门控机制

提出GateControl框架，通过门控机制在高效线性注意力模型上实现多类型可控图像生成

前置知识

线性注意力

线性注意力通过核分解替代标准注意力的指数点积，将复杂度从O(n²)降低到O(n)。公式为Attention(Q,K,V)=φ(Q)(φ(K)^T V)，其中φ是特征映射函数（如ReLU）。这显著减少了长序列处理时的计算开销，使得处理更高分辨率的图像成为可能，特别适合端侧部署场景。

本文的核心贡献就是针对线性注意力模型设计可控生成框架，理解线性注意力的工作原理和优势是理解本文技术路线的基础。

流匹配

流匹配是扩散模型的一种替代框架，将去噪过程表述为概率密度流。目标是最小化噪声速度场vt(x)和网络预测速度场ut(x)之间的L2距离：L^FM=E_{t,p_t(x)}[||vt(x)-ut(x)||²]。SANA采用基于修正流的变体，通过学习从噪声分布到数据分布的向量场来实现高效生成。

本文基于SANA模型进行实验，SANA采用流匹配框架，理解其训练目标有助于理解实验设置和损失函数设计。

LoRA微调

低秩自适应通过分解权重更新矩阵为两个低秩矩阵A和B，将更新表示为ΔW=BA，其中A∈R^{d×r}, B∈R^{r×d}且r≪d。这大幅减少了可训练参数数量，同时保持了良好的微调效果。本文使用rank=16的LoRA，仅需18.9M参数（约1.18%）就能适应新的条件输入。

本文采用LoRA而非全参数微调，这是实现高效训练和部署的关键技术选择，理解LoRA有助于理解本文的参数效率优势。

研究动机

现有强大的扩散模型（如Stable Diffusion）由于巨大的计算和内存需求，通常部署在云端服务器上。这种依赖云端的范式迫使用户上传图像、草图或其他个人数据进行推理，不可避免地引发隐私泄露担忧，限制了敏感或离线场景中可控生成的应用。虽然线性注意力模型（如SANA）提供了显著降低的内存和计算开销，适合端侧执行，但现有的可控生成框架在这些架构上表现不佳。ControlNet的特征融合机制隐式假设输入条件与潜在表示之间存在空间对齐，这在主体驱动生成等空间不对齐场景中失效。OminiControl虽然通过多模态注意力提供了更大的灵活性，但在线性注意力下收敛缓慢，在空间对齐任务上需要50k步训练，而主体驱动任务仅需15k步。

本文的目标是本文的目标是设计专门针对线性注意力扩散骨架的可控生成框架，实现强大的控制能力和安全、私有的端侧部署。具体而言，需要一种方法既能支持多种异构条件类型（包括空间对齐的边缘图、深度图等，以及空间不对齐的主体驱动生成），又能在线性注意力架构上实现快速收敛和高保真度生成。

与已有工作不同的是，本文的独特切入点是从注意力下沉现象和门控机制获得启发，认识到线性注意力存在信息压缩问题，需要通过门控来缓解。与ControlNet依赖空间对齐的特征加法不同，本文采用共享模块策略，将噪声潜在和图像条件通过相同的VAE编码器转换到共享参数空间。与OminiControl依赖完整的多模态注意力不同，本文通过最小化的门控机制实现更高效的条件注入，仅需0.09M参数就能显著加速收敛（10倍以上）。

核心方法

GateControl的整体思路是在线性注意力模型上构建一个统一、最小且高效的可控生成框架。直觉上，线性注意力虽然提高了效率，但在处理条件信息时存在信息损失问题。为此，本文设计了一个门控条件模块，通过双路径管道有效整合多种类型的条件输入，包括空间对齐和非对齐的线索。技术路线包括三个核心组件：共享模块编码减少参数、内部交互实现灵活的条件融合、门控控制加速收敛并提升性能。整个框架仅需LoRA微调和0.09M门控参数，就能实现强大的可控性。

核心创新点是引入可学习的门控模块，选择性融合本会被线性注意力交互抑制的token信息。门控机制基于输入特征计算每个token的保留强度，通过sigmoid激活将线性映射压缩到[0,1]范围，形成软动态过滤器。这种设计通过补偿线性注意力引入的信息压缩，不仅加速训练，还增强了条件特征提取。与ControlNet的特征加法（hx←hx+hc）和OminiControl的完整注意力交互不同，门控机制以最小开销实现了更灵活和高效的条件注入，同时支持空间对齐和非对齐任务。

方法步骤详情

方法步骤的完整描述如下：首先，共享模块编码将图像条件CI和噪声潜在X通过相同的VAE编码器转换到共享参数空间，然后通过相同的模型结构处理。这避免了ControlNet使用独立编码器（如CLIP）需要的额外对齐，将额外参数从590M减少到18.9M（LoRA）。其次，内部交互将噪声潜在X、文本条件CT和图像条件CI的token作为一个整体[X; CT; CI]处理，使每个块能将它们作为集成输入一起处理。这允许模型在线性注意力模块中执行双向注意力交互，实现X和CI之间的灵活信息融合，同时与文本条件通过交叉注意力交互并在Mix-FFN中融合。最后，门控控制在线性注意力层后应用自适应门控，计算公式为hX'=g(hX,X,Wg1,σ)=σ(XWg1)⊙hX，hCI'=g(hCI,CI,Wg2,σ)=σ(CIWg2)⊙hCI，然后融合为hX←hX+hCI'。门控分数对每个token独立计算，使用自注意力前的特征计算门控分数效果最佳。

技术新颖性

技术新颖性体现在多个方面。首先，这是首个专门针对线性注意力扩散骨架设计的可控生成框架，填补了端侧可控生成的空白。其次，门控机制在线性注意力上下文中具有新颖性，通过最小化的token-wise门控（0.09M参数，仅0.006%）实现了显著加速（空间任务10倍以上）和性能提升。第三，共享模块策略与现有方法形成鲜明对比，ControlNet需要复制整个骨干网络（42%额外参数），而本文通过LoRA微调和门控机制实现高效适应。第四，统一框架同时支持空间对齐和非对齐任务，而现有方法通常只擅长其中一类。第五，消融研究系统探索了门控的四个关键方面：是否使用门控、插入位置（自注意力后/交叉注意力后/Mix-FFN后）、门控类型（token-wise/element-wise/直接加法）、输入特征（自注意力前/后），为未来研究提供了深入洞察。

**Figure 3**: Exploration of different methods and comparison of results. We adopt a shared module to handle the noisy latent and image condition, thus maximally retaining the original information flow of the model.

**Figure 4**: Details of the gated control method. Even a minimal token-wise gating mechanism could effectively filter tokens by strength, enabling adaptive token fusion with negligible cost.

实验结果

核心发现包括：在五个空间对齐任务（Canny、Depth、Deblurring、Colorization、HED）上，本文方法在SANA基础上相对于ControlNet和OminiControl实现了压倒性改进。例如，在Colorization任务中，FID从24.95降至10.28，改善超过50%；在HED任务中，MSE从2320降至1168，改善超过50%。收敛行为分析显示，门控调制导致训练损失显著陡峭下降，在Canny任务中仅需1k步就超过10k步的基线，加速超过10倍。CLIP-Image评分从最早训练阶段就保持领先。在主体驱动生成任务上，本文方法在DreamBooth数据集上实现了52.9%的Identity Preservation、63.5%的Material Quality和72.4%的Natural Appearance，显著超过IP-Adapter的24.8%、30.4%和56.1%。消融研究表明，使用自注意力前的特征进行token-wise门控效果最佳，移除门控机制导致FID、SSIM和CLIP分数的显著退化。在采样步数和引导尺度的鲁棒性测试中，本文模型在低步推理和变化引导尺度下产生更好、更稳定的输出，适合低延迟场景。

**Table 1**: Additional parameters introduced by different image conditioning methods. For our method, results are reported with LoRA rank set to 16.

**Table 2**: Quantitative comparison with baseline methods on five spatially aligned tasks. We evaluate methods based on Controllability (F1-Score for Canny, MSE for others), Image Quality (SSIM, FID and MUSIQ), and Alignment (CLIP Text and CLIP Image).

**Table 3**: Ablation studies on the mechanism of gated control for information integration. Best results are in bold.

**Figure 1**: Results of our proposed method on SANA-1.0-1.6B for both subject-driven generation and spatially aligned tasks.

**Figure 2**: Convergence behavior. The introduction of gated modulation results in a substantially steeper decline in training loss compared with interaction mechanisms that rely solely on attention.

**Figure 5**: Visualization results of our model on subject-driven tasks. The conditioning images are sourced from the DreamBooth dataset.

**Figure 2 (Supplementary)**: Robustness to sampling steps and guidance scale. Our model produces better and more stable outputs than OminiControl under both low-step inference and varying guidance scales.

查看结构化数据

任务	指标	本文	基线	提升
Colorization	FID	10.28	ControlNet: 24.95	58.8%提升
HED	MSE	1168	ControlNet: 2320	49.7%提升
Deblurring	MSE	14	OminiControl: 120	88.3%提升
主体驱动-Identity Preservation	GPT-4o评分(%)	52.9	IP-Adapter: 24.8	113.3%提升
收敛速度(Canny-to-image)	达到同等性能所需训练步数	1k步	OminiControl: 10k步	10倍加速
参数效率	额外参数占原模型比例	0.006%(0.09M)	ControlNet: 36.9%(590M)	参数减少98.4%

局限与改进

局限性分析包括：作者承认门控机制在交叉注意力层后插入导致训练损失高度不稳定，在Mix-FFN层后插入只导致轻微性能波动。直接加法特征导致训练损失不稳定，凸显了动态选择相关token信息的门控的关键作用。在多条件控制实验中，当结合多个条件时可能出现冲突，例如满足几何约束可能略微改变原始主体形状。本文主要在SANA模型上验证，对其他线性注意力架构的泛化性尚需更多实验。虽然门控机制在softmax注意力上也能加速收敛，但本文对其他类型注意力机制的探索有限。在极端低步推理（如极低采样步数）场景下的性能表现需要进一步研究。

独立分析的弱点

独立分析的弱点包括：首先，门控机制虽然高效，但可能在高维特征上损失细粒度信息，element-wise门控虽然性能可比但需要200M额外参数（vs 0.09M），存在效率-精度的权衡。其次，内部交互依赖双向线性注意力，对于需要强空间对齐的任务，可能仍不如显式空间注入机制精确。第三，本文主要在图像生成任务上验证，对视频生成、3D生成等扩展任务的适用性需要探索。第四，门控分数的计算基于输入特征的线性变换，可能难以建模复杂的非线性条件关系，可以考虑引入多层感知机或自注意力增强门控能力。第五，在多条件冲突场景下缺乏显式的冲突解决机制，可以通过引入条件权重或冲突检测模块改进。

未来方向

未来研究方向包括：作者提到个性化角色创建和视频生成领域有巨大潜力，其中在保持主体身份的同时将其自然嵌入不同环境至关重要。基于本文成果可延伸的方向包括：探索门控机制在视频生成中的应用，实现时序一致的主体控制；扩展到3D生成任务，结合深度、法线等多模态条件；研究更复杂的多条件融合策略，处理条件冲突和优先级；探索动态门控网络，根据任务类型自适应调整门控策略；将门控机制应用于其他线性注意力架构，验证泛化性；研究在极端资源约束下的端侧部署优化，包括量化和剪枝；探索结合检索增强生成，利用大规模条件数据库提升多样性。

复现评估

复现评估显示，论文详细描述了实验设置：使用4张NVIDIA H200 GPU，每GPU batch size为16；采用Prodigy优化器，初始学习率1.0，权重衰减0.01，启用safeguard warmup和bias correction。空间对齐任务在Text-to-Image-2M的10k图像上训练10k步，主体驱动任务在Subject200K的1024×1024子集上训练20k步。LoRA rank默认为16。评估使用COCO 2017验证集的5000张图像。论文提供了详细的消融研究和超参数分析。然而，论文未明确说明代码和模型权重是否开源，训练数据和具体实现细节可能需要进一步确认。总体而言，复现难度中等，主要挑战在于需要大量GPU资源和特定的数据集访问。

论文图表

Figure 4 (Supplementary): Intuitive analysis of the ablation on gated application through visualizations. Without gating, the model's ability to leverage conditional information is substantially impaired.

补充材料中的Figure 4通过可视化提供了门控应用的消融直观分析。展示了Canny任务中不同配置的结果：w/o gating、Diff. input features、本文方法。不带门控时，模型利用条件信息的能力严重受损，可能导致输出不符合条件或生成不合理的人为痕迹。使用不同输入特征生成门控分数会影响图像的全局方面，如风格和颜色。这补充了Table 3的量化结果，提供了直观的质量对比。

这张图对理解论文消融研究的实际影响至关重要，因为它将抽象的量化指标转化为直观的视觉对比。通过展示w/o gating导致的不合理输出，读者可以清楚地理解门控机制的实际价值：它不仅是提升指标，更是确保生成质量和条件符合性的关键。这补充了Table 3的定量分析，提供了更全面的消融证据。