Exploring Diffusion with Test-Time Training on Efficient Image Restoration

2025/06/17 by Rongchang Lu, Lu, Rongchang, Luo, Tianduo +10
Engineering · #Computer Vision and Pattern Recognition (cs.CV) #FOS: Computer and information sciences #I.4.9 #Image Processing Techniques and Applications

paper · pdf · doi:10.48550/arxiv.2506.14541

openalex publication_date 2025/06/17 · openalex created_date 2025/10/13 · openalex updated_date 2026/07/28

Abstract

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.

Citations

• RWKV-7 "Goose" with Expressive Dynamic State Evolution
• AdvPaint: Protecting Images from Inpainting Manipulation via Adversarial Attention Disruption
• MambaIRv2: Attentive State Space Restoration
• Enhanced Control for Diffusion Bridge in Image Restoration
• Restore-RWKV: Efficient and Effective Medical Image Restoration with RWKV
• VisualRWKV: Exploring Recurrent Neural Networks for Visual Language Models
• Parallelizing Linear Transformers with the Delta Rule over Sequence Length
• Structure Matters: Tackling the Semantic Discrepancy in Diffusion Models for Image Inpainting
• MambaIR: A Simple Baseline for Image Restoration with State-Space Model
• Image Restoration Through Generalized Ornstein-Uhlenbeck Bridge
• HAT: Hybrid Attention Transformer for Image Restoration
• DiffIR: Efficient Diffusion Model for Image Restoration
• Zero-Shot Image Restoration Using Denoising Diffusion Null-Space Model
• Root-mean-square error (RMSE) or mean absolute error (MAE): when to use them or not
• High-Resolution Image Synthesis with Latent Diffusion Models
• High-Resolution Image Synthesis with Latent Diffusion Models
• Resolution-robust Large Mask Inpainting with Fourier Convolutions
• SwinIR: Image Restoration Using Swin Transformer
• SwinIR: Image Restoration Using Swin Transformer
• Aggregated Contextual Transformations for High-Resolution Image Inpainting
• High-Fidelity Pluralistic Image Completion with Transformers
• EdgeConnect: Generative Image Inpainting with Adversarial Edge Learning
• The Unreasonable Effectiveness of Deep Features as a Perceptual Metric
• Progressive Growing of GANs for Improved Quality, Stability, and Variation
• Places: An Image Database for Deep Scene Understanding
• Deep Residual Learning for Image Recognition
• A Feature-Enriched Completely Blind Image Quality Evaluator
• Image Super-Resolution Using Deep Convolutional Networks
• Image Super-Resolution Using Deep Convolutional Networks
• MatConvNet - Convolutional Neural Networks for MATLAB
• Image quality assessment: from error visibility to structural similarity
• On calculating with B-splines

Related

• Harnessing Diffusion-Yielded Score Priors for Image Restoration
• DiffIR2VR-Zero: Zero-Shot Video Restoration with Diffusion-based Image Restoration Models
• Efficient Diffusion Model for Image Restoration by Residual Shifting
• SeedVR2: One-Step Video Restoration via Diffusion Adversarial Post-Training
• Exploiting Diffusion Prior for Task-driven Image Restoration