LinFusion: 1 GPU, 1 Minute, 16K Image
Songhua Liu Weihao Yu Zhenxiong Tan Xinchao Wang



Estimating the coherence of noise in mid-scale quantum systems
    Inés de Vega
    Received ; accepted 
==============================================================







§ ABSTRACT 
Modern diffusion models, particularly those utilizing a Transformer-based UNet for denoising, rely heavily on self-attention operations to manage complex spatial relationships, thus achieving impressive generation performance. 
However, this existing paradigm faces significant challenges in generating high-resolution visual content due to its quadratic time and memory complexity with respect to the number of spatial tokens. 
To address this limitation, we aim at a novel linear attention mechanism as an alternative in this paper. 
Specifically, we begin our exploration from recently introduced models with linear complexity, e.g., Mamba2, RWKV6, Gated Linear Attention, etc, and identify two key features—attention normalization and non-causal inference—that enhance high-resolution visual generation performance. 
Building on these insights, we introduce a generalized linear attention paradigm, which serves as a low-rank approximation of a wide spectrum of popular linear token mixers. 
To save the training cost and better leverage pre-trained models, we initialize our models and distill the knowledge from pre-trained StableDiffusion (SD). 
We find that the distilled model, termed LinFusion, achieves performance on par with or superior to the original SD after only modest training, while significantly reducing time and memory complexity. 
Extensive experiments on SD-v1.5, SD-v2.1, and SD-XL demonstrate that LinFusion delivers satisfactory zero-shot cross-resolution generation performance, generating high-resolution images like 16K resolution. 
Moreover, it is highly compatible with pre-trained SD components, such as ControlNet and IP-Adapter, requiring no adaptation efforts. 
Codes are available https://github.com/Huage001/LinFusionhere. 




§ INTRODUCTION


Recent years have witnessed significant advancements in AI-generated content (AIGC) with diffusion models <cit.>. 
On the one hand, unlike classic models like GAN <cit.>, diffusion models refine noise vectors iteratively to produce high-quality results with fine details <cit.>. 
On the other hand, having trained on large-scale data pairs, these models exhibit satisfactory alignment between input conditions and output results. 
These capabilities have spurred recent advancements in text-to-image generation <cit.>. 
Benefiting from the impressive performance and the open-source community, Stable Diffusion (SD) <cit.> stands out as one of the most popular models. 


The success of models like SD can be largely attributed to their robust backbone structures for denoising. 
From UNet architectures with attention layers <cit.> to Vision Transformers <cit.>, existing designs rely heavily on self-attention mechanisms to manage complex relationships between spatial tokens. 
Despite their impressive performance, the quadratic time and memory complexity inherent in self-attention operations poses significant challenges for high-resolution visual generation. 
For instance, as illustrated in Fig. <ref>(a), using FP16 precision, SD-v1.5 fails to generate 2048-resolution images on A100, a GPU with 80GB of memory, due to out-of-memory errors, making higher resolutions or larger models even more problematic. 

To address these issues, in this paper, we aim at a novel token-mixing mechanism with linear complexity to the number of spatial tokens, offering an alternative to the classic self-attention approach. 
Inspired by recently introduced models with linear complexity, such as Mamba <cit.> and Mamba2 <cit.>, which have demonstrated significant potential in sequential generation tasks, we first investigate their applicability as token mixers in diffusion models. 


However, there are two drawbacks of Mamba diffusion models. 
On the one hand, when a diffusion model operates at a resolution different from its training scale, our theoretical analysis reveals that the feature distribution tends to shift, leading to difficulties in cross-resolution inference. 



On the other hand, diffusion models perform a denoising task rather than an auto-regressive task, allowing the model to simultaneously access all noisy spatial tokens and generate denoised tokens based on the entire input. In contrast, Mamba is fundamentally an RNN that processes tokens sequentially, meaning that the generated tokens are conditioned only on preceding tokens, a constraint termed causal restriction. Applying Mamba directly to diffusion models would impose this unnecessary causal restriction on the denoising process, which is both unwarranted and counterproductive. Although bi-directional scanning branches can somewhat alleviate this issue, the problem inevitably persists within each branch. 





Focusing on the above drawbacks of Mamba for diffusion models, we propose a generalized linear attention paradigm. 
Firstly, to tackle the distribution shift between training resolution and larger inference resolution, a normalizer for Mamba, defined by the cumulative impact of all tokens on the current token, is devised to the aggregated features, ensuring that the total impact remains consistent regardless of the input scale. 


Secondly, we aim at a non-causal version of Mamba. 
We start our exploration by simply removing the lower triangular causal mask applied on the forget gate and find that all tokens would end up with identical hidden states, which undermines the model's capacity. 
To address this issue, we introduce distinct groups of forget gates for different tokens and propose an efficient low-rank approximation, enabling the model to be elegantly implemented in a linear-attention form. 

We analyze the proposed approach technically alongside recently introduced linear-complexity token mixers such as Mamba2 <cit.>, RWKV6 <cit.>, and Gated Linear Attention <cit.> and reveal that our model can be regarded as a generalized non-causal version of these popular models. 



The proposed generalized linear attention module is integrated into the architectures of SD, replacing the original self-attention layers, and the resultant model is termed as Linear-Complexity Diffusion Model, or LinFusion in short. 
By only training the linear attention modules for 50k iterations in a knowledge distillation framework, LinFusion achieves performance on par with or even superior to the original SD, while significantly reducing time and memory complexity, as shown in Fig. <ref>. 
Meanwhile, it delivers satisfactory zero-shot cross-resolution generation performance and can generate images at 16K resolution on a single GPU. 
It is also compatible with existing components for SD, such as ControlNet <cit.> and IP-Adapter <cit.>, allowing users to inject additional controls to the proposed LinFusion flexibly without any additional training cost. 
As shown in Fig. <ref>, extensive experiments on SD-v1.5, SD-v2.1, and SD-XL validate the effectiveness of the proposed model and method. 
Our contributions can be summarized as follows: 

    
  * We investigate the non-causal and normalization-aware version of Mamba and propose a novel linear attention mechanism that addresses the challenges of high-resolution visual generation with diffusion models.
    
  * Our theoretical analysis indicates that the proposed model is technically a generalized and efficient low-rank approximation of existing popular linear-complexity token mixers.
    
  * Extensive experiments on SD demonstrate that the proposed LinFusion can achieve even better results than the original SD and exerts satisfactory zero-shot cross-resolution generation performance and compatibility with existing components for SD. To the best of our knowledge, this is the first exploration of linear-complexity token mixers on the SD series model for text-to-image generation. 




§ RELATED WORKS


In this section, we review related works from two perspectives, namely efficient diffusion architectures and linear-complexity token mixers. 



 §.§ Efficient Diffusion Architectures


There are mainly two mainstreams of works aiming at more efficient diffusion models, including efficient sampling for a reduced number of sampling time-steps <cit.> and efficient architectures for faster network inference. 
This paper focuses on the latter, which is a bottleneck for generating high-resolution visual results, particularly due to the self-attention token mixers in existing diffusion backbones. 

To mitigate the efficiency issue triggered by the quadratic time and memory complexity, a series of works, including DiS <cit.>, DiM <cit.>, DiG <cit.>, Diffusion-RWKV <cit.>, DiffuSSM <cit.>, and Zigma <cit.>. 
These works have successfully adapted recent state space models like Mamba <cit.>, RWKV <cit.>, or Linear Attention <cit.> into diffusion architectures. 



However, these architectures maintain a causal restriction for diffusion task, processing input spatial tokens one by one, with generated tokens conditioned only on preceding tokens. In contrast, the diffusion task allows models to access all noisy tokens simultaneously, making the causal restriction unnecessary. To address this, we eliminate the causal restriction and introduce a non-causal token mixer specifically designed for the diffusion model. 

Additionally, previous works have primarily focused on class-conditioned image generation. 
For text-to-image generation, <cit.> propose architectural pruning for Stable Diffusion (SD) by reducing the number of UNet stages and blocks, which is orthogonal to our focus on optimizing self-attention layers. 



 §.§ Linear-Complexity Token Mixers


Despite the widespread adoption of Transformer <cit.> across various fields due to its superior modeling capacity, the quadratic time and memory complexity of the self-attention mechanism often leads to efficiency issues in practice. 
A series of linear-complexity token mixers are thus introduced as alternatives, such as Linear Attention <cit.>, State Space Model <cit.>, and their variants including Mamba <cit.>, Mamba2 <cit.>, mLSTM <cit.>, Gated Retention <cit.>, DFW <cit.>, GateLoop <cit.>, HGRN2 <cit.>, RWKV6 <cit.>, GLA <cit.>, etc. 
These models are designed for tasks requiring sequential modeling, making it non-trivial to apply them to non-causal vision problems. 
Addressing this challenge is the main focus of our paper. 

For visual processing tasks, beyond the direct treatment of inputs as sequences, there are concurrent works focused on non-causal token mixers with linear complexity. 
MLLA <cit.> employs Linear Attention <cit.> as token mixers in vision backbones without a gating mechanism for hidden states. 
In VSSD <cit.>, various input tokens share the same group of gating values. 
In contrast, the model proposed in this paper relaxes these gating assumptions, offering a generalized non-causal version of various modern state-space models. 



§ METHODOLOGY




 §.§ Preliminary


Diffusion Models. 
As a popular model for text-to-image generation, Stable Diffusion <cit.> (SD) first learns an auto-encoder (ℰ,𝒟), where the encoder ℰ maps an image x to a lower dimensional latent space: z←ℰ(x), and the decoder 𝒟 learns to decode z back to the image space x̂←𝒟(z) such that x̂ is close to the original x. 
In the inference time, a Gaussian noise in the latent space z_T is sampled randomly and denoised by a UNet ϵ_θ for T steps. 
The denoised latent code after the final step z_0 is decoded by 𝒟 to derive a generated image. 
In training, given an image x and its corresponding text description y, ℰ is utilized to obtain its corresponding latent code, and we add a random Gaussian noise ϵ for its noisy version z_t with respect to the t-th step. 
The UNet is trained via the noise prediction loss ℒ_simple <cit.>:

    θ=_θ𝔼_z∼ℰ(x),y,ϵ∼𝒩(0,1),t[ℒ_simple]  ℒ_simple=‖ϵ-ϵ_θ(z_t,t,y)‖_2^2.


The UNet in SD contains multiple self-attention layers as token mixers to handle spatial-wise relationships and multiple cross-attention layers to handle text-image relationships. 
Given an input feature map in the UNet backbone X∈ℝ^n× d and weight parameters W_Q,W_K∈ℝ^d× d' and W_V∈ℝ^d× d, where n is the number of spatial tokens, d is the feature dimension, and d' is the attention dimension, self-attention can be formalized as:

    Y=MV,   M=softmax(QK^⊤/√(d')),    Q=XW_Q,   K=XW_K,   V=XW_V.

We can observe from Eq. <ref> that the complexity of self-attention is quadratic with respect to n since the attention matrix M∈ℝ^n× n, we mainly focus on its alternatives in this paper and are dedicated on a novel module for token mixing with linear complexity. 

Mamba. 
As an alternative to Transformer <cit.>, Mamba <cit.> is proposed to handle sequential tasks with linear complexity with respect to the sequence length. 
At the heart of Mamba lies the State Space Model (SSM), which can be written as:

    H_i=A_i⊙ H_i-1+B_i^⊤ X_i=∑_j=1^i{(∏_k=j+1^iA_k)⊙ (B_j^⊤ X_j)},   Y_i=C_iH_i,

where i is the index of the current token in a sequence, H_i denotes the hidden state, X_i and Y_i are row vectors denoting the i-th rows of the input and output matrices respectively, A_i, B_i, and C_i are input-dependent variables, and ⊙ indicates element-wise multiplication. 




 §.§ Overview


r4cm
[1]



  
    < g r a p h i c s >

  

  Overview of LinFusion. We replace self-attention layers in the original SD with our LinFusion modules and adopt knowledge distillation to optimize the parameters.
  
  






In the latest version, i.e., Mamba2 <cit.>, A_i is a scalar, B_i,C_i∈ℝ^1× d', X_i,Y_i∈ℝ^1× d, and H_i∈ℝ^d'× d. 
According to State-Space Duality (SSD), the computation in Eq. <ref> can be reformulated as the following expression, referred to as 1-Semiseparable Structured Masked Attention:

    Y=((CB^⊤)⊙Ã)X,

where Ã is a n× n lower triangular matrix and Ã_ij=∏_k=j+1^iA_k for j≤ i. 
Such a matrix Ã is known as 1-semiseparable, ensuring that Mamba2 can be implemented with linear complexity in n. 

  

In this paper, we aim at a diffusion backbone for the general text-to-image problems with linear complexity with respect to the number of image pixels. 
To this end, instead of training a novel model from scratch, we initialize and distill the model from pre-trained SD. 
Specifically, we utilize the SD-v1.5 model by default and substitute its self-attention—the primary source of quadratic complexity—with our proposed LinFusion modules. 
Only the parameters in these modules are trainable, while the rest of the model remains frozen. 
We distill knowledge from the original SD model into LinFusion such that given the same inputs, their outputs are as close as possible. 
Fig. <ref> provides an overview of this streamline. 

This approach offers two key benefits: (1) Training difficulty and computational overhead are significantly reduced, as the student model only needs to learn spatial relationships, without the added complexity of handling other aspects like text-image alignment; (2) The resulting model is highly compatible with existing components trained on the original SD models and their fine-tuned variations, since we only replace the self-attention layers with LinFusion modules, which are trained to be functionally similar to the original ones while maintaining the overall architecture. 

Technically, to derive a linear-complexity diffusion backbone, one simple solution is to replace all the self-attention blocks with Mamba2, as shown in Fig. <ref>(a). 
We apply bi-directional SSM to ensure that the current position can access information from subsequent positions. 
Moreover, the self-attention modules in Stable Diffusion do not incorporate gated operations <cit.> or RMS-Norm <cit.> as used in Mamba2. 
As shown in Fig. <ref>(b), we remove these structures to maintain the consistency and result in a slight improvement in performance. 
In the following parts of this section, we delve into the issues of applying SSM, the core module in Mamba2, to diffusion models and accordingly introduce the key features in LinFusion: normalization and non-causality in Secs. <ref> and <ref> respectively. 
Finally, in Sec. <ref>, we provide the training objectives to optimize parameters in LinFusion modules. 





 §.§ Normalization-Aware Mamba


In practice, we find that SSM-based structure shown in Fig. <ref>(b) can achieve satisfactory performance if the training and inference have consistent image resolutions. 
However, it fails when their image scales are different. 
We refer readers to Sec. <ref> for the experimental results. 
To identify the cause of this failure, we examine the channel-wise means of the input and output feature maps, which exhibit the following proposition: 

Assuming that the mean of the j-th channel in the input feature map X is μ_j, and denoting (CB^⊤)⊙Ã as M, the mean of this channel in the output feature map Y is μ_j∑_k=1^nM_ik. 

The proof is straightforward. 
We observe through Fig. <ref>(b) that there is non-negative activation applied on X, B, and C. 
Given that A is also non-negative in Mamba2, according to Prop. <ref>, the channel-wise distributions would shift if n is inconsistent in training and inference, which further leads to distorted results. 

Solving this problem requires unifying the impact of all tokens on each one to the same scale, a property inherently provided by the  function. 
In light of this, we propose normalization-aware Mamba in this paper, enforcing that the sum of attention weights from each token equals 1, i.e., ∑_k=1^nM_ik=1, which is equivalent to applying the SSM module one more time to obtain the normalization factor Z:

    Z_i=A_i⊙ Z_i-1+B_i,   C'_ij=C_ij/∑_k=1^d'{C_ik⊙ Z_ik}.

The operations are illustrated in Fig. <ref>(c). 
Experiments indicate that such normalization substantially improve the performance of zero-shot cross-resolution generalization. 





 §.§ Non-Causal Mamba




While bi-directional scanning enables a token to receive information from subsequent tokens—a crucial feature for diffusion backbones—treating feature maps as 1D sequences compromises the intrinsic spatial structures in 2D images and higher-dimensional visual content. 
To address this dilemma more effectively, we focus on developing a non-causal version of Mamba in this paper. 


Non-causality indicates that one token can access to all tokens for information mixing, which can be achieved by simply removing the lower triangular causal mask applied on Ã. 
Thus, the recursive formula in Eq. <ref> would become H_i=∑_j=1^n{(∏_k=j+1^nA_k)⊙ (B_j^⊤ X_j)}. 
We observe that H_i remains invariant with respect to i in this formula. 
This implies that the hidden states of all tokens are uniform, which fundamentally undermines the intended purpose of the forget gate A. 
To address this issue, we associate different groups of A to various input tokens. 
In this case, A is a n× n matrix and H_i=∑_j=1^n{(∏_k=j+1^nA_ik)⊙ (B_j^⊤ X_j)}. 

The Ã_ij in Eq. <ref> becomes ∏_k=j+1^nA_ik. 



Compared with that in Eq. <ref>, Ã here is not necessarily 1-semiseparable. 
To maintain linear complexity, we impose the assumption that Ã is low-rank separable, i.e., there exist input-dependent matrices F and G such that Ã=FG^⊤. 
In this way, the following proposition ensures that Eq. <ref> under this circumstance can be implemented via linear attention:

    Given that Ã=FG^⊤, F,G∈ℝ^n× r, and B,C∈ℝ^n× d', denoting C_i=c(X_i), B_i=b(X_i), F_i=f(X_i), and G_i=g(X_i), there exist corresponding functions f' and g' such that Eq. <ref> can be equivalently implemented as linear attention, expressed as Y=f'(X)g'(X)^⊤ X. 


Given existing conditions, we have: 

    (CB^⊤)⊙Ã   =[(c(X_i)b^⊤ (X_j))⊙ (f(X_i)g^⊤ (X_j))]_i,j
       =[(∑_u=1^d'{c(X_i)_ub(X_j)_u})(∑_v=1^r{f(X_i)_vg(X_j)_v})]_i,j
       =[∑_u=1^d'∑_v=1^r{(c(X_i)_uf(X_i)_v)(b(X_j)_ug(X_j)_v)}]_i,j
       =[(c(X_i)⊗ f(X_i))(b(X_j)⊗ g(X_j))^⊤]_i,j,

where ⊗ denotes Kronecker product. 
Defining f'(X_i)=c(X_i)⊗ f(X_i) and g'(X_i)=b(X_i)⊗ g(X_i), we derive Y=f'(X)g'(X)^⊤ X. 


In practice, we adopt two MLPs to mimic the functionalities of f' and g'. 
Combining with the normalization operations mentioned in Sec. <ref>, we derive an elegant structure shown in Fig. <ref>(d). 

Not only that, we further demonstrate that the form of linear attention described in Proposition <ref> can be extended to the more general case where Ã_ij is a d'-dimension vector rather than a scalar:

    Given that Ã∈ℝ^d'× n× n, if for each 1≤ u≤ d', Ã_u is low-rank separable: Ã_u=F_uG^⊤_u, where F_u,G_u∈ℝ^n× r, F_uiv=f(X_i)_uv, and G_ujv=g(X_j)_uv, there exist corresponding functions f' and g' such that the computation Y_i=C_iH_i=C_i∑_j=1^n{Ã_:ij⊙ (B_j^⊤ X_j)} can be equivalently implemented as linear attention, expressed as Y_i=f'(X_i)g'(X)^⊤ X, where Ã_:ij is a column vector and can broadcast to a d'× d matrix. 


    Given existing conditions, we have:
    
    Y_i   =∑_u=1^d'[c(X_i)_u{∑_j=1^n∑_v=1^r(f(X_i)_uvg(X_j)_uvb(X_j)_uX_j)}]
       =∑_u=1^d'∑_v=1^r[c(X_i)_uf(X_i)_uv∑_j=1^n{g(X_j)_uvb(X_j)_uX_j}]
       =vec(c(X_i)· f(X_i))[vec(b(X_j)· g(X_j))]_j^⊤X,

    where f(X_i)=F_:i: and g(X_j)=G_:j: are d'× r matrices, · denotes element-wise multiplication with broadcasting, and vec represents flatting a matrix into a row vector. 
    Defining f'(X_i)=vec(c(X_i)· f(X_i)) and g'(X_i)=vec(b(X_j)· g(X_j)), we derive Y=f'(X)g'(X)^⊤ X. 

From this point of view, the proposed structure can be deemed as a generalized linear attention and a non-causal form of recent linear-complexity sequential models, including Mamba2 <cit.>, RWKV6 <cit.>, GLA <cit.>, etc. 
In Tab. <ref>, we provide a summary of the parameterization in recent works for A_i. 


r9cm
[1]
  
  



Model     Parameterization of A_i     Causal 
 Mamba2 <cit.>     A_i∈ℝ     Yes 

mLSTM <cit.>     A_i∈ℝ     Yes 

Gated Retention <cit.>     A_i∈ℝ     Yes 
 
GateLoop <cit.>     A_i∈ℝ^d'     Yes 
 
HGRN2 <cit.>     A_i∈ℝ^d'     Yes 
 
RWKV6 <cit.>     A_i∈ℝ^d'     Yes 

Gated Linear Attention <cit.>     A_i∈ℝ^d'     Yes 
 MLLA <cit.>     A_ij=1      No 
 
VSSD <cit.>     A_ij∈ℝ     No 

Generalized Linear Attention     Ã_ij∈ℝ^d'     No 
 
A summary of the parameterization in recent linear token mixers for A_i, partially adapted from <cit.>.

  








 §.§ Training Objectives


In this paper, we replace all self-attention layers in the original SD with LinFusion modules. 
Only the parameters within these modules are trained, while all others remain frozen. 
To ensure that LinFusion closely mimics the original functionality of self-attention, we augment the standard noise prediction loss ℒ_simple in Eq. <ref> with additional losses. 
Specifically, we introduce a knowledge distillation loss ℒ_kd to align the final outputs of the student and teacher models, and a feature matching loss ℒ_feat to match the outputs of each LinFusion module and the corresponding self-attention layer. 
The training objectives can be written as:

    θ=_θ   𝔼_z∼ℰ(x),y,ϵ∼𝒩(0,1),t[ℒ_simple+αℒ_kd+βℒ_feat],
    ℒ_kd=‖ϵ_θ(z_t,t,y)-ϵ_θ_org(z_t,   t,y)‖_2^2,  ℒ_feat=1/L∑_l=1^L‖ϵ_θ^(l)(z_t,t,y)-ϵ_θ_org^(l)(z_t,t,y)‖_2^2,

where α and β are hyper-parameters controlling the weights of the respective loss terms, θ_org represents parameters of the original SD, L is the number of LinFusion/self-attention modules, and the superscript ^(l) refers to the output of the l-th one in the diffusion backbone. 














§ EXPERIMENTS




 §.§ Implementation Details







We present qualitative results on SD-v1.5, SD-v2.1, and SD-XL in Fig. <ref> 

and mainly conduct experiments on SD-v1.5 in this section. 
There are 16 self-attention layers in SD-v1.5 and we replace them with LinFusion modules proposed in this paper. 
Functions f' and g' mentioned in Proposition <ref> are implemented as MLP, which consists of a linear branch and a non-linear branch with one  block. 
Their results are added to form the outputs of f' and g'. 
The parameters of the linear branch in f' and g' are initialized as W_Q and W_K respectively, while the outputs of the non-linear branch are initialized as 0. 
We use only 169k images in LAION <cit.> with aesthetics scores larger than 6.5 for training and adopt the BLIP2 <cit.> image captioning model to regenerate the textual descriptions. 
Both hyper-parameters, α and β, are set as 0.5, following the approach taken in <cit.>, which also focuses on architectural distillation of SD. 
The model is optimized using AdamW <cit.> with a learning rate of 10^-4. 
Training is conducted on 8 RTX6000Ada GPUs with a total batch size of 96 under 512×512 resolution for 100k iterations, requiring ∼1 day to complete. 
The efficiency evaluations are conducted on a single NVIDIA A100-SXM4-80GB GPU. 



 §.§ Main Results




Ablation Studies. 
To demonstrate the effectiveness of the proposed LinFusion, we report the comparison results with alternative solutions such as those shown in Fig. <ref>(a), (b) and (c). 
We follow the convention in previous works focusing on text-to-image generation <cit.> and conduct quantitative evaluation on the COCO benchmark <cit.> containing 30k text prompts. 
The metrics are FID <cit.> against the COCO2014 test dataset and the cosine similarity in the CLIP-ViT-G feature space <cit.>. 
We also report the running time per image with 50 denoising steps and the GPU memory consumption during inference for efficiency comparisons. 
Results under 512×512 resolution are shown in Tab. <ref>. 

Mitigating Structural Difference. 
We begin our exploration from the original Mamba2 structure <cit.> with bi-directional scanning, i.e., Fig. <ref>(a), and try removing the gating and RMS-Norm, i.e., Fig. <ref>(b), to maintain a consistent holistic structure with the self-attention layer in the original SD. 
In this way, the only difference with the original SD lies on the SSM or self-attention for token mixing. 
We observe that such structural alignment is beneficial for the performance. 


Normalization and Non-Causality. 
We then apply the proposed normalization operation and the non-causal treatment sequentially, corresponding to Fig. <ref>(c) and (d). 
Although results in Tab. <ref> indicate that normalization would slightly hurt the performance, we will show in the following Tab. <ref> that it is crucial for generating images with resolutions unseen during training. 
Further adding the proposed non-causal treatment, we obtain results better than Fig. <ref>(b). 

We also compare the proposed non-causal operation with the simplified case mentioned in Sec. <ref>, achieved by directly removing the lower triangular causal mask applied on Ã, which results in a 1-rank matrix, i.e., various tokens share the same group of forget gates. 
The inferior results demonstrate the effectiveness of the proposed generalized linear attention. 



Attention Visualization. 
In Fig. <ref>, we visualize the self-attention maps yielded by various methods, including the original SD, bi-directional SSM, linear attention with shared forget gates, and generalized linear attention in LinFusion. 
Results indicate that our method works better for capturing broader-range of spatial dependency and best matches the predictions of the original SD. 

Knowledge Distillation and Feature Matching. 
We finally apply loss terms ℒ_kd and ℒ_feat in Eq. <ref>, which enhance the performance further and even surpass the SD teacher. 

Cross-Resolution Inference. 
It is desirable for diffusion model to generate images of unseen resolutions during training–a feature of the original SD. 
Since modules other than LinFusion are pre-trained and fixed in our work, normalization is a key component for this feature to maintain consistent feature distributions for training and inference. 
We report the results of 1024×1024 resolution in Tab. <ref>, which indicate that the conclusion holds for all the basic structures such as Mamba2, Mamba2 without gating and RMS-Norm, and the proposed generalized linear attention. 
Fig. <ref> shows a qualitative example, where results without normalization are meaningless. 

Ultrahigh-Resolution Generation. 
As discussed in <cit.>, directly applying diffusion models trained on low resolutions for higher-resolution generation can result in content distortion and duplication. 
In this paper, we address the challenge by handling the low resolution first, based on which higher-scale images are generated using SDEdit <cit.>. 
Note that in this work we aim at a linear-complexity model computationally enabling such ultrahigh-resolution generation as shown in Fig. <ref>. 
Dedicated designs are left as future directions. 



 §.§ Empirical Extensions


The proposed LinFusion is highly compatible with various components/plugins for SD, such as ControlNet <cit.>, IP-Adapter <cit.>, and LoRA <cit.>, without any further training or adaptation. 
We present qualitative results in Fig. <ref> and refer readers to the appendix for more results. 
The following quantitative evaluations indicate comparable performance of LinFusion with the original SD. 

ControlNet. 
ControlNet <cit.> introduces plug-and-play components to SD for additional conditions, such as edge, depth, and semantic map. 
We substitute SD with the proposed LinFusion and compare the FID, CLIP score, and the similarity between the input conditions and the extracted conditions from generated images of diffusion models with original SD. 
The results are shown in Tab. <ref>. 

IP-Adapter. 
Personalized text-to-image generation <cit.> is a popular application of SD, which focuses on generating images simultaneously following both input identities and textual descriptions. 
IP-Adapter <cit.> offers a zero-shot solution that trains a mapper from the image space to the condition space of SD, so that it can handle both image and text conditions. 
We demonstrate that IP-Adapter trained on SD can be used directly on LinFusion. 
The performance on the DreamBooth dataset <cit.>, containing 30 identities and 25 text prompts to form 750 test cases in total, is shown in Tab. <ref>. 
We use 5 random seeds for each case and report the averaged CLIP image similarity, DINO <cit.> image similarity, and CLIP text similarity. 

LoRA. 
Low-rank adapters (LoRA) <cit.> aim at low-rank matrices applied on the weights of a basic model such that it can be adapted for different tasks or purposes. 
For instance, <cit.> introduce LCM-LoRA such that the pre-trained SD can be used for LCM inference with only a few denoising steps <cit.>. 
Here, we directly apply LoRA in LCM-LoRA model to LinFusion. 
The performance on the COCO benchmark is shown in Tab. <ref>. 





§ CONCLUSION


This paper introduces a diffusion backbone termed LinFusion for text-to-image generation with linear complexity in the number of pixels. 
At the heart of LinFusion lies in a generalized linear attention mechanism, distinguished by its normalization-aware and non-causal operations—key aspects overlooked by recent linear-complexity token mixers like Mamba, Mamba2, and GLA. 
We reveal theoretically that the proposed paradigm serves as a general low-rank approximation for the non-causal variants of recent models. 
Based on Stable Diffusion (SD), LinFusion modules after knowledge distillation can seamlessly replace self-attention layers in the original model, ensuring that LinFusion is highly compatible to existing components for Stable Diffusion, like ControlNet, IP-Adapter, and LoRA, without any further training effort. 
Extensive experiments on SD-v1.5, SD-v2.1, and SD-XL demonstrate that the proposed model outperforms existing baselines and achieves performance on par with, or better than, the original SD with significantly reduced computational overhead. 
On a single GPU, it can accommodate image generation with resolutions up to 16K. 







iclr2024_conference



