论文阅读：CosAE Learnable Fourier Series for Image Restoration

这是 2024 NeurIPS 上发表的一篇文章，介绍了一种新型的基于傅里叶级数的通用编码器。

Abstract

本文介绍了余弦自动编码器（Cosine Autoencoder, CosAE），这是一种新颖的通用自动编码器，它将经典傅里叶级数与前馈神经网络无缝结合。CosAE 将输入图像表示为一系列二维余弦时间序列，每个序列由一组可学习的频率和傅里叶系数构成的元组定义。该方法与传统自动编码器形成鲜明对比，传统自动编码器在降低分辨率的瓶颈潜在空间中往往会牺牲细节。而 CosAE 在其瓶颈层对频率系数（即幅度和相位）进行编码。这种编码方式能够实现极高的空间压缩，例如，在瓶颈层将特征图下采样 64 倍，且解码时不会丢失细节。我们通过在灵活分辨率超分辨率和盲图像复原这两项极富挑战性的任务上进行大量实验，展示了 CosAE 的优势。这两项任务要求复原网络能够有效地泛化到复杂甚至未知的图像退化情况。我们的方法超越了当前最先进的技术，凸显了其学习可泛化图像复原表示的能力。

Introduction

在这里插入图片描述

图 1：与（a）传统自动编码器不同，（b）CosAE 将频域系数编码为极其狭窄的瓶颈特征图。通过一个基构造模块，它能够如实地呈现输入图像的全局结构和精细细节。

通过重建来训练图像自动编码器是图像表示学习中最常采用的方法之一。这种学习过程的核心是创建一个 information bottleneck （信息瓶颈）：网络首先将输入图像下采样为低空间维度的特征，然后再将它们上采样以重建输入图像。information bottleneck 的主要动机是增强网络捕获和保留关键数据内在模式和结构的能力，从而提高表示和泛化能力，这在数据质量对性能至关重要的视觉任务中尤为重要。在自动编码器中引入狭窄的瓶颈不仅为识别任务提供了一种学习通用表示的有效方法，而且为图像合成和操作提供了一种学习解纠缠表示的有效方法。

为了建立信息瓶颈，大多数现有的自动编码器网络将输入图像压缩到一个空间紧凑的潜在空间。它们通过提取主要的共享结构并筛选出任何噪声成分，来学习用于分类和属性解纠缠的中高级表示。然而，窄瓶颈自动编码器很少被直接设计用于图像复原任务，主要是因为传统自动编码器中使用下采样瓶颈常常会导致空间细节的丢失。因此，与更粗粒度的结构相比，这些自动编码器难以有效地表示高频细节，例如纹理。在像素回归质量起关键作用的其他应用中，例如 VQVAE 或潜在扩散模型 (LDM) ，研究人员通过在瓶颈的空间紧凑性和其保存细节的能力之间取得平衡来应对这一挑战。

在本文中，我们提出了一种新颖的自动编码器，可直接应用于图像恢复任务。它不仅拥有极其狭窄的瓶颈以提升表征能力，还能如实地保留高保真细节。我们的解决方案 —— 余弦自动编码器（CosAE），从傅里叶变换原理中汲取灵感，该原理表明有限信号可用一组谐波基函数来描述。我们遵循这一原理对频域系数（包括幅度和相位）进行编码，并在将其解码回原始图像之前，将它们与一组余弦基函数相结合（图 1（b））。这是通过一个谐波构建模块（HCM）实现的，该模块依据经典傅里叶级数的公式，将紧凑特征在空间上扩展为一系列谐波函数。与传统自动编码器的下采样特征图不同，我们发现在 CosAE 中，具有频率系数的表征在空间上可被极度压缩。然而，当这些表征与高频基函数（例如图 1 第二行中的余弦基函数）相结合时，就能有效地重建精细细节。

Approach

在这里插入图片描述

图 2：CosAE 包含一个编码器，该编码器将输入图像压缩为一个狭窄的瓶颈。根据公式（2）（见框图），每个瓶颈向量 $[A_n, \phi_{n}]$ 被转换为一组大小为 $\times T$ 的可学习谐波函数 $H$ 。然后，解码器从学习到的谐波函数中重建输入。为简便起见， $A_n$ 等同于公式（2）中的 $A_{(u_n, v_n)}$ 。

与传统自动编码器类似，CosAE 有一个编码器用于生成瓶颈潜在表示，还有一个解码器用于重建输入图像。在瓶颈层中，我们引入了一个谐波构建模块（HCM），以便在空间上将潜在表示扩展为傅里叶级数。我们先回顾经典的二维傅里叶变换理论及其谐波函数。此外，我们还深入探讨所提出的 HCM 的细节，展示如何将经典傅里叶变换整合到神经网络的潜在空间中。

Preliminary of Fourier Series

傅里叶级数是一种用一组正弦函数来表示周期信号或有限信号的有效工具。形式上，一个有限一维信号 ${x(t), t<T \}$ 的周期延拓可以用振幅 - 相位形式的傅里叶级数表示为：

$\approx \frac{A_0}{2} + \sum_{k = 1}^{\infty} \mathcal{H}(k, t) \\ \mathcal{H}(k, t) = A_k \cos(2\pi \frac{k}{T} t - \varphi_n) \tag{1}$

其中 $k$ 表示离散频率， $\mathcal{H}(k, t)$ 表示信号的第 $k$ 次谐波分量。在 $\mathcal{H}(k, t)$ 中， $A_k$ 和 $\phi_{k}$ 分别表示振幅和相位。 $A_{0}$ 是直流分量。公式（1）中的傅里叶级数很容易推广到 $n$ 维的情况。在有限二维信号（比如一幅图像）的情境下，对应于 $x$ 和 $y$ 维度、频率为 $(u, v)$ 的谐波函数表示为：

$\mathcal{H}(x, y, u, v) = A(u, v) \cdot \cos[\frac{2 \pi}{T}(ux + vy) - \phi(u, v)] \tag{2}$

其中， $x < T$ , $ y < T$, $\cos[\frac{2 \pi}{T}(ux + vy)]$ 作为余弦基函数， $\mathcal{H}(x, y, u, v)$ 称为谐波函数。

在二维空间中，输入图像的幅度和相位可通过傅里叶变换（如二维快速傅里叶变换，2D FFT）得出。然而，傅里叶系数既不可学习，也无法直接压缩以形成信息瓶颈，因此不能直接助力图像复原或视觉表征学习任务。我们不使用快速傅里叶变换，而是通过编码器将幅度和相位（分别记为 $A (u, v)$ 和 $\phi(u, v)$ ）作为瓶颈特征图进行学习，这样可实现灵活的维度配置。为了重建原始图像，我们利用谐波分量建模（HCM）生成一组经过学习的谐波函数（公式（2）），解码器则充当一个可学习的求和算子，模拟公式（1）中的重建过程。

其核心要点在于，傅里叶系数在表示图像（即便包含复杂细节的图像）时具有紧凑性。例如，一张复杂的纹理图像可能仅有少数几个重要的频率分量。通过将傅里叶系数转化为可学习的表示形式，我们能够设计出一种具有极窄瓶颈的自动编码器。

Encoding Fourier Coefficients

CosAE 的编码器负责学习瓶颈潜在空间的傅里叶系数。给定一个大小为 $\times P$ 的二维方形图像块 $I_p$ ，编码器将输入信号压缩为一个维度为 $\times c$ 的瓶颈特征向量，该向量表示对相应的幅度和相位。

$\in \mathcal{R}^{c}, \phi \in \mathcal {R}^{c}] = \mathbf{E}(I_p) \tag{3}$

为了能够适用于任意分辨率的图像，在实际操作中，我们将编码器的下采样步长设置为 $P$ ，这样幅度 $A$ 和相位 $\phi$ 都是 $c$ 通道维度的特征图，其空间分辨率是输入图像的 $1/ P$ 。在 CosAE 中，我们会采用相对较大的步长 $P$ ，以便瓶颈空间可以非常小：对于人脸图像，我们将下采样步长保持为 $P = 64$ （对于 64 x 64 的图像块，会得到一个维度为 $\mathbb{R}^{2 \times c}$ 的一维瓶颈，见公式 (3)）；而对于自然图像，我们使用的步长为 $P = 32$ 。值得注意的是，如下面所介绍的，CosAE 通过将这些系数与作为潜在特征图的余弦基函数相结合，仍然能够如实地重建高频细节。

Constructing Harmonic via HCM

我们引入一个谐波构建模块（图 2 中间框），它作用于瓶颈层，将编码器输出的预测系数转换为一组二维谐波函数，即公式（2）中的 $H$ 。以二维方形图像块 $I_{p} $为例。编码器对该图像块进行处理，生成 $c$ 对幅度 $\in \mathbb{R}^{c}$ 和相位 $\phi \in \mathbb{R}^{c}$ 。然后，这些与 $c$ 个对应的余弦基函数相结合。每个基函数都是一个 $\times T$ 的二维余弦波形，其中 $T$ 可以灵活调整，以达到所需的输出分辨率。因此，根据公式（2），每对傅里叶系数 $A (u, v)$ 和 $\phi(u, v)$ 将在空间上扩展为一个 $\times T$ 的谐波函数（空间扩展在图 2 中可视化为两个红色块之间的转换）。这 $c$ 个谐波函数构成一个 $\times T \times c$ 的特征图，随后进行解码以重建图像块 $I_{p}$ 。具体来说，在训练过程中，余弦基函数的大小 $T$ 可以定制，例如在灵活分辨率超分辨率任务中。我们将在后续章节详细介绍。

值得注意的是， $T$ 不能过小，否则无法构成有效的二维余弦波形。为了兼顾模型效率，我们设定 $T = P /2$ ，即对于人脸图像 $T = 32$ ，对于自然图像 $T = 16$ ，此时解码器的上采样步长设置为 2。决定余弦基函数的 $c$ 对频率 $\in \mathbb{R}^c，v \in \mathbb{R}^c)$ 同样是可学习的。

Learning Frequencies as Network Parameters

许多先前的研究工作也像学习傅里叶系数那样，从编码器中学习 $u$ 和 $v$ ，也就是说，频率取决于输入信号。然而，这是最优解吗？为了回答这个问题，我们简要回顾一下谐波函数的特性。

谐波。在信号处理中设置频率有两个关键原则。首先，公式（2）中的频率 $(u, v)$ 是以与信号长度成反比的增量进行采样的。其次，频率必须满足采样定理以避免混叠。为了确定频率的范围，我们对谐波函数应用傅里叶变换（为简化起见，我们采用一维情况，并忽略幅度和相位），

$F(\mathcal{H}_k(w)) = \int_{-\infty}^{\infty} \cos(\frac{2\pi}{T}kt)e^{-i2\pi wt} dt \\ = \pi [\delta (2\pi(w - \frac{k}{T})) + \delta (2\pi(w + \frac{k}{T}))] \tag{4}$

谐波的带宽 $B = k / T$ 。根据奈奎斯特 - 香农定理，对于采样率 $f_s$ ，当带宽限制 $\frac{k}{T} < \frac{f_s}{2}$ 时，可保证完美重构。对于离散信号，例如图像，其采样间隔（像素） $f_s = 1$ ，为防止混叠并确保准确的谐波表示， $k$ 可限制为 $\frac{T}{2}$ 。

基于上述经典离散傅里叶变换（DFT）的设计，我们在 $[0, T /2]$ 范围内初始化 $c$ 对 $(u, v)$ 。然而，为了实现更灵活的设计， $(u, v)$ 集合应适应不同需求，比如为提高分辨率增加频率对，或者为提高效率减少频率对。在这种情况下，均匀采样可能导致覆盖稀疏。例如，当 $T = 64$ 且基图数量设为 64 时， $u$ 和 $v$ 都会很稀疏，采样值为 $[0, 4, 8, 16, 32]$ 。由于自然图像的频率并非均匀分布，我们将 $(u, v)$ 设为所有输入图像共享的可学习参数，这样在训练过程中就能对各种频率进行有效建模。与之前基于单个输入图像调整频率的方法相比，我们的设计更紧密地遵循经典原理，使用跨越整个范围的全局共享频率。这种设计还解决了先前模型中常见的一个局限，即过于强调主导自然图像的低频部分，从而实现了更平衡、更全面的频率表示。

Decoding Harmonic for Reconstruction

一旦得到 $\mathcal{H}_{T \times T \times c}$ ，我们引入一个解码器网络，得到 RGB 像素：

$\mathcal{D}(\{\mathcal{H}(x, y, u, v) \}); (u, v) < T/2, (x, y) < T \tag{5}$

该架构的设计可以灵活多样。例如，我们的框架既可以采用隐式表示网络，在训练过程中仅对一小部分像素进行采样解码；也可以采用基于卷积或 Transformer 的解码器，这些解码器利用整个潜在空间，并能够融入诸如判别器等额外模块，以生成全尺寸的输出。

Network Implementation

FFT 用于额外输入：
我们发现将输入图像的二维 FFT 作为额外通道可以提升性能，因为这直观上与编码器在潜在空间中捕获傅里叶系数的任务相契合。具体来说，我们应用 FFT 为每个 RGB 通道生成实部和虚部的二维映射，然后将它们与原始 RGB 通道拼接在一起，为网络生成一个 9 通道的输入映射。
网络架构：
我们基于 [6] 构建编码器和解码器。为了在编码器中实现更大的下采样步长，我们添加了额外的 ResNet 块，并将注意力层放置在瓶颈附近。相应地，我们使用一个相对轻量级的解码器，其上采样因子为 2。我们在所有下游任务中都使用的图像块来训练 CosAE，并且在所有实验中设置。值得注意的是，我们对每个瓶颈特征向量进行独立处理，从而允许对任意尺寸的图像进行推理。这是通过在解码器的每个特征单元内用窗口注意力替代全局注意力来实现的。同样，在推理时，对于尺寸不为的输入，我们也会在每个不重叠的区域使用窗口注意力。
目标函数：
我们遵循 [6]，使用损失和感知损失作为重建目标来训练 CosAE，记为。此外 $L_{rec} = L_1 + L_{lpips}$ ，我们应用基于图像块的判别器，记为 $L_{GAN}$ ，以增强重建图像的视觉质量。整个网络连同 $c$ 对可学习频率 $\in \mathbb{R}^{2 \times c}$ ，使用 $L_{rec} + \lambda L_{GAN}$ 进行训练。我们对自适应权重 $\lambda$ 保持与 [6] 相同的设置。

Experimental Results

在本文中，我们重点关注灵活分辨率超分辨率和盲图像去噪的应用，而将其他与自动编码器相关的潜在应用留作未来的研究。我们在人脸和自然图像领域对这两项任务进行评估。

对于人脸和自然图像领域，我们使用从 1 到 8 的放大因子来训练 CosAE。在训练期间，我们从训练集中采样一个高分辨率（HR）图像 $I_{HR} \in \mathcal{R}^{H \times W \times 3}$ ，并使用均匀采样的比率 $(1, 8)$ 对其进行下采样，以获得低分辨率（LR）输入图像 $I_{LR} \in \mathcal{R}^{H/r \times W/r \times 3}$ 。

为了训练具有灵活输出分辨率的超分辨率模型，我们接着采样一个放大比率 $\in (1, r)$ ，让 CosAE 将输入图像放大为 $I_{out} \in \mathcal{R}^{urH \times urW \times 3}$ 。这是通过设置 $T = T_{max} u/r$ 来实现的，其中 $T_{max}$ 指的是前面节中的 $T$ ，它将低分辨率图像放大到高分辨率图像的原始尺寸。