# Exploring Diffusion with Test-Time Training on Efficient Image Restoration

arXiv: https://arxiv.org/abs/2506.14541 · HTML: https://mlsanity.com/paper/2506.14541/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 26 of its 40 references that are in the map.
- First posted on arXiv 17 Jun 2025.

## Builds on (most cited first, 26 of 26)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [Progressive Growing of GANs for Improved Quality, Stability, and Variation](https://mlsanity.com/paper/1710.10196/) (2017)
- [SwinIR: Image Restoration Using Swin Transformer](https://mlsanity.com/paper/2108.10257/) (2021)
- [Image Super-Resolution Using Deep Convolutional Networks](https://mlsanity.com/paper/1501.00092/) (2014)
- [Resolution-robust Large Mask Inpainting with Fourier Convolutions](https://mlsanity.com/paper/2109.07161/) (2021)
- [Zero-Shot Image Restoration Using Denoising Diffusion Null-Space Model](https://mlsanity.com/paper/2212.00490/) (2022)
- [Parallelizing Linear Transformers with the Delta Rule over Sequence Length](https://mlsanity.com/paper/2406.06484/) (2024)
- [MambaIR: A Simple Baseline for Image Restoration with State-Space Model](https://mlsanity.com/paper/2402.15648/) (2024)
- [Sketch-based Manga Retrieval using Manga109 Dataset](https://mlsanity.com/paper/1510.04389/) (2015)
- [DiffIR: Efficient Diffusion Model for Image Restoration](https://mlsanity.com/paper/2303.09472/) (2023)
- [EdgeConnect: Generative Image Inpainting with Adversarial Edge Learning](https://mlsanity.com/paper/1901.00212/) (2019)
- [RWKV-7 "Goose" with Expressive Dynamic State Evolution](https://mlsanity.com/paper/2503.14456/) (2025)
- [High-Fidelity Pluralistic Image Completion with Transformers](https://mlsanity.com/paper/2103.14031/) (2021)
- [MambaIRv2: Attentive State Space Restoration](https://mlsanity.com/paper/2411.15269/) (2024)
- [Aggregated Contextual Transformations for High-Resolution Image Inpainting](https://mlsanity.com/paper/2104.01431/) (2021)
- [Places: An Image Database for Deep Scene Understanding](https://mlsanity.com/paper/1610.02055/) (2016)
- [HAT: Hybrid Attention Transformer for Image Restoration](https://mlsanity.com/paper/2309.05239/) (2023)
- [Image Restoration Through Generalized Ornstein-Uhlenbeck Bridge](https://mlsanity.com/paper/2312.10299/) (2023)
- [Structure Matters: Tackling the Semantic Discrepancy in Diffusion Models for Image Inpainting](https://mlsanity.com/paper/2403.19898/) (2024)
- [MatConvNet - Convolutional Neural Networks for MATLAB](https://mlsanity.com/paper/1412.4564/) (2014)
- [Restore-RWKV: Efficient and Effective Medical Image Restoration with RWKV](https://mlsanity.com/paper/2407.11087/) (2024)
- [AdvPaint: Protecting Images from Inpainting Manipulation via Adversarial Attention Disruption](https://mlsanity.com/paper/2503.10081/) (2025)
- [VisualRWKV: Exploring Recurrent Neural Networks for Visual Language Models](https://mlsanity.com/paper/2406.13362/) (2024)
- [Enhanced Control for Diffusion Bridge in Image Restoration](https://mlsanity.com/paper/2408.16303/) (2024)

## Abstract (arXiv)

Image restoration faces challenges including ineffective feature fusion, computational bottlenecks and inefficient diffusion processes. To address these, we propose DiffRWKVIR, a novel framework unifying Test-Time Training (TTT) with efficient diffusion. Our approach introduces three key innovations: (1) Omni-Scale 2D State Evolution extends RWKV's location-dependent parameterization to hierarchical multi-directional 2D scanning, enabling global contextual awareness with linear complexity O(L); (2) Chunk-Optimized Flash Processing accelerates intra-chunk parallelism by 3.2x via contiguous chunk processing (O(LCd) complexity), reducing sequential dependencies and computational overhead; (3) Prior-Guided Efficient Diffusion extracts a compact Image Prior Representation (IPR) in only 5-20 steps, proving 45% faster training/inference than DiffIR while solving computational inefficiency in denoising. Evaluated across super-resolution and inpainting benchmarks (Set5, Set14, BSD100, Urban100, Places365), DiffRWKVIR outperforms SwinIR, HAT, and MambaIR/v2 in PSNR, SSIM, LPIPS, and efficiency metrics. Our method establishes a new paradigm for adaptive, high-efficiency image restoration with optimized hardware utilization.