UNITY: Attention Flow Networks for Adaptive Conditioning in Diffusion

2026-06-18

AI Summary

This work addresses the limitations of diffusion models in multimodal conditional generation—specifically, the need for separate adapter training per condition, high computational overhead, and poor generalization—by introducing the UNITY framework. UNITY employs a two-stage training strategy comprising universal semantic learning followed by modality-specific fine-tuning, coupled with a lightweight adapter architecture that enables efficient condition fusion without modifying the base model. Its key innovations include the Morphable Attention Flow network and the Morph Wrapper module, which leverage learnable flow fields and channel–spatial adaptive attention to maintain constant computational complexity under both single and composite conditions. Experiments demonstrate that UNITY achieves state-of-the-art image generation fidelity across multiple datasets while significantly reducing inference latency and memory consumption.

Abstract

We introduce UNITY, a Universal-to-Specialized adapter for efficient and scalable composite conditioning in diffusion based image generation. Unlike prior methods that train separate adapters for each conditioning modality, UNITY jointly learns shared semantics across multiple conditioning types and subsequently specializes without modifying the underlying architecture. The proposed two stage training paradigm consists of a Universal Stage that captures cross modal representations across all conditioning modalities using half of the total training steps, followed by a Specialization Stage that refines modality specific features using the remaining training budget. At the core of UNITY are the Morphable Attention Flow (MAF) Network and Morph Wrapper modules, which enable channel aware and spatially adaptive feature alignment through learnable flow fields and attention based fusion. This constant complexity formulation supports flexible operation under both single and composite conditioning settings while significantly reducing inference latency and memory consumption. Extensive experiments across multiple datasets demonstrate that UNITY achieves state of the art image fidelity while maintaining superior memory efficiency. Code: https://github.com/arya-domain/UNITY

Problem

• diffusion
• composite conditioning
• adapter
• image generation
• modality

Innovation

• UNITY
• Morphable Attention Flow
• composite conditioning
• diffusion models
• two-stage training