2D小波分解高频分量在扩散模型中的应用：技术选型与避坑指南

1次阅读

没有评论

共计 2036 个字符，预计需要花费 6 分钟才能阅读完成。

小波分解（Wavelet Decomposition）是图像处理中的经典工具，通过高通和低通滤波器将图像分解为不同频率的分量。传统应用包括去噪（如 Donoho et al., 1995）和压缩（JPEG2000 标准）。其数学表达为：

$$
W_f(a,b) = \frac{1}{\sqrt{a}} \int_{-\infty}^{\infty} f(t)\psi\left(\frac{t-b}{a}\right)dt
$$

其中 $\psi$ 是小波基函数，$a$ 和 $b$ 分别控制尺度和位移。2D 情况下通过行列分离实现，得到 LL（低频）、LH（水平高频）、HL（垂直高频）、HH（对角线高频）四个子带。

直接将高频分量输入扩散模型（如 DDPM）会引发以下问题：

训练不稳定 ：高频分量包含大量噪声样成分，导致梯度幅值剧烈波动（参见 Song et al., ICLR 2021）
细节失真 ：高频信息在扩散过程后期被过度平滑，生成图像出现伪影（如图 1 中的棋盘效应）
收敛困难 ：高频能量分布不均匀导致损失函数存在局部极小值（类似模式崩溃现象）

通过 PCA 或自编码器将高频分量映射到低维空间。以 CelebA 为例，原始 512×512 图像经 3 层小波分解后高频维度为 1792×1792，经 PCA 可压缩至 256 维。核心代码片段：

from sklearn.decomposition import IncrementalPCA

pca = IncrementalPCA(n_components=256)
hf_components = pca.fit_transform(hf_wavelet.reshape(-1, patch_size**2))

优点：计算效率高
缺点：可能损失高频细节

参考 CVPR 2022 的 WaveDiffusion 方案，构建金字塔结构：

对原始图像进行 3 级小波分解
在各尺度上分别应用扩散模型
通过门控机制融合特征：

$$
F_{fusion} = \sigma(W_g \cdot [F_{low}, F_{high}]) \odot F_{high} + (1-\sigma) \odot F_{low}
$$

在 UNet 中插入频域注意力模块（完整实现见下节），结构如图 2 所示。相比空间注意力，额外计算频域能量图：

$$
E_{freq} = \sum_{i=1}^C |\mathcal{F}(U_i)|^2
$$

其中 $\mathcal{F}$ 表示 FFT 变换。

import torch
import torch.nn as nn
import torch.fft

class FreqAttention(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.query = nn.Conv2d(channels, channels//8, 1)
        self.key = nn.Conv2d(channels, channels//8, 1)
        self.value = nn.Conv2d(channels, channels, 1)

    def forward(self, x):
        B, C, H, W = x.shape
        # 空间注意力分支
        q = self.query(x).view(B, -1, H*W)
        k = self.key(x).view(B, -1, H*W)
        v = self.value(x).view(B, -1, H*W)
        spatial_attn = torch.softmax(q @ k.transpose(1,2), dim=-1)

        # 频域注意力分支
        fft = torch.fft.rfft2(x, norm='ortho')
        energy = (fft.abs() ** 2).mean(dim=1)
        freq_attn = torch.sigmoid(energy).unsqueeze(1)

        # 特征融合
        out = (spatial_attn @ v).view(B, C, H, W)
        return out * freq_attn

在 CelebA-HQ 上测试 256×256 图像生成任务（batch_size=32）：