PaperlayerView exact revision on arXiv

Source-linked AI summary

LinFusion: 1 GPU, 1 Minute, 16K Image

Songhua Liu, Weihao Yu, Zhenxiong Tan, Xinchao Wang

arXiv:2409.02097v3cs.CVcs.LG

TL;DR

Self-attention’s quadratic time and memory costs challenge high-resolution diffusion generation. LinFusion replaces self-attention with distilled, normalized non-causal linear attention; the paper reports SD-comparable or better performance and zero-shot generation up to 16K on one GPU.

  • Problem

    Self-attention’s quadratic time and memory complexity makes high-resolution visual generation challenging.

    Source: §1 · p. 2

  • Method

    LinFusion replaces Stable Diffusion’s self-attention layers with generalized, normalized non-causal linear attention modules trained through knowledge distillation.

    Source: §1 · p. 2, §2.2 · p. 4

  • Results

    LinFusion achieves performance on par with or superior to original SD while reducing time and memory complexity, with satisfactory zero-shot generation up to 16K on a single GPU.

    Source: §1 · p. 3, §4 · p. 10

  • Takeaways & Limitations

    LinFusion modules can work with existing Stable Diffusion components and pipelines without additional training cost.

    Source: §1 · p. 3, §2.2 · p. 4

    Evaluation uses PyTorch 1.13 to reflect theoretical complexity; higher versions apply block-wise memory strategies, but time efficiency remains a problem.

    Source: §1 · p. 2

Abstract

from arXiv · show

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 at https://github.com/Huage001/LinFusion.

1 INTRODUCTION

Self-attention supports diffusion models’ spatial relationships, but its quadratic cost challenges high-resolution generation. LinFusion proposes normalized, non-causal linear attention, distilled from Stable Diffusion, as an alternative.

  • Self-attention’s quadratic time and memory complexity challenges high-resolution generation; SD-v1.5 runs out of memory at 2048 resolution on an 80GB A100 in FP16.The paper motivates a token mixer with linear complexity in the number of spatial tokens.Source: §1 · p. 2
  • Mamba diffusion models face feature-distribution shift across resolutions and a causal restriction that conflicts with denoising from all spatial tokens.Bidirectional scanning can alleviate but does not eliminate the restriction within each branch.Source: §1 · p. 2
  • LinFusion normalizes aggregated features to keep total token impact consistent across input scales and uses distinct forget-gate groups with a low-rank approximation for non-causal linear attention.The paper characterizes the approach as a generalized low-rank approximation of popular linear-complexity token mixers.Source: §1 · pp. 2–3
  • The generalized linear attention module replaces self-attention layers in Stable Diffusion architectures.The resulting model is called LinFusion.Source: §1 · pp. 2–3
  • Evaluation uses PyTorch 1.13 to reflect theoretical complexity; higher versions use block-wise memory strategies, but time efficiency remains a problem.Source: §1 · p. 2
  • 50k iterations of distillation yield performance on par with or superior to original SD while significantly reducing time and memory complexity.Figure 2 compares generation speed and GPU memory across resolutions for LinFusion and SD-v1.5.Source: Figure 2 · p. 3, §1 · p. 3
  • LinFusion supports satisfactory zero-shot cross-resolution generation, including 16K images on a single GPU, and compatibility with SD components and pipelines.The paper reports experiments on SD-v1.5, SD-v2.1, and SD-XL.Source: §1 · p. 3

2 METHODOLOGY

LinFusion replaces Stable Diffusion’s quadratic self-attention with a distilled linear-attention module. Its design adds token normalization and non-causal information mixing, using low-rank separability to retain linear complexity.

2.1 PRELIMINARY

  • SD's UNet uses self-attention for spatial relations and cross-attention for text-image relations; its n-token attention matrix gives self-attention quadratic complexity.Source: §2.1 · pp. 3–4

2.2 OVERVIEW

  • LinFusion replaces SD self-attention with trainable modules distilled to match the original model's outputs, leaving the rest of the model frozen.This limits training to spatial relationships and preserves compatibility with existing SD components and fine-tuned variations.Source: §2.2 · p. 4

2.3 NORMALIZATION-AWARE MAMBA

  • When training and inference image scales differ, SSM-based Mamba2 can fail because channel-wise feature distributions shift with token count.LinFusion normalizes attention weights so their sum for each token equals 1, substantially improving zero-shot cross-resolution generalization.Source: §2.3 · pp. 4–5

2.4 NON-CAUSAL MAMBA

  • Removing the causal mask lets tokens access all tokens, but requires assigning different groups of A to input tokens to avoid uniform hidden states.Source: §2.4 · p. 5
  • Assuming the non-causal mask is low-rank separable allows its computation to be expressed as linear attention; MLPs approximate the resulting functions.The paper extends this form to vector-valued masks and frames it as generalized linear attention covering recent linear-complexity token mixers.Source: §2.4 · pp. 5–6

2.5 TRAINING OBJECTIVES

  • The training objective weights its respective loss terms with hyperparameters α and β and refers to outputs across the LinFusion/self-attention modules.Source: §2.5 · p. 6

3 EXPERIMENTS

The experiments evaluate LinFusion’s performance, cross-resolution behavior, and compatibility with Stable Diffusion components and pipelines. They report gains from distillation and feature matching, plus efficient generation at very high resolutions.

3.1 IMPLEMENTATION DETAILS

  • The SD-v1.5 experiments replace its 16 self-attention layers with LinFusion modules and train on 169k LAION images; efficiency is evaluated on one NVIDIA A100 GPU.Training uses 100k iterations on eight RTX6000Ada GPUs at 512 × 512 resolution.Source: §3.1 · pp. 6, 8

3.2 MAIN RESULTS

  • Removing Mamba2’s gating and RMS-Norm to align its structure with original Stable Diffusion self-attention benefits performance; proposed normalization slightly hurts, while adding proposed non-causal treatment outperforms Fig. 4(b).The simplified non-causal operation removes the lower-triangular causal mask, making various tokens share the same group of forget gates; it performs worse than proposed generalized linear attention.Source: §3.2 · p. 8
  • LinFusion’s attention maps capture broader spatial dependencies and best match the original Stable Diffusion predictions among the compared methods.The comparison includes the original model, bidirectional SSM, shared-forget-gate linear attention, and LinFusion’s generalized linear attention.Source: §3.2 · p. 8, Figure 6 · p. 8
  • Knowledge distillation and feature matching enhance LinFusion further, enabling it to surpass the Stable Diffusion teacher.The experiments apply loss terms Lkd and Lfeat.Source: §3.2 · p. 9
  • At 1024 × 1024, normalization is crucial for cross-resolution generation across the tested basic structures.This resolution was unseen during training; without normalization, the qualitative result is meaningless.Source: §3.2 · p. 9, Table 3 · p. 9

3.3 EMPIRICAL EXTENSIONS

  • LinFusion is reported as comparable to the original Stable Diffusion and compatible with its components and pipelines without further training or adaptation.The extensions include ControlNet, IP-Adapter, LoRA, DemoFusion, and DistriFusion.Source: §3.3 · p. 9
  • ∼2× acceleration at 2048 × 2048 is reported for LinFusion in the high-resolution pipeline settings.The pipeline also skips the former 60% of denoising steps while maintaining quality, and supports generation up to 16K on one GPU.Source: §3.3 · p. 10, Table 5 · p. 10
  • Linear complexity makes LinFusion suitable for distributed parallel inference, with constant communication cost independent of image resolution.The reported setup achieves satisfactory acceleration without NVLink hardware.Source: §3.3 · p. 10, Table 8 · p. 10

4 CONCLUSION

LinFusion replaces Stable Diffusion’s self-attention with a distilled, generalized linear-attention backbone. Experiments report performance on par with or better than Stable Diffusion, reduced computational overhead, compatibility with SD components, and generation up to 16K on one GPU.

4 CONCLUSION

  • LinFusion is a text-to-image diffusion backbone with linear complexity in the number of pixels, built around normalized, non-causal generalized linear attention.The paper characterizes the mechanism as a general low-rank approximation of non-causal variants of recent token mixers.Source: §4 · p. 10
  • Experiments on SD-v1.5, SD-v2.1, and SD-XL report performance on par with or better than the original Stable Diffusion, with reduced computational overhead.The distilled LinFusion modules replace the original model’s self-attention layers and are compatible with existing SD components and pipelines without further training.Source: §4 · p. 10
  • LinFusion accommodates image generation up to 16K resolution on a single GPU.Source: §4 · p. 10

A RELATED WORKS

The related-works section reviews efficient diffusion architectures and linear-complexity token mixers.

A RELATED WORKS

  • The review organizes prior work around efficient diffusion architectures and linear-complexity token mixers.Source: A RELATED WORKS · p. 16

A.1 EFFICIENT DIFFUSION ARCHITECTURES

This section surveys efficient diffusion architectures, contrasting reduced-step sampling with faster network inference. It focuses on architecture efficiency and the use of linear-complexity token mixers.

A.1 EFFICIENT DIFFUSION ARCHITECTURES

  • Efficient diffusion research includes both reducing sampling steps and designing faster network inference; this paper focuses on the latter bottleneck for high-resolution results.Source: §A.1 · p. 16
  • Prior state-space and linear-attention diffusion architectures retain causal token processing, although diffusion can access all noisy spatial tokens simultaneously.The paper removes this causal restriction and introduces a non-causal token mixer for diffusion.Source: §A.1 · p. 16
  • Architectural pruning of Stable Diffusion’s UNet stages and blocks addresses a focus orthogonal to this paper’s optimization of self-attention layers.The cited pruning approach targets text-to-image generation.Source: §A.1 · p. 16

A.2 LINEAR-COMPLEXITY TOKEN MIXERS

Linear-complexity token mixers offer alternatives to self-attention, but applying sequential models to non-causal vision tasks is challenging. LinFusion’s approach relaxes gating assumptions used by other non-causal vision mixers.

A.2 LINEAR-COMPLEXITY TOKEN MIXERS

  • Sequential token mixers such as Mamba, Mamba2, mLSTM, RWKV6, and GLA have linear complexity but are not designed for non-causal vision tasks.The section identifies adapting these sequential models to non-causal vision as the paper’s focus.Source: §A.2 · p. 16
  • MLLA uses Linear Attention without hidden-state gating, while VSSD shares one group of gating values across input tokens.Source: §A.2 · p. 16
  • The proposed model relaxes those gating assumptions to generalize non-causal versions of modern state-space models.Source: §A.2 · p. 16

B THEORETICAL PROOF

The theoretical proof establishes that low-rank-separable token-mixing computations can be equivalently expressed as linear attention. It gives this equivalence for multiple forms of the mixing factors.

B THEORETICAL PROOF

  • When the mixing matrix is low-rank separable, the resulting computation can be written as linear attention with corresponding functions f′ and g′.Source: B THEORETICAL PROOF · pp. 16–17
  • The proof extends the equivalence to channel-wise low-rank-separable mixing, expressing each output Yi through linear attention.Source: B THEORETICAL PROOF · p. 17
  • For one formulation, defining f′(X) and g′(X) from the input-dependent factors yields Y = f′(X)g′(X)⊤X.Source: B THEORETICAL PROOF · p. 17

C ADDITIONAL EXPERIMENTS

Additional experiments examine LinFusion in ultrahigh-resolution, distributed, and alternative-architecture settings. They report whole-image inference and comparable distributed results, alongside qualifications for SD-v2.1, SD-XL, and PixArt-Sigma.

Ultrahigh-Resolution Generation

  • LinFusion enables whole-image denoising on a single GPU in the DemoFusion pipeline, avoiding patch-wise treatment that can miss holistic text-image relationships.The paper illustrates the contrast with stars appearing on the astronaut’s body in the patch-wise DemoFusion result.Source: C ADDITIONAL EXPERIMENTS · p. 17
  • For DemoFusion, the paper proposes skipping initial high-resolution denoising steps because low-resolution denoising has already produced overall image structure.The authors report improved efficiency and greater robustness to high-resolution layout turbulence.Source: C ADDITIONAL EXPERIMENTS · p. 17, Figure 8 · p. 18

Distributed Parallel Inference

Results on More Architectures

  • The experiments cover SD-v1.5, SD-v2.1, SD-XL, and PixArt-Sigma, with quantitative results listed in Table 9 for the COCO benchmark.Source: C ADDITIONAL EXPERIMENTS · p. 18, Table 9 · p. 18
  • On SD-v2.1 and SD-XL, LinFusion performs slightly worse; the authors speculate that its ∼160K mostly sub-512 × 512 training samples may explain the gap.They suggest that more high-quality samples could benefit performance.Source: C ADDITIONAL EXPERIMENTS · p. 18
  • On PixArt-Sigma, replacing all self-attention layers yields unnatural results, while retaining 25% of them could largely alleviate the issue.The paper attributes the challenge speculatively to self-attention being DiT’s core and sole mechanism for managing token relationships.Source: C ADDITIONAL EXPERIMENTS · pp. 18–19, Figure 10 · p. 19

D LIMITATIONS

The paper treats replacing all self-attention layers as an exploratory configuration, not necessarily the optimal design. It identifies hybrid structures as a possible future direction.

D LIMITATIONS

  • Replacing all self-attention layers is not necessarily optimal; the paper proposes exploring hybrid structures instead.Source: D LIMITATIONS · p. 19
  • One suggested hybrid places attention on deep features with fewer tokens and more feature channels.The paper describes this as potentially important for DiT-based architectures and a future direction.Source: D LIMITATIONS · p. 19
  • The proposed hybrid design is presented as a promising direction to explore, not an established optimum.Source: D LIMITATIONS · p. 19
Loading 2409.02097v3…