Efficient One-Step Diffusion Restoration Model with Compact Token Compression and Linear Attention

arXiv:2605.23451 · cs.CV · Submitted 2026-05-22 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: We're looking at the paper: "Efficient One-Step Diffusion Restoration Model with Compact Token Compression and Linear Attention".

Jane: Here's the gist of the paper. The gist The SANA-SR framework proposes an efficient one-step restoration model that achieves strong quality and efficiency trade-offs for real-world image super-resolution by employing compact latent representation and linear complexity modeling How it works SANA-SR is designed to address the bottleneck of excessive token redundancy and costly token interactions in high-resolution restoration by revisiting the task from two perspectives: compact latent representation and efficient token interaction The framework first leverages a "degradation-aware deep compression autoencoder with a 32× compression ratio, which significantly reduces the number of latent tokens while preserving restoration-relevant structural and textural information" to map both the low-quality input and the high-quality target into a compact latent space, yielding zL and zH This results in fewer tokens than most existing diffusion-based restorers which operate on dense 4× or 8× VAE latents, leading to high per-step computation Once the restoration problem is transferred to this compact latent space, the model adapts a pretrained SANA LinearDiT prior with lightweight LoRA updates rather than training a new restoration transformer from scratch To ensure consistency with clean image generation rather than degraded input appearance, the model is conditioned on a "degradation-aware refined prompt" constructed from three components: (i) a tag-based extractor Π(·), that emits content tags from xL (e.g.

Tom: tree, foliage, leaf, branch); (ii) a fixed quality template ptpl = “clean, sharp, best quality, detailed, 8K, high-resolution” that biases the prior toward the clean regime; and (iii) an attention mask m that suppresses any residual degradation-related token The one-step LinearDiT backbone utilizes "linear self-attention" in each block, where queries Q, keys K, and values V are three learned linear projections of this latent token sequence so that LA(Q, K, V) = ϕ(Q)ϕ(K)⊤Vϕ(Q)ϕ(K)⊤1 + εatt The model predicts a latent residual and updates zL to the restored latent zˆ using the update equation zˆ = zL − στg fθ(zL, τg, c, m), where στg>0 is the scheduler coefficient at τg Training and Regularization The training objective is composed of three terms: pixel loss Lrec = λ2/3HW∥xˆ − xH∥2 + λp LPIPS(ˆx, xH), which combines pixel and perceptual losses The framework further regularizes one-step adaptation with frozen-prior alignment (Lalign) and adapter consistency (Lcons) The frozen-prior alignment branch evaluates the same frozen backbone on perturbed latents to align channel-wise Gaussian summaries, and the consistency branch compares the adapter-on and adapteroff responses on the same perturbed restored latent Only the LoRA parameters in fθ are optimized, while the VAE, text encoder, and pretrained SANA remain frozen throughout training Efficient Deployment via Pruning After LoRA adaptation, the model is further compressed through "prompt-aware structured pruning" to select a subset of transformer blocks under a deployment budget P⋆ The saliency of block Bl is defined as Sl = X wi∈Bl w 2i(Fi + εp), where Fi is the reciprocal-curvature-weighted score derived from the task loss This strategy achieves the best quality–efficiency tradeoff under a 0.35B budget, reducing latency and MACs substantially while preserving most of the restoration quality Experimental Performance SANA-SR demonstrates strong performance across various benchmarks, achieving "the best PSNR, SSIM, and MUSIQ on DIV2KVal" and ranking second on DISTS and NIQE On real-world data like RealSR, the method achieves "the best MANIQA and CLIPIQA," showing a "strong overall balance between fidelity and perceptual quality on authentic degradations" compared to recent one-step baselines In terms of efficiency, the deployed model runs in 0.019 s with 407.95 G MACs and 344M parameters, which is faster (0.019 s vs.

Jane: 0.03 s), lighter in computation (407.95 G vs.

Tom: 496 G MACs), and uses fewer parameters (344 M vs.

Paper summary: Tom: Looking at the title of "Efficient One-Step Diffusion Restoration Model with Compact Token Compression and Linear Attention," it really captures the core idea—efficiency through compression and linear attention in this new diffusion restoration model.

Tom: So, when we look at "Efficient One-Step Diffusion Restoration Model with Compact Token Compression and Linear Attention," the real point is that they managed to get high fidelity without relying on massive computational resources or overly complex math.

Lu: The authors show how this works by using a degradation-aware deep compression autoencoder with a thirty-two times compression ratio, which significantly reduces the number of latent tokens while preserving restoration-relevant structural and textural information <ref:2605.23451#pg1>.

Jane: It's interesting to hear that they are mapping both the low-quality input and the high-quality target into a compact latent space, yielding zL and zH <ref:2605.23451#pg2>.

Meng: I wonder if that thirty-two times compression ratio is sustainable across different types of real-world degradation; sometimes aggressive compression loses important details <ref:2605.23451#pg1>.

Jane: That combination of techniques—compression and linear attention—is what makes the difference in complexity compared to older quadratic methods <ref:2605.23451#pg1>.

Lalam: The paper shows that SANA-SR achieves this by mapping both the low-quality input and the high-quality target into a compact latent space, yielding zL and zH <ref:2605.23451#pg2>.

Conclusion: Tom: So, to wrap up this paper by SANA-SR, the authors are showing us how to tackle real-world image super-resolution without needing massive computational resources or overly complicated math in the process.

Jane: What they're really doing is moving away from those old quadratic models that use dense latent representations and instead focusing on compact latent representation and linear complexity modeling.

Lu: The paper argues that the bottleneck isn't a lack of restoration priors, but rather excessive token redundancy and costly token interactions during high-resolution restoration.

Meng: That means they are cutting down on the amount of data the AI has to process at every step, which is crucial when you're trying to deploy something that actually runs on a mobile device or a standard server.

Lalam: The method uses a deep compression autoencoder with a thirty-two times compression ratio to drastically reduce the number of latent tokens while still keeping all the important structural and textural information for restoration.

Tom: That thirty-two times compression ratio is what makes them drastically reduce the computation compared to methods that handle dense, high-resolution latents.

Jane: And on top of having this compact latent space, they introduce a linear-attention DiT with LoRA fine-tuning to handle the restoration efficiently without needing a full re-training process.

Lu: They condition this model on a degradation-aware refined prompt which includes content tags and a quality template that biases the prior toward clean, sharp outputs.

Meng: I'm interested in how they manage that linear attention mixing across all those compressed tokens to keep the complexity low while still getting high fidelity results.

Tom: The results show SANA-SR achieves highly competitive or even superior quantitative performance on benchmarks, like DIV2KVal and DRealSR, while also delivering much faster inference speeds.

Jane: It's interesting because they managed to get this strong balance between restoration quality and efficiency that’s hard to find in other one-step methods.

Lu: This suggests that the structure of how we represent and interact with tokens can be significantly changed to make diffusion restoration practical for real applications.

Tom: If this approach scales, it could mean much faster image enhancement pipelines for things like real-time video processing or mobile photo editing apps.

Bingtian Qiao, Yue Shi, Yingjie Zhou, Yong Guo, Guangtao Zhai, Jiezhang Cao

Shanghai Jiao Tong University · Fuzhou University · Shanghai AI Laboratory

cs.CV

Submitted: 2026-05-22

Updated: 2026-10-20

Comments: NeurIPS 2026; code: https://github.com/Seiun-Sky-debug/SANASR

Code: https://github.com/Seiun-Sky-debug/SANASR

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 82/100

The gist: The gist The SANA-SR framework proposes an efficient one-step restoration model that achieves strong quality and efficiency trade-offs for real-world image super-resolution by employing compact

Key concepts

Deep Compression Autoencoder: This component is used to compress both the low-quality input image and the high-quality target image into a very small, compact latent space. By using a 32x compression ratio, it significantly reduces the number of tokens needed compared to standard methods, making the subsequent restoration process much faster and less memory-intensive.
Linear Self-Attention: Instead of traditional quadratic self-attention which scales poorly with high resolution, SANA-SR employs linear self-attention. This technique uses three learned linear projections (Q, K, V) to perform attention calculations in a way that scales linearly with the number of tokens. This allows the model to process high-resolution data efficiently without incurring excessive computational overhead.
Prompt-Aware Refined Prompt: The model is conditioned not just on an image, but on a carefully constructed prompt. This prompt includes content tags from the input, a fixed quality template like 'clean, sharp, best quality,' and an attention mask to suppress noise. This guidance helps the linear DiT backbone focus its restoration effort precisely where it is needed for realistic texture recovery.
Prompt-Aware Structured Pruning: After training, the model undergoes structured pruning based on block saliency derived from task loss. This process selects a subset of transformer blocks under a specific budget to significantly reduce the model's size and computational requirements while preserving most of the restoration quality.

Terminology

Latent Representation — 2× in this paper · explained in 2 episodes
Benchmark — 1× in this paper · explained in 2 episodes
Fine-Tuning — 1× in this paper · explained in 5 episodes
Latent Representations — 1× in this paper · explained in 3 episodes
Latent Space — 1× in this paper · explained in 7 episodes
LoRA Fine-Tuning — 1× in this paper · explained in 3 episodes
Pruning — 1× in this paper · explained in 4 episodes
Redundancy — 1× in this paper · explained in 2 episodes

Summary

The gist The SANA-SR framework proposes an efficient one-step restoration model that achieves strong quality and efficiency trade-offs for real-world image super-resolution by employing compact latent representation and linear complexity modeling

How it works

SANA-SR is designed to address the bottleneck of excessive token redundancy and costly token interactions in high-resolution restoration by revisiting the task from two perspectives: compact latent representation and efficient token interaction The framework first leverages a degradation-aware deep compression autoencoder with a 32× compression ratio, which significantly reduces the number of latent tokens while preserving restoration-relevant structural and textural information to map both the low-quality input and the high-quality target into a compact latent space, yielding zL and zH This results in fewer tokens than most existing diffusion-based restorers which operate on dense 4× or 8× VAE latents, leading to high per-step computation

Once the restoration problem is transferred to this compact latent space, the model adapts a pretrained SANA LinearDiT prior with lightweight LoRA updates rather than training a new restoration transformer from scratch To ensure consistency with clean image generation rather than degraded input appearance, the model is conditioned on a degradation-aware refined prompt constructed from three components: (i) a tag-based extractor Π(·), that emits content tags from xL (e.g. tree, foliage, leaf, branch); (ii) a fixed quality template ptpl = “clean, sharp, best quality, detailed, 8K, high-resolution” that biases the prior toward the clean regime; and (iii) an attention mask m that suppresses any residual degradation-related token The one-step LinearDiT backbone utilizes linear self-attention in each block, where queries Q, keys K, and values V are three learned linear projections of this latent token sequence so that LA(Q, K, V) = ϕ(Q)ϕ(K)⊤Vϕ(Q)ϕ(K)⊤1 + εatt The model predicts a latent residual and updates zL to the restored latent zˆ using the update equation zˆ = zL − στg fθ(zL, τg, c, m), where στg>0 is the scheduler coefficient at τg

Training and Regularization

The training objective is composed of three terms: pixel loss Lrec = λ2/3HW∥xˆ − xH∥2 + λp LPIPS(ˆx, xH), which combines pixel and perceptual losses The framework further regularizes one-step adaptation with frozen-prior alignment (Lalign) and adapter consistency (Lcons) The frozen-prior alignment branch evaluates the same frozen backbone on perturbed latents to align channel-wise Gaussian summaries, and the consistency branch compares the adapter-on and adapteroff responses on the same perturbed restored latent Only the LoRA parameters in fθ are optimized, while the VAE, text encoder, and pretrained SANA remain frozen throughout training

Efficient Deployment via Pruning

After LoRA adaptation, the model is further compressed through prompt-aware structured pruning to select a subset of transformer blocks under a deployment budget P⋆ The saliency of block Bl is defined as Sl = X wi∈Bl w 2i(Fi + εp), where Fi is the reciprocal-curvature-weighted score derived from the task loss This strategy achieves the best quality–efficiency tradeoff under a 0.35B budget, reducing latency and MACs substantially while preserving most of the restoration quality

Experimental Performance

SANA-SR demonstrates strong performance across various benchmarks, achieving the best PSNR, SSIM, and MUSIQ on DIV2KVal and ranking second on DISTS and NIQE On real-world data like RealSR, the method achieves the best MANIQA and CLIPIQA, showing a strong overall balance between fidelity and perceptual quality on authentic degradations compared to recent one-step baselines In terms of efficiency, the deployed model runs in 0.019 s with 407.95 G MACs and 344M parameters, which is faster (0.019 s vs. 0.03 s), lighter in computation (407.95 G vs. 496 G MACs), and uses fewer parameters (344 M vs. 456 M) compared to the strongest recent one-step baseline AdcSR

Ablation Insights

Ablation studies confirm that frozen-prior alignment is the most critical component in training, as removing Lalign causes the largest drop in PSNR, LPIPS, and MUSIQ Similarly, prompt design has a consistent impact on both fidelity and perceptual quality, with DAPE giving the best overall result The choice of generation timestep τg = 900 is chosen as the best balance between fidelity and perception

Conclusion

SANA-SR achieves a favorable trade-off between restoration quality and efficiency, delivering highly competitive or superior quantitative results while running in 0.019s with 407.95G MACs and 344M parameters across synthetic and real-image benchmarks

--- Page 18 ---

SANA-SR (ours) TSD-SR SwinIR StableSR LR SeeSR ResShift OSEDiff AdcSR SANA-SR (ours) TSD-SR SwinIR StableSR LR SeeSR ResShift OSEDiff AdcSR Figure 5: Additional qualitative comparison on challenging examples from DRealSR. PSNR SSIM MUSIQCLIPIQA LQ HQ HQ+S HQ+S+D 24.00 0.716 68.69 741 24.32 0.729 68.74 748 HQ+S+D SANA-SR Ours – TSD-SR SwinIR StableSR LR SeeSR ResShift OSEDiff AdcSR Figure 3: Qualitative comparison on challenging examples from DRealSR PSNR SSIM MUSIQCLIPIQA LQ HQ HQ+S HQ+S+D 25.96 0.747 70.73 766

--- Page 19 ---

Table 4: Quantitative comparison on DRealSR with best and second-best results Prior-based multi-step / flexible-step methods StableSR DiffBIR SeeSR PASD ResShift SUPIR DreamClear InvSR LinearSR Efficient / one-step methods SinSR OSEDiff S3Diff AddSR D3SR FLUX-SR AdcSR TSD-SR Ours – 27.99 0.8001 0.6073 7179

--- Page 20 ---

Table 5: Efficiency comparison of different kinds of methods Prior-based multi-step / flexible-step One-step methods Metric StableSR DiffBIR SeeSR PASD ResShift SinSR OSEDiff S3Diff AdcSR Ours – 200 50 50 20 15 1 1 1

--- Page 21 ---

Table 6: Comparisons of different pruning strategies Strategy Params Time(s) Full model full full model full Tail pruning Full model full Tail pruning Random pruning Full model full Random pruning Unstructured pruning ∼0.35B Full model full Unstructured pruning ∼0.35B Ours 0.344B

--- Page 22 ---

Table 7: Quantitative comparison on DRealSR with best and second-best results Prior-based multi-step / flexible-step methods StableSR DiffBIR SeeSR PASD ResShift SinSR OSEDiff S3Diff AddSR D3SR FLUX-SR AdcSR TSD-SR Ours – 27.99 0.8001 0.

Improvements for AI systems

The system can achieve high-quality image super-resolution (Real-ISR) with a strong efficiency trade-off by employing a degradation-aware deep compression autoencoder with a 32× compression ratio to drastically reduce latent tokens while preserving restoration-relevant structural and textural information.
The architecture can utilize linear self-attention within the DiT backbone, which reduces per-image complexity from quadratic to linear, leading to a significant reduction in MACs and inference time, as evidenced by the report that SANA-SR achieves an asymptotic per-image cost of O(N d 2) with N ≈256.
The model can be adapted efficiently using LoRA fine-tuning instead of training from scratch, allowing for effective restoration learning with moderate training overhead, and further stabilized by a loss term that ensures compatibility between the predicted latent and the frozen prior: Lalign = 1/2BC X B b=1 X C c=1 log sb,c(qH) + ε / sb,c(qzˆ) + µb,c(qzˆ) − µb,c(qH).
The deployed model can be optimized for practical mobile deployment through prompt-aware structured pruning, which selects blocks based on a reciprocal-curvature-weighted score to ensure that the retained blocks preserve prompt-conditioned restoration behavior.