LinFusion: 1 GPU, 1 Minute, 16K Image

## Abstract

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 enables satisfactory and efficient zero-shot cross-resolution generation, accommodating ultra-resolution images like 16K on a single GPU.
Moreover, it is highly compatible with pre-trained SD components and pipelines, such as ControlNet, IP-Adapter, DemoFusion, DistriFusion, etc , requiring no adaptation efforts.
Codes are available here .

## 1 Introduction

Recent years have witnessed significant advancements in AI-generated content (AIGC) with diffusion models Croitoru et al. (); Yang et al. ().
On the one hand, unlike classic models like GAN Goodfellow et al. (), diffusion models refine noise vectors iteratively to produce high-quality results with fine details Nichol & Dhariwal (); Dhariwal & Nichol (); Rombach et al. (); Ho et al. ().
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 Balaji et al. (); Ding et al. (); Nichol et al. (); Ramesh et al. (); Betker et al. (); Rombach et al. (); Saharia et al. ().
Benefiting from the impressive performance and the open-source community, Stable Diffusion (SD) Rombach et al. () 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 Ronneberger et al. (); Rombach et al. () to Vision Transformers Peebles & Xie (); Bao et al. (); Chen et al. (); Esser et al. (), 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. 2(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(^1^11PyTorch 1.13 is adopted here for evaluation to reflect the theoretical complexity of various architectures. On higher versions of PyTorch, block-wise strategies are applied for memory efficient attention. However, the time efficiency is still a problem.).

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 Gu & Dao () and Mamba2 Dao & Gu (), 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 Dao & Gu (), RWKV6 Peng et al. (), and Gated Linear Attention Yang et al. () and reveal that our model can be regarded as a generalized non-causal version of these popular models.

Figure: Figure 2: (a) and (b): Comparisons of the proposed LinFusion with original SD-v1.5 under various resolutions in terms of generation speed using 8 steps and GPU memory consumption. The dashed lines denote estimated values using quadratic functions due to out-of-memory error. (c) and (d): Efficiency comparisons on various architectures under their default resolutions.
Refer to caption: x2.png

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. 2.
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 and pipelines for SD, such as ControlNet Zhang et al. (), IP-Adapter Ye et al. (), DemoFusion Du et al. (), DistriFusion Li et al. (), etc, allowing users to achieve various purposes with the proposed LinFusion flexibly without any additional training cost.
As shown in Fig. 1, 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 and pipelines 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.

## 2 Methodology

### 2.1 Preliminary

Diffusion Models.
As a popular model for text-to-image generation, Stable Diffusion Rombach et al. () (SD) first learns an auto-encoder $(\mathcal{E},\mathcal{D})$, where the encoder $\mathcal{E}$ maps an image $x$ to a lower dimensional latent space: $z\leftarrow\mathcal{E}(x)$, and the decoder $\mathcal{D}$ learns to decode $z$ back to the image space $\hat{x}\leftarrow\mathcal{D}(z)$ such that $\hat{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 $\epsilon_{\theta}$ for $T$ steps.
The denoised latent code after the final step $z_{0}$ is decoded by $\mathcal{D}$ to derive a generated image.
In training, given an image $x$ and its corresponding text description $y$, $\mathcal{E}$ is utilized to obtain its corresponding latent code, and we add a random Gaussian noise $\epsilon$ for its noisy version $z_{t}$ with respect to the $t$-th step.
The UNet is trained via the noise prediction loss $\mathcal{L}_{simple}$ Ho et al. (); Nichol & Dhariwal ():

$$ $\theta=\operatorname*{arg\,min}_{\theta}\mathbb{E}_{z\sim\mathcal{E}(x),y, \epsilon\sim\mathcal{N}(0,1),t}[\mathcal{L}_{simple}]\quad\mathcal{L}_{simple} =\|\epsilon-\epsilon_{\theta}(z_{t},t,y)\|_{2}^{2}.$ (1) $$

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\in\mathbb{R}^{n\times d}$ and weight parameters $W_{Q},W_{K}\in\mathbb{R}^{d\times d^{\prime}}$ and $W_{V}\in\mathbb{R}^{d\times d}$, where $n$ is the number of spatial tokens, $d$ is the feature dimension, and $d^{\prime}$ is the attention dimension, self-attention can be formalized as:

$$ $Y=MV,\quad M=\mathrm{softmax}(QK^{\top}/\sqrt{d^{\prime}}),\quad Q=XW_{Q}, \quad K=XW_{K},\quad V=XW_{V}.$ (2) $$

We can observe from Eq. 2 that the complexity of self-attention is quadratic with respect to $n$ since the attention matrix $M\in\mathbb{R}^{n\times 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 Vaswani et al. (), Mamba Gu & Dao () 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}\odot H_{i-1}+B_{i}^{\top}X_{i}=\sum_{j=1}^{i}\{(\prod_{k=j+1}^{i}A _{k})\odot(B_{j}^{\top}X_{j})\},\quad Y_{i}=C_{i}H_{i},$ (3) $$

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 $\odot$ indicates element-wise multiplication.

### 2.2 Overview

Figure: Figure 3: Overview of LinFusion. We replace self-attention layers in the original SD with our LinFusion modules and adopt knowledge distillation to optimize the parameters.
Refer to caption: x3.png

In the latest version, i.e., Mamba2 Dao & Gu (), $A_{i}$ is a scalar, $B_{i},C_{i}\in\mathbb{R}^{1\times d^{\prime}}$, $X_{i},Y_{i}\in\mathbb{R}^{1\times d}$, and $H_{i}\in\mathbb{R}^{d^{\prime}\times d}$.
According to State-Space Duality (SSD), the computation in Eq. 3 can be reformulated as the following expression, referred to as 1-Semiseparable Structured Masked Attention:

$$ $Y=((CB^{\top})\odot\tilde{A})X,$ (4) $$

where $\tilde{A}$ is a $n\times n$ lower triangular matrix and $\tilde{A}_{ij}=\prod_{k=j+1}^{i}A_{k}$ for $j\leq i$.
Such a matrix $\tilde{A}$ 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. 3 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. 4(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 Hochreiter & Schmidhuber (); Cho () or RMS-Norm Zhang & Sennrich () as used in Mamba2.
As shown in Fig. 4(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. 2.3 and 2.4 respectively.
Finally, in Sec. 2.5, we provide the training objectives to optimize parameters in LinFusion modules.

### 2.3 Normalization-Aware Mamba

In practice, we find that SSM-based structure shown in Fig. 4(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. 3.2 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:

###### Proposition 1 .

Assuming that the mean of the $j$-th channel in the input feature map $X$ is $\mu_{j}$, and denoting $(CB^{\top})\odot\tilde{A}$ as $M$, the mean of this channel in the output feature map $Y$ is $\mu_{j}\sum_{k=1}^{n}M_{ik}$.

The proof is straightforward.
We observe through Fig. 4(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. 1, 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 Softmax 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., $\sum_{k=1}^{n}M_{ik}=1$, which is equivalent to applying the SSM module one more time to obtain the normalization factor $Z$:

$$ $Z_{i}=A_{i}\odot Z_{i-1}+B_{i},\quad C^{\prime}_{ij}=\frac{C_{ij}}{\sum_{k=1}^ {d^{\prime}}\{C_{ik}\odot Z_{ik}\}}.$ (5) $$

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

Figure: Figure 4: (a) The architecture of Mamba2. Bi-directional SSM is additionally involved here. (b) Mamba2 without gating and RMS-Norm. (c) Normalization-aware Mamba2. (d) The proposed LinFusion module with generalized linear attention.
Refer to caption: x4.png

###### Proposition 1 .

### 2.4 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 $\tilde{A}$.
Thus, the recursive formula in Eq. 3 would become $H_{i}=\sum_{j=1}^{n}\{(\prod_{k=j+1}^{n}A_{k})\odot(B_{j}^{\top}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\times n$ matrix and $H_{i}=\sum_{j=1}^{n}\{(\prod_{k=j+1}^{n}A_{ik})\odot(B_{j}^{\top}X_{j})\}$.
The $\tilde{A}_{ij}$ in Eq. 4 becomes $\prod_{k=j+1}^{n}A_{ik}$.
Compared with that in Eq. 4, $\tilde{A}$ here is not necessarily 1-semiseparable.
To maintain linear complexity, we impose the assumption that $\tilde{A}$ is low-rank separable, i.e., there exist input-dependent matrices $F$ and $G$ such that $\tilde{A}=FG^{\top}$.
In this way, the following proposition ensures that Eq. 4 under this circumstance can be implemented via linear attention:

###### Proposition 2 .

Given that $\tilde{A}=FG^{\top}$, $F,G\in\mathbb{R}^{n\times r}$, and $B,C\in\mathbb{R}^{n\times d^{\prime}}$, 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^{\prime}$ and $g^{\prime}$ such that Eq. 4 can be equivalently implemented as linear attention, expressed as $Y=f^{\prime}(X)g^{\prime}(X)^{\top}X$.

The proof can be found in the appendix.
In practice, we adopt two MLPs to mimic the functionalities of $f^{\prime}$ and $g^{\prime}$.
Combining with the normalization operations mentioned in Sec. 2.3, we derive an elegant structure shown in Fig. 4(d).

Not only that, we further demonstrate that the form of linear attention described in Proposition 2 can be extended to the more general case where $\tilde{A}_{ij}$ is a $d^{\prime}$-dimension vector rather than a scalar:

###### Proposition 3 .

Given that $\tilde{A}\in\mathbb{R}^{d^{\prime}\times n\times n}$, if for each $1\leq u\leq d^{\prime}$, $\tilde{A}_{u}$ is low-rank separable: $\tilde{A}_{u}=F_{u}G^{\top}_{u}$, where $F_{u},G_{u}\in\mathbb{R}^{n\times r}$, $F_{uiv}=f(X_{i})_{uv}$, and $G_{ujv}=g(X_{j})_{uv}$, there exist corresponding functions $f^{\prime}$ and $g^{\prime}$ such that the computation $Y_{i}=C_{i}H_{i}=C_{i}\sum_{j=1}^{n}\{\tilde{A}_{:ij}\odot(B_{j}^{\top}X_{j})\}$ can be equivalently implemented as linear attention, expressed as $Y_{i}=f^{\prime}(X_{i})g^{\prime}(X)^{\top}X$, where $\tilde{A}_{:ij}$ is a column vector and can broadcast to a $d^{\prime}\times d$ matrix.

The proof is provided in the appendix.
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 Dao & Gu (), RWKV6 Peng et al. (), GLA Yang et al. (), etc.
In Tab. 1, we provide a summary of the parameterization in recent works for $A_{i}$.

Table: Table 1: A summary of the parameterization in recent linear token mixers for $A_{i}$, partially adapted from Yang et al. ().

###### Proposition 2 .

###### Proposition 3 .

### 2.5 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 $\mathcal{L}_{simple}$ in Eq. 1 with additional losses.
Specifically, we introduce a knowledge distillation loss $\mathcal{L}_{kd}$ to align the final outputs of the student and teacher models, and a feature matching loss $\mathcal{L}_{feat}$ to match the outputs of each LinFusion module and the corresponding self-attention layer.
The training objectives can be written as:

$$ $\displaystyle\theta=\operatorname*{arg\,min}_{\theta}$ $\displaystyle\mathbb{E}_{z\sim\mathcal{E}(x),y,\epsilon\sim\mathcal{N}(0,1),t} [\mathcal{L}_{simple}+\alpha\mathcal{L}_{kd}+\beta\mathcal{L}_{feat}],$ (6) $\displaystyle\mathcal{L}_{kd}=\|\epsilon_{\theta}(z_{t},t,y)-\epsilon_{\theta_ {org}}(z_{t},$ $\displaystyle t,y)\|_{2}^{2},\quad\mathcal{L}_{feat}=\frac{1}{L}\sum_{l=1}^{L} \|\epsilon_{\theta}^{(l)}(z_{t},t,y)-\epsilon_{\theta_{org}}^{(l)}(z_{t},t,y) \|_{2}^{2},$ $$

where $\alpha$ and $\beta$ are hyper-parameters controlling the weights of the respective loss terms, $\theta_{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.

Figure: Figure 5: Qualitative text-to-image results by LinFusion based on various architectures.
Refer to caption: x5.png

## 3 Experiments

### 3.1 Implementation Details

Table: Table 2: Performance and efficiency comparisons with various baselines on the COCO benchmark.

Figure: Figure 6: Visualization of attention maps by various architectures. The prompt is “Astronaut in a jungle, cold color palette, muted colors, detailed, 8k”.
Refer to caption: x6.png

We present qualitative results on SD-v1.5, SD-v2.1, and SD-XL in Fig. 5
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^{\prime}$ and $g^{\prime}$ mentioned in Proposition 2 are implemented as MLP, which consists of a linear branch and a non-linear branch with one Linear-LayerNorm-LeakyReLU block.
Their results are added to form the outputs of $f^{\prime}$ and $g^{\prime}$.
The parameters of the linear branch in $f^{\prime}$ and $g^{\prime}$ are initialized as $W_{Q}$ and $W_{K}$ respectively, while the outputs of the non-linear branch are initialized as 0 0.
We use only 169k images in LAION Schuhmann et al. () with aesthetics scores larger than $6.5$ for training and adopt the BLIP2 Li et al. () image captioning model to regenerate the textual descriptions.
Both hyper-parameters, $\alpha$ and $\beta$, are set as $0.5$, following the approach taken in Kim et al. (), which also focuses on architectural distillation of SD.
The model is optimized using AdamW Loshchilov & Hutter () with a learning rate of $10^{-4}$.
Training is conducted on 8 RTX6000Ada GPUs with a total batch size of 96 under $512\times 512$ resolution for 100k iterations, requiring $\sim 1$ day to complete.
The efficiency evaluations are conducted on a single NVIDIA A100-SXM4-80GB GPU.

### 3.2 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. 4(a), (b) and (c).
We follow the convention in previous works focusing on text-to-image generation Kang et al. () and conduct quantitative evaluation on the COCO benchmark Lin et al. () containing 30k text prompts.
The metrics are FID Heusel et al. () against the COCO2014 test dataset and the cosine similarity in the CLIP-ViT-G feature space Radford et al. ().
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\times 512$ resolution are shown in Tab. 2.

Mitigating Structural Difference.
We begin our exploration from the original Mamba2 structure Dao & Gu () with bi-directional scanning, i.e., Fig. 4(a), and try removing the gating and RMS-Norm, i.e., Fig. 4(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. 4(c) and (d).
Although results in Tab. 2 indicate that normalization would slightly hurt the performance, we will show in the following Tab. 8 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. 4(b).

We also compare the proposed non-causal operation with the simplified case mentioned in Sec. 2.4, achieved by directly removing the lower triangular causal mask applied on $\tilde{A}$, 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.

Figure: Figure 7: Normalization is crucial for cross-resolution generation as demonstrated by the results on the COCO benchmark under $1024\times 1024$ resolution, which is unseen in training.
Refer to caption: x7.png

Attention Visualization.
In Fig. 6, 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 $\mathcal{L}_{kd}$ and $\mathcal{L}_{feat}$ in Eq. 6, 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\times 1024$ resolution in Tab. 8, 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. 8 shows a qualitative example, where results without normalization are meaningless.

### 3.3 Empirical Extensions

The proposed LinFusion is highly compatible with various components/pipelines for SD, such as ControlNet Zhang et al. (), IP-Adapter Ye et al. (), LoRA Hu et al. (), DemoFusion Du et al. (), DistriFusion Li et al. (), etc, without any further training or adaptation.
We present some qualitative results in Fig. 5 and refer readers to the appendix for more results.
The overall performance of LinFusion is comparable with the original SD.

ControlNet.
ControlNet Zhang et al. () 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. 10.

IP-Adapter.
Personalized text-to-image generation Gal et al. () is a popular application of SD, which focuses on generating images simultaneously following both input identities and textual descriptions.
IP-Adapter Ye et al. () 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 Ruiz et al. (), containing 30 identities and 25 text prompts to form 750 test cases in total, is shown in Tab. 12.
We use 5 random seeds for each case and report the averaged CLIP image similarity, DINO Caron et al. () image similarity, and CLIP text similarity.

LoRA.
Low-rank adapters (LoRA) Hu et al. () 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, Luo et al. () introduce LCM-LoRA such that the pre-trained SD can be used for LCM inference with only a few denoising steps Luo et al. ().
Here, we directly apply LoRA in LCM-LoRA model to LinFusion.
The performance on the COCO benchmark is shown in Tab. 12.

Ultrahigh-Resolution Generation.
As discussed in Huang et al. (); He et al. (), directly applying diffusion models trained on low resolutions for higher-resolution generation can result in content distortion and duplication.
A series of works are dedicated for higher-resolution image generation by leveraging off-the-shelf diffusion models Du et al. (); Lin et al. (; ); Haji-Ali et al. ().
However, limited by the quadratic-complexity self-attention, when applied for ultrahigh-resolution generation, existing approaches turn to patch-wise strategies to overcome the heavy computation burden Bar-Tal et al. (), which leads to inferior results, as shown in Settings A and B of Tab. 10.

Complementary to these methods, LinFusion addresses the computational overhead via generalized linear attention.
As shown in Settings C and D of Tab. 10, LinFusion achieves $\sim 2\times$ acceleration under $2048\times 2048$ resolution.
Instead of going through full denoising steps in the original DemoFusion Du et al. (), tricks in SDEdit Meng et al. () are additionally applied here so that the former 60% steps are skipped, which further enhances the efficiency without scarifying the quality.
Please refer to the appendix for more analysis.
Backed up by the linear-complexity LinFusion, such strategies enable ultrahigh-resolution generation up to 16K on a single GPU as shown in Fig. 1.

Distributed Parallel Inference.
LinFusion is friendly for distributed parallel inference benefiting from its linear complexity, given that the communication cost is constant with respect to image resolution.
Specifically, unlike the original DistriFusion Li et al. () requiring transmitting all the key and value tokens for self-attention communication, the transmission in LinFusion is $g^{\prime}(X)^{\top}X\in\mathbb{R}^{c^{\prime}\times c}$, which is not related with the number of image tokens.
In consequence, as shown in Tab. 3, LinFusion does not require NVLink hardware to achieve satisfactory acceleration.
Please refer to the appendix for qualitative examples.

## 4 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 or pipelines for Stable Diffusion, like ControlNet, IP-Adapter, LoRA, DemoFusion, DistriFusion, etc, 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.

Figure: Figure 9: Results of ControlNet on the original SD-v1.5 and LinFusion.

Figure: Figure 11: Results of IP-Adapter on the original SD-v1.5 and LinFusion.

Table: Table 3: Results of distributed parallel inference on a server with 8 RTX 4090 D GPUs. Benefiting from its linear complexity and constant communication cost among various patches, LinFusion is readily for distributed parallel inference with multiple GPUs. Compared with DistriFusion, it achieves more significant acceleration even without NVLink.

## Appendix A Related Works

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

### A.1 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 Song et al. (); Luo et al. (); Kim et al. (); Ma et al. (); Zhou et al. () 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 Fei et al. (), DiM Teng et al. (), DiG Zhu et al. (), Diffusion-RWKV Fei et al. (), DiffuSSM Yan et al. (), and Zigma Hu et al. ().
These works have successfully adapted recent state space models like Mamba Gu & Dao (), RWKV Peng et al. (), or Linear Attention Katharopoulos et al. () into diffusion architectures.
However, these architectures maintain a causal restriction for diffusion tasks, 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, Kim et al. () 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.

### A.2 Linear-Complexity Token Mixers

Despite the widespread adoption of Transformer Vaswani et al. () 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 Katharopoulos et al. (), State Space Model Gu et al. (), and their variants including Mamba Gu & Dao (), Mamba2 Dao & Gu (), mLSTM Beck et al. (); Peng et al. (), Gated Retention Sun et al. (), DFW Mao (); Pramanik et al. (), GateLoop Katsch (), HGRN2 Qin et al. (), RWKV6 Peng et al. (), GLA Yang et al. (), 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 Han et al. () employs Linear Attention Katharopoulos et al. () as token mixers in vision backbones without a gating mechanism for hidden states.
In VSSD Shi et al. (), 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.

## Appendix B Theoretical Proof

###### Proposition 1 .

Assuming that the mean of the $j$-th channel in the input feature map $X$ is $\mu_{j}$, and denoting $(CB^{\top})\odot\tilde{A}$ as $M$, the mean of this channel in the output feature map $Y$ is $\mu_{j}\sum_{k=1}^{n}M_{ik}$.

The proof is straightforward.

###### Proposition 2 .

Given that $\tilde{A}=FG^{\top}$, $F,G\in\mathbb{R}^{n\times r}$, and $B,C\in\mathbb{R}^{n\times d^{\prime}}$, 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^{\prime}$ and $g^{\prime}$ such that Eq. 4 of the main manuscript can be equivalently implemented as linear attention, expressed as $Y=f^{\prime}(X)g^{\prime}(X)^{\top}X$.

###### Proof.

Given existing conditions, we have:

$$ $\displaystyle(CB^{\top})\odot\tilde{A}$ $\displaystyle=[(c(X_{i})b^{\top}(X_{j}))\odot(f(X_{i})g^{\top}(X_{j}))]_{i,j}$ (7) $\displaystyle=[(\sum_{u=1}^{d^{\prime}}\{c(X_{i})_{u}b(X_{j})_{u}\})(\sum_{v=1 }^{r}\{f(X_{i})_{v}g(X_{j})_{v}\})]_{i,j}$ $\displaystyle=[\sum_{u=1}^{d^{\prime}}\sum_{v=1}^{r}\{(c(X_{i})_{u}f(X_{i})_{v })(b(X_{j})_{u}g(X_{j})_{v})\}]_{i,j}$ $\displaystyle=[(c(X_{i})\otimes f(X_{i}))(b(X_{j})\otimes g(X_{j}))^{\top}]_{i ,j},$ $$

where $\otimes$ denotes Kronecker product.
Defining $f^{\prime}(X_{i})=c(X_{i})\otimes f(X_{i})$ and $g^{\prime}(X_{i})=b(X_{i})\otimes g(X_{i})$, we derive $Y=f^{\prime}(X)g^{\prime}(X)^{\top}X$.
∎

###### Proposition 3 .

Given that $\tilde{A}\in\mathbb{R}^{d^{\prime}\times n\times n}$, if for each $1\leq u\leq d^{\prime}$, $\tilde{A}_{u}$ is low-rank separable: $\tilde{A}_{u}=F_{u}G^{\top}_{u}$, where $F_{u},G_{u}\in\mathbb{R}^{n\times r}$, $F_{uiv}=f(X_{i})_{uv}$, and $G_{ujv}=g(X_{j})_{uv}$, there exist corresponding functions $f^{\prime}$ and $g^{\prime}$ such that the computation $Y_{i}=C_{i}H_{i}=C_{i}\sum_{j=1}^{n}\{\tilde{A}_{:ij}\odot(B_{j}^{\top}X_{j})\}$ can be equivalently implemented as linear attention, expressed as $Y_{i}=f^{\prime}(X_{i})g^{\prime}(X)^{\top}X$, where $\tilde{A}_{:ij}$ is a column vector and can broadcast to a $d^{\prime}\times d$ matrix.

###### Proof.

Given existing conditions, we have:

$$ $\displaystyle Y_{i}$ $\displaystyle=\sum_{u=1}^{d^{\prime}}[c(X_{i})_{u}\{\sum_{j=1}^{n}\sum_{v=1}^{ r}(f(X_{i})_{uv}g(X_{j})_{uv}b(X_{j})_{u}X_{j})\}]$ (8) $\displaystyle=\sum_{u=1}^{d^{\prime}}\sum_{v=1}^{r}[c(X_{i})_{u}f(X_{i})_{uv} \sum_{j=1}^{n}\{g(X_{j})_{uv}b(X_{j})_{u}X_{j}\}]$ $\displaystyle=\mathrm{vec}(c(X_{i})\cdot f(X_{i}))[\mathrm{vec}(b(X_{j})\cdot g (X_{j}))]_{j}^{\top}X,$ $$

where $f(X_{i})=F_{:i:}$ and $g(X_{j})=G_{:j:}$ are $d^{\prime}\times r$ matrices, $\cdot$ denotes element-wise multiplication with broadcasting, and $\mathrm{vec}$ represents flatting a matrix into a row vector.
Defining $f^{\prime}(X_{i})=\mathrm{vec}(c(X_{i})\cdot f(X_{i}))$ and $g^{\prime}(X_{i})=\mathrm{vec}(b(X_{j})\cdot g(X_{j}))$, we derive $Y=f^{\prime}(X)g^{\prime}(X)^{\top}X$.
∎

###### Proposition 1 .

###### Proposition 2 .

###### Proof.

###### Proposition 3 .

###### Proof.

## Appendix C Additional Experiments

Figure: Figure 13: LinFusion is complementary to pipelines dedicated to high-resolution generation like DemoFusion. To enhance the performance, instead of working patch-by-patch, we handle the image as a whole benefiting from the efficient LinFusion architecture. Moreover, we reveal that skipping part of the denoising steps in the high-resolution stage can further improve the efficiency without hurting the performance. The prompt is “An astronaut floating in space. Beautiful view of the stars and the universe in the background.”.
Refer to caption: x8.png

Figure: Figure 14: LinFusion is complementary to pipelines for high-resolution generation like DistriFusion. Using constant communication cost among various GPU, it achieves highly comparable performance with single-GPU inference. The prompt is “Astronaut in a jungle, cold color palette, muted colors, detailed, 8k”.
Refer to caption: x9.png

Figure: Figure 15: By default, LinFusion replaces all self-attention layers in a diffusion backbone. When applied to DiT-based structures like PixArt-Sigma, this configuration often struggles to generate smooth results. Leaving a small part of original self-attention layers unmoved, e.g., 25%, could largely alleviate this challenge. The prompts are “A photo of beautiful mountain with realistic sunset and blue lake, highly detailed, masterpiece” and “dog” respectively.
Refer to caption: x10.png

Ultrahigh-Resolution Generation.
We present qualitative examples to illustrate the effectiveness of LinFusion on ultrahigh-resolution generation in Fig. 13.
We build LinFusion upon DemoFusion Du et al. (), a pipeline dedicated to high-resolution generation.
In the original implementation, for efficiency, DemoFusion handles a high-resolution image patch-by-patch and averages the outputs of overlapped areas Bar-Tal et al. ().
However, we find that such a patch-wise treatment largely ignores the holistic text-image relationships.
As shown in Fig. 13(DemoFusion), there are stars on the body of the astronaut.
With an efficient architecture introduced by LinFusion, we do not have to conduct inference patch-by-patch.
Instead, the whole image can be accommodated to a single GPU for denoising, which addresses the above limitation effectively as shown in Fig. 13(Full Steps).

Moreover, DemoFusion has to conduct full steps in the high-resolution denoising stage, which would introduce significant latency.
Motivated from the insight in SDEdit Meng et al. () that early denoising steps tend to take over the overall image layouts, we propose to skip some initial steps in the high-resolution stage, given that the overall image structures have been produced in the low-resolution stage.
We find that it not only improves the efficiency but also makes the pipeline more robust to the turbulence on image layout in the high-resolution stage, as shown in Fig. 13(40% Steps).

Distributed Parallel Inference.
We supplement qualitative results of distributed parallel inference by building LinFusion upon DistriFusion Li et al. () in Fig. 14.
Using constant communication cost among various GPUs, it achieves highly comparable performance with single-GPU inference.
Unlike the original DistriFusion, LinFusion offers significant acceleration using multiple GPUs without the dependency on NVLink.

**Table 4: Performance of LinFusion built upon various models on the COCO benchmark.**
| Setting | FID($\downarrow$) | CLIP-T($\uparrow$) |
| --- | --- | --- |
| SD-v1.5 | 12.86 | 0.321 |
| w. LinFusion | 12.57 | 0.323 |
| SD-v2.1 | 12.84 | 0.333 |
| w. LinFusion | 13.84 | 0.329 |
| SD-XL | 14.74 | 0.340 |
| w. LinFusion | 15.72 | 0.338 |
| PixArt-Sigma | 26.32 | 0.334 |
| w. 75% LinFusion | 24.32 | 0.327 |

Results on More Architectures.
We conduct experiments on a variety of diffusion architectures in this paper, including SD-v1.5, SD-v2.1 Rombach et al. (), SD-XL Podell et al. (), and PixArt-Sigma Chen et al. ().
The former three adopt transformer-based UNet while the last one is based on DiT Peebles & Xie (), a pure-Transformer structure.
Their quantitative results are listed in Tab. 4.

We find that on SD-v2.1 and SD-XL, LinFusion leads to slightly inferior results.
We speculate that the reason lies in the training data used for LinFusion, which consists of only $\sim$160K relatively low-resolution samples, the majority of which are below $512\times 512$ resolution.
Involving more high-quality samples can benefit the performance.

On PixArt-Sigma, we find that replacing all the self-attention layers in the DiT would result in unnatural results, as shown in Fig. 15.
We speculate that the challenge arises because self-attention is the core and sole mechanism for managing token relationships in DiT.
Replacing these layers entirely with LinFusion may create a significant divergence from the original architecture, leading to difficulties during training.
As shown in Fig. 15, leaving a small part of original self-attention layers unchanged, e.g., 25%, could largely alleviate this challenge.

## Appendix D Limitations

The motivation of LinFusion is to explore a linear-complexity diffusion architecture by experimentally replacing all the self-attention layers with the proposed generalized linear attention.
This may not be the optimal configuration in practice.
For example, it could be promising to explore hybrid structures and apply attention on deep features with a relatively smaller number of tokens but a large number of feature channels, which appears to be crucial for DiT-based architectures and could be a meaningful future direction.

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 6
presentation: 3
questions: - Why does removing the patchification operation in DemoFusion in Table 5 (A -> B) increase the sampling speed? Since patchification typically reduces training and sampling costs, it seems counterintuitive.
- How the 25% of unremoved self-attention layers in PixArt-Sigma selected? Are they from shallow layers, deep layers or just randomly sampled? Can distillation twice alleviate this problem (e.g., replace 50% in the first time and train the model as described, followed by replacing the other half in the second time)? 
- Minor typo: In Table 3, it should read "Bi-Directional Mamba2 w/o Gating & RMS-Norm + Normalization."
soundness: 3
strengths: - The motivation is clear and well-founded, with a thorough analysis of existing linear attention mechanisms to identify the key factors contributing to their effectiveness in diffusion.
- Extensive experiments across various applications support the claims that LinFusion is both efficient and generalizable to different diffusion models as well as existing training and testing pipelines.
- Overall, the writing is fluent and easy to follow, with informative figures that provide ample supporting information.
summary: This paper introduces LinFusion, a versatile pipeline designed to enhance GPU memory efficiency and boost sampling speed across various diffusion models for image generation. Specifically, LinFusion investigates recent linear-attention mechanisms to identify key factors that enable their effectiveness in diffusion models and subsequently proposes an improved, generalized linear attention to replace standard self-attention. To simplify training, the models are not trained from scratch; instead, LinFusion selectively distills its linear attention module from the original diffusion models, keeping all other weights fixed. Additional supervision is applied to align both the final output and intermediate feature representations. Extensive experiments demonstrate that LinFusion can be effectively integrated with different diffusion models, significantly accelerating image generation.
weaknesses: - The comparisons are conducted only during the sampling stage. Since the proposed LinFusion module may also provide similar benefits during training, are there any metrics available for this stage?
- Related to the previous point, the paper includes only fine-tuning experiments. It would be valuable to investigate whether training a diffusion model from scratch with LinFusion replacing self-attention results in any performance drop. If so, what is the extent of this drop? Experiments on a class-conditional image generation task would be informative, even without a large-scale text-to-image model.
- The memory and efficiency comparisons primarily utilize PyTorch 1.13, which does not incorporate memory-efficient methods like flash-attention or flash-attention v2. How significant is the difference in memory consumption and sampling efficiency when these newer techniques are considered?
- The method by which LinFusion generates ultra-high-resolution images is somewhat unclear. Can LinFusion directly generate 16K-resolution images, thereby avoiding patch-wise splitting, or does it produce a lower-resolution image that is later upsampled with techniques like SDEdit?

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 	Please address questions in "Weaknesses".
soundness: 3
strengths: 	Compared to the original SD-v1.5, LinFusion offers significant advantages in speed and GPU memory usage for generating high-resolution images.
	The extensive amount of open-sourcing and experiment reproducibility is greatly appreciated.
summary: This work introduces a generalized linear attention paradigm and extracts knowledge from Stable Diffusion to develop a distilled model called LinFusion. LinFusion avoids the quadratic increase in complexity associated with traditional attention mechanisms as the number of tokens grows, enabling the efficient generation of high-resolution visual content. Extensive experiments demonstrate that LinFusion achieves satisfactory and efficient zero-shot cross-resolution generation.
weaknesses: 	The comparison experiments in the paper are not comprehensive; for instance, the experimental section lacks an analysis of parameters and data size. 
	It is unclear whether LinFusion can outperform the latest lightweight diffusion methods, such as BK-SDM[1], on the COCO 256×256 30K dataset.
	In Table 7, LinFusion shows a significant decrease in FID scores. In contrast, LinFusion exhibits better compatibility with other components and pipelines of SD, which would be better to analyze why this occurs.
	The visual quality of the generated images does not seem particularly impressive.
[1]Kim B K, Song H K, Castells T, et al. Bk-sdm: A lightweight, fast, and cheap version of stable diffusion[J]. arXiv preprint arXiv:2305.15798, 2023.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: Please referring to the Weaknesses above.
soundness: 3
strengths: - The paper presents an efficient text-to-image model, LinFusion, which innovatively addresses the computational inefficiencies inherent in high-resolution image generation with diffusion models.
- Two notable innovations of LinFusion are normalization-aware mamba and non-causal mamba, which significantly improving the model's performance.
- The authors have conducted an extensive set of experiments, demonstrating LinFusion's effectiveness across various resolutions and showcasing its superior capability in generating ultra-high-resolution images like 16K on a single GPU.
- The writing is clear and methodical, effectively guiding readers through the complex technical details while maintaining a focus on the practical implications of the research.
- The paper stands out for its thorough experimental validation, which not only benchmarks LinFusion against existing models but also integrates it with various components and pipelines, highlighting its versatility and compatibility in real-world applications.
summary: This paper introduces a novel text-to-image model named LinFusion, addressing the challenge of generating high-resolution visual content with diffusion models. To optimize this, the authors propose to ultilize the popular linear attention and present methods for normalization-aware and non-causal operations, achieving performance on par with or even superior
to the original diffusion model while significantly reducing time and memory complexity. Experiments demonstrate that LinFusion achieves comparable or superior performance to the original Stable Diffusion on tasks like zero-shot cross-resolution generation, with excellent results on MS COCO.
weaknesses: - While LinFusion demonstrates significant improvements in computational efficiency, mamba2 is designed for language models. Could you give more comparison with state-of-the-art linear attention methods[1,2,3] in computer vision.

- The results of the experiment are unconvincing. Could provide a more holistic assessment of LinFusion's performance across different aspects of image generation., such as HPSv2，T2I_Combench，DPG？

- Could the linear attention combined with the MM-DiT blocks，which are popular in state-of-the-art diffusion models，SD3 and FLUX?

- The training cost of this method is low. Is the training sufficient? Will longer training or increasing the training resolution (1k or 2k) further improve the model effect?

[1] Cai, Han, et al. "Efficientvit: Multi-scale linear attention for high-resolution dense prediction." arXiv preprint arXiv:2205.14756 (2022).

[2] Zhu, Lianghui, et al. "DiG: Scalable and Efficient Diffusion Models with Gated Linear Attention." arXiv preprint arXiv:2405.18428 (2024).

[3] Zhu, Lianghui, et al. "Vision mamba: Efficient visual representation learning with bidirectional state space model." arXiv preprint arXiv:2401.09417 (2024).

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
questions: 1. What about the influences of each loss in the training objective?
2. In L512 - L516, LinFusion uses SDEdit tricks for higher-resolution generation, and comparing with DemoFusion. Is this comparison fair? LinFusion is a specific base model. A better comparison should be to use the same settings as DemoFusion.
soundness: 3
strengths: 1. The usage of linear attention and network distillation is an interesting direction in efficient text-to-image generation.
2. Different Mamba architecture is considered and ablated, providing a reference for other related topics. 
3. This paper is well-written. The detailed experiments show the advantages of the proposed method and even some down-streaming plug-and-play applications.
summary: This paper presents a new method for efficient image generation in linear computation complex. To achieve this goal, a generalized linear attention paradigm is introduced. For training, they distill the knowledge from the pre-trained SD. This distilled LinFusion model achieves better or on-par performance than the original teacher. This model also enables many down-streaming plugins, which makes this paper far more interesting.
weaknesses: My main concern about his paper is the experiments.

1. This paper introduces three loss functions, i.e., $\mathcal{L}_{simple}$.
  $\mathcal{L}{kd}$, $\mathcal{L}{feat}$ and two hyper-parameters ($\alpha$, $\beta$). However, there is no ablation study for each of them.
2. Several down-streaming extensions (LoRA, ControlNet, etc.) have been evaluated. However, there is no further discussion of why the original network extensions work and why some of the results are better than the baseline and others are not.

Meta-review: This paper introduces LinFusion, a versatile pipeline designed to enhance GPU memory efficiency and accelerate sampling speeds in various diffusion models for image generation. LinFusion investigates recent linear-attention mechanisms to identify key factors contributing to their effectiveness in diffusion models. It then proposes a generalized and improved linear attention module as a replacement for standard self-attention. To simplify training, LinFusion avoids training models from scratch. Instead, it selectively distills the linear attention module from existing diffusion models, keeping all other weights fixed. Additional supervision is applied to align both the final outputs and intermediate feature representations.


This work received mixed reviews from four reviewers before rebuttal : one reviewer supported acceptance, one reviewer leaned toward borderline acceptance, and two reviewers gave a borderline rejection. After rebuttal, one reviewer raise their score to borderline acceptance and one reviewer maintain their score as borderline rejection. 

After carefully evaluating the reviews, rebuttal, and refined draft, the Area Chair (AC) agrees with the concerns raised by Reviewer WZLt and the weaknesses highlighted by Reviewer cB7B, particularly regarding unconvincing experimental results. Furthermore, the quality of high-resolution generated images remains subpar, weakening the submission despite the demonstrated improvements in inference speed through linear attention.

Thus, AC recommends rejection to this work.

Decision: Reject