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.
Abstract (translated)
图像恢复面临包括无效特征融合、计算瓶颈和低效扩散过程等挑战。为解决这些问题，我们提出了一种新的框架——DiffRWKVIR，它将测试时间训练（TTT）与高效的扩散过程结合起来。我们的方法引入了三个关键创新： 1. **Omni-Scale 2D State Evolution**：扩展了RWKV的位置依赖参数化到层次化的多方向二维扫描，使得全局上下文感知具有线性复杂度O(L)。 2. **Chunk-Optimized Flash Processing**：通过连续的块处理加速了跨块并行性（复杂度为O(LCd)），减少了顺序依赖和计算开销，并提升了3.2倍的速度。 3. **Prior-Guided Efficient Diffusion**：在仅需5到20步的情况下提取紧凑的图像先验表示（IPR），相比DiffIR，训练和推理速度快45%，同时解决了去噪中的计算效率问题。 我们对超分辨率和补全基准（Set5、Set14、BSD100、Urban100、Places365）进行了评估。结果表明，DiffRWKVIR在PSNR、SSIM、LPIPS以及效率指标上均优于SwinIR、HAT和MambaIR/v2。 我们的方法为自适应且高效的图像恢复建立了新的范例，并实现了优化的硬件利用率。
URL
https://arxiv.org/abs/2506.14541
