An Efficient ControlNet for Diffusion Transformers with 85% Fewer Parameters

Existing controllableDiffusion Transformer methods, while making significant progress in advancing text-to-image and video generation, also introduce a large number of parameters and computational overhead.

At the same time, these methods fail to effectively consider the correlation of control information between different Transformer layers, resulting in inefficient resource allocation.

The research team of 360 Artificial Intelligence Research Institute proposed a new generation of efficient controllable generation framework guided by control correlation - RelaCtrl.

This framework achieves more efficient and resource-optimized control in Diffusion Transformer by optimizing the integration of control signals, thus effectively solving the above problems.

The team conducted experimental verification on the common T2I Diffusion Transformer and achieved remarkable results.

The first paper is Cao Ke, a second-year master's student from the University of Science and Technology of China. The co-author and project leader of the paper is Ma Ao, an AIGC researcher from 360 AI Research.

At present, relevant research papers, project homepages and code repositories have been officially opened and will be open source soon.

Research motivation

Currently, controllable generation methods based on Diffusion Transformer (DiT) (such as PixArt-δ, OminiControl, etc.) make full use of its powerful architecture and scalability, significantly improving the fidelity of the generated results and the consistency with the given text description.

However, the controlled generation method of DiT still faces two major problems:

First, these methods introduce a large number of additional parameters and computations, resulting in increased training and inference burdens.

Secondly, the correlation of control information between different network layers is often ignored, resulting in inefficient allocation of computing resources.

In related research in the field of model quantification and pruning, it has been found that the LLM model with a cascading Transformer structure has different influences on different layers, so detailed pruning strategies for different layers can be designed.

Inspired by this, the researchers measured the quality and control accuracy of the generated images by skipping different layers in the control network during the inference stage, thereby obtaining the "ControlNet correlation score", and used this to guide subsequent conditional injection of controllable generation.

As shown in the figure, there are significant differences in the correlation of different layers to control information in DiT. This correlation shows a trend of first increasing and then decreasing. The front center layer shows a higher correlation, while it is lower in deeper layers.

Existing methods fail to fully account for this variation and instead apply uniform settings to all layers, resulting in inefficient parameter and computational resource allocation, especially in redundant parameters and computations in less relevant layers.

Therefore, the research team proposed an efficient and controllable generation framework based on correlation guidance. Through precise parameter allocation and control strategies, combined with lightweight designed control modules, it can effectively reduce model parameters and reduce computational complexity.

Efficient and controllable generation framework RelaCtrl

The research team first identified and selected the 11 most critical control positions (sorted from high to low by relevance) for integrated control modules based on the derived ControlNet correlation score rankings.

With this approach, the team achieved comparable control performance to PixArt-δ using 13 replicated blocks, while reducing the number of parameters by approximately 15%.

Although this approach effectively reduces the model size and computational overhead, there is still a certain degree of redundancy in the internal design of the control module.

MetaFormer attributes the effectiveness of Transformers to two key components:

**The token mixer implemented through the self-attention mechanism and the channel mixer implemented through the feedforward network (FFN) layer. **

However, research shows that although FFN consumes a considerable portion of model parameters, it is often highly redundant. To this end, the research team introduced a lightweight module - Relevance Guided Lightweight Control Block (RGLC), which unifies token mixing and channel mixing into one operation. Specifically, the authors design a novel two-dimensional shuffle mixer (TDSM) that replaces the attention layer and FFN layer in the original PixArt Transformer block, thereby simplifying the architecture and improving efficiency.

This method promotes information interaction and modeling between token and channel dimensions, and greatly reduces the number of parameters and computational requirements for copy blocks.

The implementation of the two main parts will be described in detail below:

DiT-ControlNet Relevance Prior

To systematically evaluate the correlation of various layers in DiT-ControlNet with generation quality and control accuracy, the research team trained a fully controlled PixArt-α network containing 27 replicated modules.

During inference, researchers systematically skip each control block layer and evaluate its impact on the final generated result.

For quantitative evaluation, Fréchet Inception Distance (FID) is introduced to measure the quality of image generation, and Hausdorff Distance (HDD) is used to evaluate control accuracy.

These metrics help analyze the impact of skipping individual layers in the control branch on overall performance and provide a relevance score for each control block. Finally, based on the combination of these two indicators, the research team calculated the correlation score of ControlNet, and selected the 11 most critical control positions (sorted from high to low by relevance) for the integrated control module, as shown in the white serial number in the figure.

Two Dimensions Shuffle Mixer

In order to reduce the computational burden of control branches, the research team proposed a method of grouping calculations on tokens, and adopted specific strategies to enhance the interaction and modeling capabilities between token groups.

Specifically, this operation first performs random channel selection, then randomly scrambles the input sequence in the token dimension space, and then performs local self-attention calculation.

Although the subsequent attention mechanism is limited to a fixed group, the tokens involved may come from outside this group.

This operation effectively disrupts the inherent relationship between tokens and introduces information between other local windows to a certain extent, thereby breaking the interaction constraints usually imposed by local attention.

In addition, in order to ensure that the partition operation will not affect the semantic information embedded in the latent code during the recovery process, the research team designed to apply the inverse recovery operation on the token and channel dimensions after the calculation of self-attention.

This holistic approach that combines shuffle and recovery, called a two-dimensional shuffle mixer (TDSM), ensures that information is retained during self-attention calculations by leveraging the ability of reversible transformation pairs to achieve efficient non-local information interaction across channel and token dimensions.

Experimental results

The authors conducted qualitative and quantitative evaluations on controllable generation tasks under various conditions to verify the effect of RelaCtrl.

Qualitative comparisonAs shown in the figure, RelaCtrl shows very competitive performance in the visual effect comparison and is able to generate photo-realistic images that are highly consistent with real images under the constraints of given conditions. Even under different conditions, it also demonstrates the powerful ability to generate images.

Quantitative Analysis

The research team conducted a quantitative evaluation of RelaCtrl and comparison methods on the COCO validation set. The results show that RelaCtrl achieves better performance in control indicators, visual quality and text similarity, verifying its effectiveness in all aspects.

Computational complexity analysis

As shown in the table below, RelaCtrl based on PixArt-α only brings about a 7.38% increase in parameters and an 8.61% increase in computational complexity, which is significantly lower than the increment of the ControlNet method, which increases both parameters and complexity by nearly 50%. Experiments on Flux.1-dev further verified this. RelaCtrl achieved a balance between the increase in parameter volume and computational complexity, and at the same time significantly outperformed the ControlNet method on multiple indicators. This advantage is also reflected in the actual inference time, with RelaCtrl showing higher efficiency.

ablation study.

Compatibility with community models:

The research team used PixArt weights fine-tuned with Lora for inference. RelaCtrl can effectively utilize these community weights. The image below shows the paint, oil painting, antique and pixel style images produced by the model under specified conditions.

Paper address: http://arxiv.org/abs/2502.14377 Project homepage: https://360cvgroup.github.io/RelaCtrl/ Code repository: https://github.com/360CVGroup/RelaCtrl

— End —