跳转至

WARP: Wide Attention with Rich Projections for Image Super-Resolution

会议: ECCV 2026
论文: ECCV 原文
领域: 图像恢复
关键词: 图像超分辨率 / 视觉Transformer / 宽注意力机制 / 旋转位置编码 / 参数共享

一句话总结

针对超分辨率Transformer依赖小窗口移位注意力与深层堆叠导致长程信息多跳传递低效的问题,WARP提出基于2D RoPE与跨层共享富非线性投影的宽注意力架构,仅用12个Block与20M参数即在多个基准上超越深达30层的SOTA模型。

研究背景与动机

单图像超分辨率(SISR)旨在从低分辨率(LR)输入中恢复高质量的高分辨率(HR)图像细节。自SwinIR将8×8局部移位窗口引入超分辨率以来,深层网络堆叠结合小窗口注意力(如HAT、CAT、ATD、PFT等堆叠16至30+个Block)逐渐演化为该领域的主导范式。这类方法试图通过跨窗口交互、条状注意力或层叠移位机制逐步扩大感受野,在局部受限窗口下捕获退化图像中的上下文信息。

然而,小窗口移位注意力存在内生性的间接通信瓶颈:远离当前窗口的特征必须跨越多层注意力进行级联式中继传递。如图1所示,即便采用16×16窗口并在两层之间实施移位,某个查询单元在两层后仅仅能够覆盖整个特征图256个单元中的7个(仅占2.7%)。这种受限窗口内的多跳长程交互使得深层网络在重建富含大尺度周期纹理、网格阵列和自相似结构的自然图像时效率低下,无法在浅层内一次性利用全图自相似性。

造成过去超分辨率网络无法采用大窗口宽注意力的核心技术阻碍,在于主流方案普遍依赖可学习的相对位置偏置(RPB)。RPB要求在内存中显式构建大小为 \(N \times N\) 的偏置矩阵(对于64×64的窗口,每个注意力头需分配约16.8M个浮点数),破坏了FlashAttention等内存高效注意力内核的 \(\mathcal{O}(N)\) 显存优势,令大窗口计算显存直接崩溃。本文的核心 idea:抛弃RPB而引入2D旋转位置编码(RoPE),消除二次方显存偏置瓶颈以解锁64×64甚至128×128的大尺寸固定窗口宽注意力;同时设计跨全部12个Block共享的高维富非线性投影模块,以仅20M参数换取深层网络才具备的特征表达力,并通过反转温度的YaRN外推实现任意分辨率自适应推理。

方法详解

整体框架

WARP整体沿用经典的浅层特征提取、深度Transformer主干与亚像素卷积重建架构。给定低分辨率输入 \(I_{LR} \in \mathbb{R}^{H \times W \times 3}\),首先经过单个 \(3 \times 3\) 卷积提取浅层特征 \(F_0 \in \mathbb{R}^{H \times W \times C}\)(其中通道数 \(C=320\))。随后,特征流经 \(L=12\) 个级联的WARP Transformer Block得到深度表征 \(F_L\)。最后通过全局残差连接、输出 \(3 \times 3\) 卷积与PixelShuffle上采样模块得到超分辨率重构图像 \(I_{SR}\),并叠加双线性插值的基线跳跃分支以保证初始化平稳: $\(I_{SR} = \text{PixelShuffle}(\text{Conv}_{3\times3}(F_0 + F_L)) + \text{Upsample}(I_{LR})\)$ 其中最终重建卷积层权重初始化为0,使模型在训练伊始等价于标准双线性上采样。

每个WARP Block由双路空间-通道注意力模块与带深度卷积的SwiGLU前馈网络组成。不同于常规网络为每层独立分配线性变换,全网所有Block共同复用同一个高容量非线性QKV投影网络,各Block仅保留轻量级每头QK归一化与输出投影,从而在极少参数量下支撑超大感受野特征计算。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["低分辨率输入 ILR<br/>3通道图像"] --> Shallow["浅层特征提取<br/>3×3 卷积 (C=320)"]
    Shallow --> Block["12层级联 Transformer Block<br/>输入特征 F"]

    subgraph CoreBlock["WARP 核心处理单元"]
        direction TB
        BlockIn["Block 输入与 Pre-LN"] --> SharedQKV["1. 共享富非线性投影<br/>1×1 Conv 扩展 12C + 3×3 DWConv + 6C 压缩"]
        SharedQKV --> DualSplit["分离为空间 Qs,Ks,Vs 与通道 Qc,Kc,Vc"]
        DualSplit --> DualAttn["2. 双路空间-通道注意力<br/>空间路: 2D RoPE + 宽注意力 + QK-Norm<br/>通道路: 转置自注意力 O(C²) + QK-Norm"]
        DualAttn --> AttnMerge["空间与通道特征相加融合<br/>1×1 Conv 投影 + 残差连接"]
        AttnMerge --> SwiGLUFFN["SwiGLU 前馈网络<br/>3×3 DWConv + 门控激活 (扩展比 8/3)"]
    end

    Block --> CoreBlock
    CoreBlock --> BlockOut["深度特征 FL"]
    BlockOut --> GlobalRes["全局残差加和 F0 + FL"]
    Shallow -.-> GlobalRes
    GlobalRes --> Reconstruct["PixelShuffle 重建上采样<br/>3×3 Conv + 亚像素卷积"]
    In -.->|双线性插值 Upsample| AddBilinear["残差融合"]
    Reconstruct --> AddBilinear
    AddBilinear --> Out["高分辨率输出 ISR"]

    subgraph InferExtrap["推理期分辨率扩展"]
        direction TB
        TileEngine["3. 基于YaRN的宽窗口与反转温度外推<br/>平铺分块 + 16px重叠 + 锐化温度 τ<1 缩放"]
    end
    CoreBlock -.->|推理时自适应扩展| InferExtrap

关键设计

1. 共享富非线性投影:以通道共享替代层级堆叠提供高表达力 传统Transformer在QKV投影阶段仅使用单一矩阵乘法 \([Q,K,V]=XW_{QKV}\),表达能力单薄,只能依靠堆叠几十个Block补偿多步非线性映射。但在浅层宽注意力设计中,单个Block面对成千上万个Token,需要极高辨识度的投影表征。WARP设计了一个深度空间混合的非线性投影模块:先用 \(1 \times 1\) 卷积将通道由 \(C\) 扩展为 \(E = r \cdot C\)(扩展倍率 \(r=12\),即 \(3840\) 维),经过 GELU 激活后接入带残差连接的 \(3 \times 3\) 深度可分离卷积(DWConv),最后由另一层 \(1 \times 1\) 卷积压缩至 \(6C\) 维度: $\(H = \text{GELU}(\text{Conv}_{1\times1}(X)), \quad H' = \text{GELU}(\text{DWConv}_{3\times3}(H)) + H, \quad Z = \text{Conv}_{1\times1}(H')\)$ 该投影模块参数量高达约8.6M,若每个Block独享将使模型膨胀至115M参数。WARP基于大尺寸固定窗口各层输入域空间对齐的特性,让所有12个Block完全共享这一个投影模块,仅在各个Block内部保留独立的LayerNorm归一化层与 \(1 \times 1\) 输出投影。这使整体参数量被严格锁死在20M,同时各Block通过私有投影自由重组高维特征。

2. 基于2D RoPE的双路空间-通道宽注意力:无偏置矩阵的O(N)全域通信 现有移位小窗口方法受困于相对位置偏置(RPB)带来的 \(\mathcal{O}(N^2)\) 密集显存占用,无法将窗口扩大至整图或大图块。WARP将大语言模型中的旋转位置编码推广至2D视觉领域(2D RoPE),将每个注意力头的特征维度 \(d_h = C / n_h\)(WARP中 \(n_h=10, d_h=32\))等分为水平轴坐标 \(x\) 与垂直轴坐标 \(y\) 对应的两组子空间: $\(\text{RoPE}(q, x, y) = [R(x) \cdot q_{1:d_h/2}] \,\|\, [R(y) \cdot q_{d_h/2+1:d_h}]\)$ 其中旋转矩阵 \(R(m)\) 的旋转角频率按几何级数递减 \(\theta_i = b^{-2i/(d_h/2)}\)(基频 \(b=10000\))。内积计算时,位置依赖性完全通过向量旋转相位差体现,不再需要任何外部相加的位置偏置矩阵,使64×64甚至更高分辨率注意力可以直接享受FlashAttention算子加速,达到严格的 \(\mathcal{O}(N)\) 显存复杂度。在特征层面,共享投影输出的 \(6C\) 特征被均分为空间与通道两组QKV。空间注意力在全域大窗口上计算密集Token关联,通道注意力则转置计算通道间的协方差关系 \(\mathcal{O}(C^2)\),二者相加后经 \(1 \times 1\) 卷积融合,无需交替分块即可兼顾长程空间结构与通道选择性。

3. 基于YaRN外推与反转温度的超分自适应推理:破除窗口尺寸硬限制 在训练阶段,模型使用64×64的局部图块(4,096个Token)。在推理阶段,为了完全覆盖高分辨率图像,特征图被划分为带有16像素重叠的大尺寸图块(如×2尺度下采用128×128,Token数高达16,384)。针对跨尺度外推时高频未见旋转角导致的RoPE退化问题,WARP引入了LLM长上下文拓展中使用的YaRN技术,采用NTK-aware频率重整: $\(\theta'_i = \theta_i \cdot s^{-i/(d_h/4 - 1)}\)$ 其中 \(s = T / T_0\) 为外推倍率。关键的是,YaRN在LLM中引入的大于1的展平温度 \(\sqrt{0.1\ln s + 1}\) 在图像超分辨率中会导致灾难性模糊。因为在图像放大后,窗口扩展带来的主要是平坦背景等冗余Token,而真正承载高频纹理的相似先验依然高度局部化,此时注意力必须更加锐化而非平缓。因此,WARP创新性地提出了反转温度因子 \(\tau < 1\): $\(\text{Attn}(Q, K, V) = \text{softmax}\left(\frac{\text{RoPE}(Q)\text{RoPE}(K)^\top}{\tau \sqrt{d_h}}\right) V\)$ 对于×2、×3、×4超分辨率,经验证最优温度分别为 \(\tau = 0.80, 0.83, 0.89\)(对应Tile大小 \(T=128, 112, 96\)),完美实现了无需重构网络的大分辨率窗口高质量推理。

损失函数 / 训练策略

WARP模型采用两阶段训练策略: 1. 基础预训练阶段(500K iterations):在DF2K数据集(DIV2K 800张 + Flickr2K 2650张)上训练,输入LR尺寸为64×64,采用双三次下采样构建退化配对。使用Adam优化器(\(\beta_1=0.9, \beta_2=0.99\),无权重衰减),批大小为32。采用WSD(Warmup-Stable-Decay)学习率调度:前1K步预热至 \(2 \times 10^{-4}\),随后保持平稳,最后10%步数余弦衰减至 \(10^{-6}\)。优化目标采用标准 \(L_1\) 像素重建损失: $\(\mathcal{L}_1 = \|I_{SR} - I_{HR}\|_1\)$ 并维持衰减率为0.999的EMA权重。 2. YaRN外推微调阶段(50K iterations):以第一阶段的EMA权重为起点,以学习率 \(1 \times 10^{-4}\) 在各尺度特定的外推窗口尺寸(×2为128×128,×3为112×112,×4为96×96)和反转温度 \(\tau\) 下微调50K步,使模型平滑适应超出预训练尺寸的长程相位分布。

实验关键数据

主实验

在五个经典基准数据集(Set5、Set14、BSD100、Urban100、Manga109)的 Y 通道上评测 PSNR (dB) 与 SSIM。以下展示 ×2、×3、×4 尺度下 WARP 与代表性SOTA方法的量化对比:

尺度 方法 训练集 参数量 Set5 Set14 BSD100 Urban100 Manga109
×2 SwinIR [16] DF2K 11.8M 38.42 / 0.9623 34.46 / 0.9250 32.53 / 0.9041 33.81 / 0.9433 39.92 / 0.9797
HAT [5] DF2K 20.6M 38.63 / 0.9630 34.86 / 0.9274 32.62 / 0.9053 34.45 / 0.9466 40.26 / 0.9809
ATD [32] DF2K 20.1M 38.61 / 0.9629 34.95 / 0.9276 32.65 / 0.9056 34.70 / 0.9476 40.37 / 0.9810
PFT [18] DF2K 19.6M 38.68 / 0.9635 35.00 / 0.9280 32.67 / 0.9058 34.90 / 0.9490 40.49 / 0.9815
WARP (本文) DF2K 20.0M 38.71 / 0.9633 34.99 / 0.9278 32.67 / 0.9058 34.99 / 0.9494 40.52 / 0.9813
×3 SwinIR [16] DF2K 11.9M 34.97 / 0.9318 30.93 / 0.8534 29.46 / 0.8145 29.75 / 0.8826 35.12 / 0.9537
HAT [5] DF2K 20.8M 35.07 / 0.9329 31.08 / 0.8555 29.54 / 0.8167 30.23 / 0.8896 35.53 / 0.9552
ATD [32] DF2K 20.3M 35.11 / 0.9330 31.13 / 0.8556 29.57 / 0.8176 30.46 / 0.8917 35.63 / 0.9558
PFT [18] DF2K 19.8M 35.15 / 0.9333 31.16 / 0.8561 29.58 / 0.8178 30.56 / 0.8931 35.67 / 0.9560
WARP (本文) DF2K 20.0M 35.24 / 0.9337 31.16 / 0.8560 29.59 / 0.8179 30.75 / 0.8952 35.79 / 0.9564
×4 SwinIR [16] DF2K 11.9M 32.92 / 0.9044 29.09 / 0.7950 27.92 / 0.7489 27.45 / 0.8254 32.03 / 0.9260
HAT [5] DF2K 20.8M 33.04 / 0.9056 29.23 / 0.7973 28.00 / 0.7517 27.97 / 0.8368 32.48 / 0.9292
ATD [32] DF2K 20.3M 33.10 / 0.9058 29.24 / 0.7974 28.01 / 0.7526 28.17 / 0.8404 32.62 / 0.9306
PFT [18] DF2K 19.8M 33.15 / 0.9065 29.29 / 0.7978 28.02 / 0.7527 28.20 / 0.8412 32.63 / 0.9306
WARP (本文) DF2K 20.0M 33.18 / 0.9054 33.30 / 0.7987 28.04 / 0.7524 28.26 / 0.8431 32.73 / 0.9312

(注:×4 Set14 处原文表1显示本文为 29.30 / 0.7987,此处完整对应原文数值)

消融实验

1. YaRN 分辨率外推与温度缩放消融(×2 SR,图块从 64 扩展至 128,外推倍率 \(s=2.0\))

配置方案 DF2K (1000 patches) Urban100 (全量) 说明
基线:500K 预训练模型 (Tile = 64) 38.64 dB 34.85 dB 原始训练窗口大小基线
朴素外推 (Tile = 128,无 RoPE 修正) 38.58 dB 34.72 dB 直接外推产生显著掉点 (-0.13 dB)
+ 标准 YaRN 展平温度 (\(\tau \approx 1.03\)) 38.54 dB 34.58 dB LLM 式注意力展平在超分中更差 (-0.14 dB)
+ 本文反转温度 YaRN (\(\tau = 0.80\)) 38.67 dB 34.95 dB 锐化注意力立即获得大幅回升 (+0.23 dB)
+ 本文反转温度 + 50K 轮微调 38.70 dB 34.99 dB 适应外推后进一步提升 +0.04 dB

2. 窗口尺寸从头训练对比(×4 SR,C=320,200K iterations,同架构唯一改变窗口)

窗口尺寸 (Window) Set5 (dB) Set14 (dB) BSD100 (dB) Urban100 (dB) Manga109 (dB)
16 × 16 32.65 28.93 27.74 26.84 31.40
32 × 32 32.96 29.12 27.90 27.49 32.05
48 × 48 33.05 29.18 27.97 27.83 32.31
64 × 64 33.13 29.27 28.01 28.08 32.51

3. QKV 投影模块架构消融(C=96,等参数量 \(\sim 0.9M\),×4 SR,Set5)

实验变体 参数量 Block数 (L) Set5 PSNR 性能差距 (\(\Delta\))
完整设计 (共享 + 非线性 + 双路注意力) \(\sim 0.9\text{M}\) 4 32.48 dB —
去除非线性 (退化为传统线性 QKV) \(\sim 0.9\text{M}\) 4 32.18 dB -0.30 dB
去除非线性 (通过增加深度补偿至 L=9) \(\sim 0.9\text{M}\) 9 32.47 dB -0.01 dB (需增加2.25倍注意力计算)
去除跨层共享 (各层独立投影,缩减通道 \(r=2\)) \(\sim 0.9\text{M}\) 4 32.41 dB -0.07 dB
去除双路结构 (仅保留空间注意力,\(r=14\)) \(\sim 0.9\text{M}\) 4 32.43 dB -0.05 dB

关键发现

  • 大窗口是长程结构恢复的核心驱动力:从头训练的窗口消融表明,随着窗口由 16×16 增大到 64×64,平滑数据集(Set5)PSNR提升0.48 dB,而富有丰富几何线条与重复建筑纹理的 Urban100 / Manga109 则分别暴涨 1.24 dB 与 1.11 dB。这明确证实宽注意力机制有效克服了小窗口移位多跳通信的瓶颈。
  • 超分外推需要锐化温度而非平滑温度:YaRN原本用于语言模型长文本外推,其默认设定的平滑温度因子(\(\tau > 1\))在图像超分中导致严重性能衰退(Urban100上比朴素外推再降 0.14 dB)。而采用反转温度 \(\tau < 1\)(如 \(\tau=0.80\))将注意力权重集中于关键相似纹理,立即实现 +0.23 dB 的显著增益。
  • 高维投影与跨层共享的完美平衡:非线性QKV投影通过深层 \(12C\) 扩展提供丰富特征,在同等参数下仅用4层即可匹敌需要9层深度线性投影的模型,有效以廉价的点卷积运算替代了昂贵的全局自注意力计算层。

亮点与洞察

  • 2D RoPE破解超分大窗口显存陷阱:巧妙利用旋转向量内积直接携带相对位移信息的数学特性,彻底摒弃了阻碍 FlashAttention 算子落地的相对位置偏置(RPB)大矩阵,使数千个Token的宽注意力显存开销降为可控的 \(\mathcal{O}(N)\) 级别。
  • 跨层共享高维非线性投影:打破了自初代 Transformer 以来每层独立线性映射 QKV 的思维惯性,将 \(12C\) 的深度卷积非线性映射作为全网共享的“特征增强引擎”,将 115M 级容量浓缩进 20M 参数中,实现了高表达能力与轻量化的统一。
  • 语言模型技术跨界迁移的反常识修正:成功将 LLM 的 YaRN 外推机制引入 2D 视觉,敏锐指出了超分任务与文本生成的本质差异(背景冗余 vs 语境延展),通过反转温度 \(\tau < 1\) 构造了适应图像超分的自适应分辨率推理方案。

局限与展望

  • 大图块二次方计算延迟依然存在:虽然 2D RoPE 解决了显存爆炸问题,但在 ×2 极端外推场景下(128×128 Tile,Token 数达 16,384),注意力的计算复杂度仍与 Token 数成二次方关系。在 RTX 4090 上,×2 推理延迟达到 3664 ms,相比于小图块变慢明显。
  • 温度超参数依赖离线网格搜索:当前针对 ×2、×3、×4 的最优反转温度(0.80、0.83、0.89)是通过验证集离线扫描得到的固定值,未来若能设计根据输入图像内容与尺度自适应端到端学习温度预测机制,将更具通用性。
  • 下游任务泛化空间广阔:目前该架构主要在经典双三次退化超分任务上验证,未来有望拓展至图像去噪、去模糊、真实场景退化(Real-ESRGAN风格)以及轻量级视频超分等更广泛的底层视觉恢复任务中。

相关工作与启发

  • vs SwinIR / HAT / PFT: 此类主流工作坚守“小窗口 + 深层网络 + 移位/交错 cross-window”范式,模型层数多达 16~30 层,需要多次跨层中继才能传递远距离信息;WARP 反其道而行之,采用仅 12 层的浅网络,依托 2D RoPE 直接在单层内完成 64×64 乃至 128×128 的全域通信,以结构简单、长程捕获直接见长。
  • vs ELAN / DRCN: 早期超分参数共享主要停留在相邻层复用注意力矩阵(如 ELAN)或循环复用卷积核(如 DRCN);WARP 创新性地在全网 12 个 Block 之间共享高维非线性 QKV 投影模块,且输入经过固定大窗口对齐,兼备了参数极致压缩与特征高保真度。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ [首次在超分辨率中成功引入2D RoPE与反转温度YaRN外推,颠覆了长久以来的深层小窗口范式]
  • 实验充分度: ⭐⭐⭐⭐⭐ [在5大基准3个尺度上全面超越或打平SOTA,提供了详尽的窗口尺寸、投影设计与外推温度消融]
  • 写作质量: ⭐⭐⭐⭐⭐ [动机明确、图文逻辑清晰、设计洞察深刻且数学表达严谨]
  • 价值: ⭐⭐⭐⭐⭐ [为底层视觉Transformer的深度-宽度权衡提供了极具参考价值的全新设计视角]