PaperlayerView exact revision on arXiv

Source-linked AI summary

CLEAR: Conv-Like Linearization Revs Pre-Trained Diffusion Transformers Up

Songhua Liu, Zhenxiong Tan, Xinchao Wang

arXiv:2412.16112v1cs.CV

TL;DR

Quadratic attention makes high-resolution DiT generation slow, while it remains unclear whether existing efficient attention mechanisms can linearize pre-trained DiTs. CLEAR applies local attention and, after fine-tuning, reports results comparable to FLUX.1-dev with reduced attention computation and faster 8K generation.

  • Problem

    Quadratic attention creates substantial latency in high-resolution image generation, and applying existing efficient attention mechanisms to pre-trained DiTs remains uncertain.

    Source: §1 · p. 2

  • Method

    CLEAR linearizes pre-trained DiTs by restricting each query's attention to a local neighborhood.

    Source: §1 · p. 2

  • Results

    99.5% fewer attention computations and 6.3× acceleration for 8K-resolution generation accompany results comparable to original FLUX.1-dev after 10K iterations on 10K self-generated samples.

    Source: §1 · p. 2, §5 · p. 8

  • Takeaways & Limitations

    The distilled attention layers generalize zero-shot across models and plugins and support multi-GPU parallel inference.

    Source: §1 · p. 2, §5 · p. 8

    Practical acceleration falls short of FLOPS-based expectations, is less significant at low resolutions, and can be slower than the original DiT below 1024 × 1024.

    Source: §5 · p. 8

Abstract

from arXiv · show

Diffusion Transformers (DiT) have become a leading architecture in image generation. However, the quadratic complexity of attention mechanisms, which are responsible for modeling token-wise relationships, results in significant latency when generating high-resolution images. To address this issue, we aim at a linear attention mechanism in this paper that reduces the complexity of pre-trained DiTs to linear. We begin our exploration with a comprehensive summary of existing efficient attention mechanisms and identify four key factors crucial for successful linearization of pre-trained DiTs: locality, formulation consistency, high-rank attention maps, and feature integrity. Based on these insights, we introduce a convolution-like local attention strategy termed CLEAR, which limits feature interactions to a local window around each query token, and thus achieves linear complexity. Our experiments indicate that, by fine-tuning the attention layer on merely 10K self-generated samples for 10K iterations, we can effectively transfer knowledge from a pre-trained DiT to a student model with linear complexity, yielding results comparable to the teacher model. Simultaneously, it reduces attention computations by 99.5% and accelerates generation by 6.3 times for generating 8K-resolution images. Furthermore, we investigate favorable properties in the distilled attention layers, such as zero-shot generalization cross various models and plugins, and improved support for multi-GPU parallel inference. Models and codes are available here: https://github.com/Huage001/CLEAR.

1. Introduction

Quadratic attention makes high-resolution DiT generation slow, and whether existing efficient attention mechanisms can linearize pre-trained DiTs is unclear. CLEAR uses local attention and reports results comparable to FLUX.1-dev alongside substantial 8K-generation efficiency gains.

1. Introduction

  • Initial substitutions show formulation-variation methods do not transfer as they do in scratch-trained models, key-value compression distorts details, and sampling highlights the need for local tokens.Source: §1 · p. 2
  • The paper identifies locality, formulation consistency, high-rank attention maps, and feature integrity as crucial for linearizing pre-trained DiTs.Source: §1 · p. 2
  • CLEAR restricts each query to tokens within a predefined distance r, producing linear complexity with respect to image resolution.The number of key-value tokens interacting with each query remains fixed.Source: §1 · p. 2
  • CLEAR reduces attention computations by 99.5% and accelerates the original DiT by 6.3 times for 8K-resolution generation.Results are comparable to original FLUX.1-dev after 10K fine-tuning iterations on 10K self-generated samples.Source: §1 · p. 2
  • Distilled attention layers work across FLUX variants and ControlNet plugins without adaptation, while local interactions support multi-GPU parallel inference.Source: §1 · p. 2

2. Efficient Attention: A Taxonomic Overview

The section groups efficient attention approaches into formulation variation, key-value compression, and key-value sampling. These methods seek to reduce the quadratic cost of constructing pairwise attention maps.

2. Efficient Attention: A Taxonomic Overview

  • Self-attention constructs n × n attention maps, motivating efficient methods grouped into formulation variation, key-value compression, and key-value sampling.Source: §2 · p. 3

2.1. Formulation Variation

  • Formulation-variation methods omit softmax to compute KᵀV with linear time in n, then use kernel functions or alternative operations to approximate softmax.Examples include Mamba2, Gated Linear Attention, sigmoid, relu2, and Nyström-based approximation.Source: §2.1 · p. 3

2.2. Key-Value Compression

  • Key-value compression reduces complexity by making the key-value token count m smaller than the query count n, using downsampling, agent tokens, slots, or learned maps.Source: §2.2 · p. 3

2.3. Key-Value Sampling

  • Key-value sampling prunes original tokens per query, with approaches including regular-interval sampling, grouped routing, local windows, and combinations with random or global tokens.Source: §2.3 · p. 3

3. Methods

The section identifies four requirements for linearizing pre-trained DiTs, then introduces CLEAR’s local attention design and its training and parallel-inference procedures.

3.1. What are Crucial for Linearizing DiTs?

  • The four criteria are locality, softmax-formulation consistency, sufficiently high-rank attention maps, and integrity of raw query, key, and value features.The supplied comparisons report distorted textures with KV-compression methods, failures or convergence problems with alternative formulations, and blocky patterns for methods with low-rank attention maps.Source: §3.1 · pp. 3–4
  • Pre-trained DiTs rely on local features: perturbing remote features has little impact, while local perturbations can collapse image quality.Attention maps show significant scores concentrated near each query; for a 64 × 64 feature map, remote-distance clipping remains reasonable at r=8, while local perturbation collapses results at r=2.Source: §3.1 · pp. 3–4, Figure 5 · p. 4

3.2. Conv-Like Linearization

  • CLEAR keeps text queries global and lets each image query attend to all text tokens and image tokens within a circular local window, yielding linear complexity as resolution increases.The circular window includes key-value tokens within Euclidean distance r; its computation overhead is about π/4 that of the corresponding square window.Source: §3.2 · p. 4, Figure 6 · p. 5

3.3. Training and Optimization

  • Fine-tuning uses flow-matching and teacher-consistency losses on predictions and attention outputs, with only attention-layer parameters trainable.The authors report that samples generated by the original DiT yield significantly better training results than a real-image dataset, even when the latter contains more higher-quality data.Source: §3.3 · p. 5

3.4. Multi-GPU Parallel Inference

  • Local attention supports patch-parallel inference with communication limited to image-patch boundaries, costing O(r × W) between adjacent GPUs rather than O(H × W).Text-token attention still requires global key-value synchronization in the exact computation, but the authors report that averaging patch-wise attention approximates it without hurting performance too much.Source: §3.4 · p. 5

4. Experiments

CLEAR is evaluated as a linearized FLUX model against the original and in resolution extrapolation, cross-model and plugin settings, and parallel inference. The experiments also compare training data and report implementation and evaluation settings.

4.1. Settings and Implementation Details

  • The study replaces FLUX-1.dev attention layers with CLEAR and fine-tunes on 10K self-generated 1024 × 1024 samples for 10K iterations; evaluation uses 5,000 COCO2014 images and several image-quality and similarity metrics.Experiments use three window sizes, r = 8, 16, and 32; training uses batch size 32 on four H100 GPUs.Source: §4.1 · p. 6

4.2. Main Comparisons

  • CLEAR achieves performance comparable to or better than original FLUX-1.dev while using less computation; its linearized outputs preserve layout, texture, and tone.At r = 32, the CLIP Image score exceeds 90.Source: §4.2 · pp. 6–7, Figure 8 · p. 7

4.3. Empirical Studies

  • MS-SSIM reaches 0.9 in high-resolution comparisons; with SDEdit, editing strength controls the trade-off between fine detail and content preservation.The comparison uses 1,000 COCO2014 validation images at 2048 × 2048 and 4096 × 4096.Source: §4.3 · p. 7, Table 3 · p. 7, Figure 9 · p. 8
  • CLEAR layers trained on FLUX-1.dev transfer zero-shot to FLUX-1.schnell, producing results similar to the original FLUX-1.schnell.The target model supports four-step inference.Source: §4.3 · p. 7, Table 4 · p. 7
  • CLEAR supports the pre-trained ControlNet with grayscale conditions, with quantitative results reported as comparable to the original model.The evaluation uses 1,000 COCO2014 validation images.Source: §4.3 · p. 7, Table 5 · p. 8, Figure 9 · p. 8
  • Patch-wise average approximates information gathering for multi-GPU inference, and the reported semantic correlation before and after approximation is high.The results are evaluated across different numbers of patches.Source: §4.3 · p. 7, Table 6 · p. 8
  • Fine-tuning on FLUX-1.dev-generated samples outperforms fine-tuning on real LAION5B data, and synthetic-data training produces much lower loss.The authors speculate that a distribution mismatch between LAION and FLUX makes fine-tuning more difficult.Source: §4.3 · p. 8

5. Conclusions

The conclusion presents CLEAR as a convolution-like local attention strategy for linearizing pre-trained DiTs, with results comparable to the original model and substantial 8K efficiency gains. It also notes that practical acceleration can fall short of FLOPS-based expectations.

5. Conclusions

  • Fine-tuning on real data performs worse than fine-tuning on self-generated synthetic data.Source: Figure 10 · p. 8
  • CLEAR restricts attention to a circular local window, achieving linear complexity while retaining results comparable to the original model.The approach is grounded in locality, formulation consistency, high-rank attention maps, and feature integrity.Source: §5 · p. 8
  • Fine-tuning on 10K self-generated samples reduces attention computations by 99.5% and accelerates 8K image generation by 6.3×.The conclusion describes this as efficient knowledge transfer to a student model.Source: §5 · p. 8
  • CLEAR’s distilled attention layers support zero-shot generalization across models and plugins and improve multi-GPU parallel inference capabilities.The conclusion presents these properties as broadening applicability and scalability.Source: §5 · p. 8
  • Practical acceleration falls short of FLOPS-based expectations and can be less significant at lower resolutions, or slower than the original below 1024 × 1024.The authors attribute part of the gap to the greater difficulty of hardware optimization for sparse attention than for FlashAttention.Source: §5 · p. 8

A. Details of Efficient Attention Alternatives

This section surveys efficient-attention alternatives, including token compression, sparse attention, and changes to the attention formulation. It notes that several alternatives remain unsatisfactory when used to fully linearize a pre-trained DiT.

Linear Attention

  • Linear attention avoids softmax and uses associative multiplication to compute K^T V first, yielding linear complexity.The described implementation applies ELU-based non-negative functions to Q and K and normalizes each query's similarity scores.Source: A. Details of Efficient Attention… · p. 12

Sigmoid Attention

  • Sigmoid Attention replaces softmax with a sigmoid formulation to remove softmax normalization and accelerate attention.The implementation follows FlashSigmoid's official hardware-aware optimization.Source: A. Details of Efficient Attention… · p. 12

PixArt-Sigma

  • PixArt-Sigma downsamples key-value token maps with learnable group-wise Conv4×4 kernels at stride 4, initialized to act as average pooling.The passage reports that such compression can work in relatively deep DiT layers, but remains unsatisfactory for completely linearizing a pre-trained DiT.Source: A. Details of Efficient Attention… · p. 12

Agent Attention

  • Agent Attention uses downsampled query tokens to form agent tokens, which then serve as value tokens.This is presented as an efficient attention operation using agent tokens.Source: A. Details of Efficient Attention… · p. 12

Slot Attention

  • Slot Attention adaptively aggregates key-value tokens into s memory slots, using learnable P to model each input token's writing intensity.The resulting slots replace the original key-value tokens for attention computation.Source: A. Details of Efficient Attention… · p. 12

Sparse Attention

  • Strided Attention samples tokens at regular intervals, while Swin uses independent attention within sliding windows; the described Swin setup uses size 16 and shift 8 on odd-indexed layers.The image-to-image mask rank equals the number of windows, which challenges the high-rank requirement for linearizing DiTs.Source: A. Details of Efficient Attention… · p. 12

B. Training Dynamics

Training-loss curves for efficient-attention alternatives on FLUX-1.dev support the paper's conclusion about convergence.

Training Dynamics

  • Alternatives meeting locality, formulation consistency, high-rank attention-map, and feature-integrity requirements yield the most satisfactory training convergence.The supplementary curves are reported for FLUX-1.dev.Source: B. Training Dynamics · p. 13

C. Raw Data for Efficiency Comparisons

The supplement provides raw data for the efficiency comparisons in Figure 2.

Raw Data for Efficiency Comparisons

D. Results on More DiTs

CLEAR is also evaluated on StableDiffusion3.5-Large using the COCO2014 validation dataset. The reported qualitative comparisons show overall comparable performance, while the authors note a possible weakness in capturing global structure.

Results on More DiTs

  • CLEAR is evaluated on StableDiffusion3.5-Large with default r = 16, which experiments identify as the best quality-efficiency trade-off, and results are reported on COCO2014 validation.The results are presented in Table 8.Source: D. Results on More DiTs · p. 13
  • Qualitative comparisons with original FLUX-1.dev and SD3.5-L indicate overall comparable performance, but CLEAR may underperform on global structure such as symmetry.The authors suggest that adding global tokens could potentially mitigate this issue, leaving further design exploration to future work.Source: D. Results on More DiTs · p. 13

High-Resolution Inference

  • CLEAR-linearized DiTs are described as compatible with pipelines designed for high-resolution inference.Source: Figure 12 · p. 13

E. More High-Resolution Results

CLEAR is compatible with resolution-extrapolation pipelines beyond the paper’s SDEdit setup, including I-Max, where it may balance low-resolution content preservation and high-resolution detail.

E. More High-Resolution Results

  • CLEAR is compatible with resolution-extrapolation pipelines beyond the paper’s SDEdit setup, including the training-free I-Max pipeline.The authors report that this combination may better balance preserving low-resolution content and capturing high-resolution details.Source: E. More High-Resolution Results · p. 13
  • In the I-Max example, dress textures from the low-resolution result are preserved with minimal variation while clear high-resolution details are enhanced.The authors present this as an example of the balance observed with CLEAR and I-Max.Source: E. More High-Resolution Results · p. 13
  • The paper’s main CLEAR setup uses SDEdit for coarse-to-fine high-resolution generation from a conditional image.SDEdit is described as a simple yet effective strategy for this setting.Source: E. More High-Resolution Results · p. 13

F. More ControlNet Results

CLEAR is evaluated zero-shot with pre-trained ControlNets under tiled and blur image conditions. For multi-GPU inference, the paper reports greater acceleration as resolution increases, while the original DiT encounters out-of-memory errors.

F. More ControlNet Results

  • CLEAR layers are tested zero-shot with pre-trained ControlNets under tiled and blur image conditions.The quantitative evaluation uses 1,000 COCO2014 validation images; RMSE is computed against condition images.Source: F. More ControlNet Results · p. 13, Table 9 · p. 14

G. Efficient Multi-GPU Parallel Inference

  • Multi-GPU acceleration becomes more significant as image resolution increases, while the original DiT encounters out-of-memory errors because it must cache all key-value tokens.Efficiency is measured in seconds per 50 denoising steps on an HGX H100 8-GPU server; the table notes that CLEAR with r = 16 at 1024 × 1024 is unavailable because each GPU’s patch is smaller than the boundary size.Source: G. Efficient Multi-GPU Parallel… · p. 13, Table 10 · p. 14
  • The appendix shows qualitative comparisons of original and CLEAR-linearized FLUX-1.dev and SD3.5-Large, plus additional 4K examples from CLEAR-linearized FLUX-1.dev.Prompts are listed in Fig. 16.Source: Figure 13 · p. 15, Figure 14 · p. 16

Related papers

Loading 2412.16112v1…