扩散模型在视频压缩中的创新应用与优化

暗茧

1. 项目概述：当视频压缩遇上扩散模型

在视频生成领域，自动编码器扮演着"信息蒸馏器"的关键角色——它需要将高维像素数据（如1080p视频每帧约200万像素）压缩为紧凑的潜在表示。传统3D-CNN架构的自动编码器存在三个致命伤：固定压缩率导致简单视频浪费token资源、CNN的刚性结构阻碍变长建模、确定性解码器难以恢复丢失的细节。这就像用固定尺寸的集装箱运输不同体积的货物，小物件被迫填满整个集装箱，而大物件却被强行切割。

One-DVA的创新在于将Transformer的灵活性、扩散模型的生成能力与自适应压缩机制相结合。其核心突破点包括：

动态token分配：通过可学习评分机制，为简单视频分配较少token（如200个），复杂视频分配更多（如1000个）
混合潜在表示：结构潜在（structural latent）捕获空间低频信息，1D潜在序列编码时空高频细节
生成式解码：将重建视为条件生成任务，利用扩散模型补全压缩丢失的细节

实测表明，在4×16×16压缩比下，One-DVA的PSNR达到36.48dB，优于HunyuanVideo（35.54dB）和Wanx2.2（35.23dB）等SOTA模型。更关键的是，它能根据视频复杂度动态调整token数量，平均节省44.2%的token时仅损失1dB PSNR。

2. 核心架构设计解析

2.1 查询式Transformer编码器

编码器采用ViT变体，其创新点在于双路特征提取机制：

python复制class QueryViTEncoder(nn.Module):
    def __init__(self):
        self.patch_embed = LinearPatchifier(patch_size=16)  # 16×16块嵌入
        self.queries = nn.Parameter(torch.randn(N, D))  # 可学习查询向量
        self.transformer_blocks = TransformerBlock(depth=12)
        
    def forward(self, x):
        # 路径1：时空特征提取
        patches = self.patch_embed(x)  # [B, T×H×W, D]
        patches = patches + positional_encoding(patches)
        
        # 路径2：查询交互
        queries = self.queries.unsqueeze(0).expand(B, -1, -1)
        queries = queries + query_positional_encoding(queries)
        
        # 双路特征融合
        features = torch.cat([patches, queries], dim=1)
        features = self.transformer_blocks(features)
        
        # 分离输出
        patch_features = features[:, :patches.size(1)]
        query_features = features[:, patches.size(1):]
        
        return patch_features, query_features

关键设计考量：

多分辨率训练：在输入序列中添加特殊token标识视频的H/W/T/长宽比，使模型适应不同分辨率
通道压缩：将2048维特征降至64维，减少后续扩散模型的condition维度
混合潜在构造：结构潜在来自ViT特征的区域池化，1D潜在来自前K个查询向量（K根据压缩比动态调整）

2.2 变长丢弃机制

该模块实现动态token分配的核心算法：

运动评分：计算连续帧间像素差异的L2距离，归一化为[0,1]
丢弃策略：从序列尾部向前随机丢弃token，丢弃比例l ~ Beta(α=1+score, β=3-score)
极端情况：10%样本保留全部token，10%仅使用结构潜在（强制模型学习退化情况）

实践发现：对运动剧烈的视频（score>0.7），保留60%以上token才能维持PSNR>35dB；静态场景（score<0.3）仅需20%token

2.3 像素空间扩散解码器

解码器采用DiT架构，但有三处关键改进：

组件	传统DiT	One-DVA改进	优势
条件注入方式	Cross-Attention	特征拼接+自适应层归一化	保留空间结构信息
噪声调度	线性	对数正态分布	厚尾分布促进细节生成
上采样模块	转置卷积	像素洗牌+长跳跃连接	减少棋盘伪影

损失函数组合体现生成与重建的平衡：

math复制L = 0.8L_{diff} + 0.1L_{perceptual} + 0.05L_{kl} + 0.05L_{repa}

其中REPA损失（Representation Alignment）确保1D潜在与结构潜在的空间一致性，这对后续LDM训练至关重要。

3. 训练策略与技巧

3.1 两阶段训练法

阶段一：确定性预训练（415K步）

解码器输入纯噪声（t≡1）
禁用变长丢弃（固定使用100% token）
目标：迫使编码器提取最大信息量

阶段二：随机后训练（800K步）

启用对数正态timestep采样
引入变长丢弃（平均保留55.8% token）
10%样本使用单步采样（t=1→0），其余多步扩散

对比实验：端到端训练导致rFVD升高115.6→230.1，证明分阶段策略的必要性

3.2 生成适配微调

为优化下游LDM性能，两个关键操作：

潜在空间对齐

python复制def align_loss(structural_latent, query_latent):
    # 计算每个查询向量与空间区域的余弦相似度
    sim_matrix = F.cosine_similarity(
        query_latent.unsqueeze(2), 
        structural_latent.flatten(2).unsqueeze(1),
        dim=-1)
    # 强制每个查询关注最匹配的空间区域
    top1_sim = sim_matrix.max(dim=2)[0].mean()
    return 1 - top1_sim