PFP: Pretraining Frame Preservation in Autoregressive Video Memory Compression

Authors: Lvmin Zhang, Shengqu Cai, Muyang Li, Chong Zeng, Beijia Lu, Anyi Rao, Song Han, Gordon Wetzstein, Maneesh Agrawala Affiliations: Stanford University, MIT, Carnegie Mellon University, HKUST

1. Motivation (研究动机)

核心问题: 自回归视频生成中 Context Length 与 Quality 的根本矛盾

在自回归视频生成中, 模型将历史帧作为 context 来生成未来帧。随着视频长度增加, 历史 context 的长度线性增长, 带来两个核心瓶颈:

  1. 计算瓶颈: 一个 60 秒、480p、24fps 的视频, 经过标准 Hunyuan/Wan VAE 和 patchify 后, history context length 达到 832/16 × 480/16 × 60 × 24/4 = 561,600, 在消费级 GPU 上根本无法处理
  2. 质量-长度权衡 (Quality-Length Trade-off): 各种压缩方案 (sliding window、token merging、compact VAE、FramePack 的多级 patchify) 都在缩短 context 的同时丢失高频细节, 这是一个尚未找到最优平衡点的根本性 trade-off

现有方案的不足

方案问题
Naive sliding window保持固定 context length, 但完全丢失长程历史
Token merging合并率越高, 细节丢失越严重
Compact VAE (LTXV, DC-AE)可压缩到更紧凑的 latent, 但牺牲高频信息
FramePack 多级 patchify可在多层级压缩, 但同样以高频图像细节为代价
Sparse/Linear attention减少计算开销, 但 linear layer 仍有训练/推理成本

PFP 的核心洞察

一个好的视频压缩机制的关键能力指标是: 能否在任意时间位置高质量检索 (retrieve) 单帧的高频细节。当压缩率很高时, 完美检索不可能, 因此目标变为 最大化任意帧的检索质量

这导出了一个自然的预训练目标: 先独立训练一个 memory compression model, 使其学会将长视频压缩为短 context, 同时保持对任意时间位置帧的高保真检索能力。训练完成后, 再将其作为 memory encoder 接入自回归视频 diffusion 模型进行微调。


2. Idea (核心思想)

PFP 的核心思想是 两阶段训练: 先预训练压缩模型, 再微调自回归生成模型。

2.1 阶段一: Memory Compression Model 预训练 (Frame Retrieval)

将视频压缩建模为一个 帧检索 (frame retrieval) 任务:

  • 输入: 一段 >20 秒的长视频历史 H
  • 压缩: 通过 compression model φ(·) 将 H 压缩为 ~5k length 的 context φ(H)
  • 目标: 从 φ(H) 中检索任意时间位置 t 的帧, 使其与原始帧尽可能一致

关键设计: 使用 随机帧选择 + noise-as-mask 策略, 随机保留一些帧不变, mask 掉其余帧, 让 diffusion model 重建被 mask 的帧。这防止模型学到 “作弊” 解法 (如只编码首尾帧)。

2.2 阶段二: 自回归视频模型微调

将预训练好的 compression model φ(·) 作为 memory encoder, 与 DiT (加 LoRA) 一起微调:

  • History 经过 φ(·) 压缩后与 target context concat
  • DiT 通过 LoRA 学习利用压缩后的历史 context 生成下一段视频
  • 推理时自回归地将生成结果加入历史, 持续生成

2.3 网络架构: 双分支 + 残差增强

压缩模型采用轻量级双分支架构:

  • 低分辨率分支: 低帧率低分辨率视频 → VAE → DiT patchify + first projection → DiT context (提供语义信息)
  • 高分辨率分支: 高帧率高分辨率视频 → VAE → 3D Conv 压缩 → 残差增强向量 (提供高频细节)

关键: 高分辨率分支的特征 不经过 VAE 的 16-channel bottleneck, 而是直接在 DiT 的 inner channel (如 3072 for WAN-5B) 上输出, 最大限度保留保真度。


3. Method (方法)

3.1 预训练 Memory Compression Model

Figure 2 解读: Memory compression model 的预训练流程。左侧: 从长视频 (>20秒) 中随机挑选若干帧保持可见, 其余帧被 mask (通过加噪实现 noise-as-mask)。可见帧和 masked 帧一起通过 VAE 编码, 送入 Memory compression model φ(·) 压缩为 ~5k 长度的 context。右侧: 被选中的帧被 clone 作为 estimation target, 经过 VAE 编码后加噪, 与压缩 context concat 后送入 DiT (带 LoRA)。DiT 学习从压缩 context 中重建目标帧。可选地, 原始 history 可通过 Enhance 路径直接输入 DiT 底层辅助训练。

Preliminaries: Flow Matching 公式

自回归视频 diffusion 模型使用 rectified flow matching:

其中 是 diffusion timestep, 是 clean latent。学习目标:

压缩预训练目标

给定 history H, compression model φ(·), 随机帧索引集合 Ω ⊆ N:

其中 是在索引 Ω 处保留的帧, 其余帧被 mask。模型需要从压缩 context φ(H) 中检索 Ω 位置的帧。

Noise-as-mask 策略: mask 方式为从 (shifted logit-normal) 采样噪声水平加到被 mask 的帧上, 本质上让 diffusion 系统尝试在任意位置重建目标帧。

随机帧选择的重要性: 如果帧选择不随机 (如只取首尾帧), 模型会学到只编码少数帧的 “作弊” 策略。随机选择迫使模型将所有帧信息编码到压缩 context 中。

3.2 压缩模型网络架构

Figure 3 解读: Memory compression model 的双分支架构。上方高分辨率分支: 480p 高帧率视频 (如 h480w832f480, 20秒) 经 VAE 编码为 HR latent (60×104×120×16), 通过级联 3D Conv + SiLU 逐步压缩通道 (64→128→256→512), 最终输出 7×13×60×3072 的残差增强向量。下方低分辨率分支: 120p 低帧率视频经 VAE 编码为 LR latent (15×26×60×16), 经 DiT 的 patchify + first projection 变为 h2w2f1×3072, 输出 5460×3072 的 context。两个分支通过 add 操作融合。可选地, 可通过 cross-attention 将 encoder 的 hidden state 注入 DiT 的每个 block, 进一步增强一致性。

架构细节:

  • 3D Conv 压缩层: 先压缩时间维度, 再压缩空间维度 (通过 stride)
  • 隐藏通道: 64 → 128 → 256 → 512, 最终通过 1×1 conv 投射到 DiT inner channel (3072 或 5120)
  • 压缩率标记: H×W×T, 如 4×4×2 表示在 latent space 上空间压缩 4×4、时间压缩 2×
  • 压缩模型几乎全为卷积, 支持 on-the-fly concatenation (无需重算历史)

3.3 微调自回归视频模型

Figure 4 解读: 自回归视频模型的微调与推理。(a) 微调: History 经过预训练好的 memory compression model 压缩, 与 future target (加噪) concat 后送入带 LoRA 的 DiT 计算 loss。(b) 推理: 自回归地生成, 每次将新生成的段落加入 history, 经压缩模型压缩后作为下一次生成的 context。

微调目标:

3.4 替代架构扩展

添加 sliding window: 与小 sliding window (3 latent frames) 结合, 可调节镜头切换频率, 支持连续长镜头生成 (代价: 略微增加 context length)

Cross-attention 增强: 将 encoder 倒数第二层特征通过 cross-attention 注入 DiT 每个 block (类似 IP-Adapter), 进一步增强细节一致性 (如超市货架排列顺序), 代价是额外计算开销

多压缩模型组合: 同时使用多个 compression model, 如标准 4×4×2 + 2×2×8 (高时间压缩 + 高空间保留), 以 2× context length 为代价获得更好的细节一致性

4. Experimental Setup (实验设置)

4.1 实现细节

  • 预训练: 8×H100 GPU clusters
  • 微调: 1×A100-80G (batch 64, window size 2-3) 或 batch 32 (window 4-5)
  • DiT LoRA rank: 128
  • 基础模型: HunyuanVideo 12.8B, Wan 2.2 (5B, 14B)
  • 数据: ~500 万互联网视频, 由 Gemini-2.5-flash / QwenVL 生成 storyboard 格式 caption (带时间戳)
  • 测试集: 1000 条 Gemini-2.5-pro 生成的 storyboard prompt + 4096 条 unseen 视频

5. Experimental Results (实验结果)

5.1 消融实验: 压缩架构对比

Figure 5 解读: 不同压缩架构在不同压缩率下的重建效果对比。Origin 是原始图像。Large Patchifier (等价于 FramePack) 通过增大 patchify kernel 实现压缩, 产生较大的结构变化。Only LR 只保留低分辨率分支, Without LR 只保留高分辨率分支, 两者都偏离原始图像较多。Proposed (4×4×2) 在高压缩率下仍保持大部分图像外观。Proposed (2×2×2) 和 (2×2×1) 随压缩率降低细节保留更好。

Table 1: 压缩结构定量对比

MethodPSNR ↑SSIM ↑LPIPS ↓
Large Patchifier* (4×4×2)12.930.4120.365
Only LR (4×4×2)15.210.4720.212
Without LR (4×4×2)15.730.4230.198
Proposed (4×4×2)17.410.5960.171
Proposed (2×2×2)19.120.6830.152
Proposed (2×2×4)18.630.6370.153
Proposed (2×2×1)20.190.7050.121

核心发现:

  • 双分支架构 (Proposed) 在所有指标上大幅超越单分支和 Large Patchifier
  • 高分辨率残差增强分支对保真度至关重要 (绕过 VAE 16-channel bottleneck)
  • 压缩率越低 (如 2×2×1), 检索质量越好, 但 context length 越长

5.2 预训练的影响

Figure 6 解读: 有无预训练的效果对比。相同 20 秒历史输入, 相同架构, 分别展示 with pretraining (proposed) 和 without pretraining 的中间帧。With pretraining 模型保持了面部特征、服装、全局视频风格、叙事线和镜头运动的强一致性。Without pretraining 模型在身份、服装等方面出现明显不一致。

5.3 视频内容一致性定量结果

Table 2: 不同方案的视频一致性对比

MethodCloth ↑Identity ↑Instance ↑ELO ↑
WanI2V + QwenEdit (1p)94.1068.4585.21/
WanI2V + QwenEdit (2p)95.0968.2291.191198
WanI2V + QwenEdit (3p)94.2867.3080.34/
Only LR (4×4×2)91.9869.2285.321194
Without LR (4×4×2)89.6467.4182.85/
Without Pretrain (4×4×2)87.1266.9981.13/
Proposed (4×4×2)96.1270.7389.891216
Proposed (2×2×2)96.7172.1290.271218

核心发现:

  • Proposed 方法在 Cloth 和 Identity 一致性上领先, User Study ELO 也最高
  • 预训练对一致性提升显著: Without Pretrain 在各项指标上明显落后
  • 2×2×2 压缩率进一步提升身份一致性 (Identity 72.12 vs 70.73)

5.4 不同基础模型对比

Table 3: 不同基础模型的性能

MethodAesthetic ↑Clarity ↑Dynamics ↑Semantic ↑ELO ↑
HunyunVideo 12.8B (4×4×2)61.2767.4971.2226.291189
Wan 2.2 14B* (4×4×2)67.2269.3769.8127.421234
Wan 2.2 5B (4×4×2)66.2569.0165.1325.991215
Wan 2.2 5B (2×2×2)66.3768.9566.2926.131224

Wan 14B 在美学和清晰度上表现最佳, User Study ELO 最高。

5.5 Error Accumulation (Drifting)

一个重要发现: error accumulation 高度依赖训练数据集。当训练数据以密集镜头切换的 Short-style 视频为主时, drifting 几乎不可见。对于需要长单镜头连续生成的场景, 可结合 sliding window、Self-Forcing 等策略缓解。

5.6 Storyboard 生成质量

Figure 7 解读: Storyboard 驱动的多镜头生成结果。展示了烘焙、记者采访等场景, 模型在多个 prompt 之间保持了角色身份、服装和物体的一致性。Storyboard prompt 由 Gemini-2.5-pro 生成。


6. Conclusion

核心贡献

  1. 首次提出将视频压缩建模为帧检索预训练任务: 通过显式优化任意帧的高保真检索, 学习紧凑且保留高频细节的 context 表示
  2. 双分支残差增强架构: 低分辨率语义分支 + 高分辨率残差分支 (绕过 VAE bottleneck), 在相同压缩率下大幅提升保真度
  3. 两阶段训练范式: 先预训练 compression model (数据利用效率高, 可独立评估), 再微调 autoregressive model (训练成本降低)
  4. 系统性消融研究: 定量分析了 context length 与 quality 的 trade-off, 不同架构、压缩率、基础模型的影响
  5. 实用性: 20 秒视频压缩为 ~5k context, 支持 RTX 4070 12GB 级别的消费级 GPU

局限性

  • 压缩率与质量仍存在 trade-off, 高压缩率下面部等细节仍有损失
  • Error accumulation 在长单镜头场景仍需额外策略处理
  • 当前实验主要基于 Wan/HunyunVideo, 对其他架构的泛化性待验证

7. Figure 解读总结

图号文件内容
Fig.1teaser.svg自回归生成 22+ 秒视频的帧序列, 展示全程身份/服装一致性
Fig.2f2.svgMemory compression model 预训练流程: 随机帧选择 + noise-as-mask + VAE + DiT 重建
Fig.3f3.svg双分支压缩架构: HR 3D Conv 分支 + LR patchify 分支, 通过 add 融合
Fig.4f1.svg自回归微调与推理流程: compression model + LoRA DiT
Fig.5ab1.png不同架构的重建质量对比: Proposed 双分支显著优于单分支和 Large Patchifier
Fig.6ab2.svg有无预训练对比: 预训练显著提升面部/服装/风格一致性
Fig.7qual.svgStoryboard 多镜头生成结果, 跨镜头保持角色一致性
Fig.8sli.svg添加 sliding window 实现连续长镜头 (减少镜头切换)
Fig.9cross.svgCross-attention 增强在货架等高细节场景提升一致性
Fig.10mixing.svg多 compression model 组合: 4×4×2 + 2×2×8, 以更长 context 换更好细节

8. Python Pseudocode

import torch
import torch.nn as nn
 
# ============================================================
# Stage 1: Memory Compression Model (预训练)
# ============================================================
 
class MemoryCompressionModel(nn.Module):
    """双分支视频压缩模型: HR残差增强 + LR语义"""
 
    def __init__(self, dit_inner_channels=3072):
        super().__init__()
        # 高分辨率分支: 3D Conv 级联压缩
        self.hr_encoder = nn.Sequential(
            nn.Conv3d(16, 64, kernel_size=3, stride=(2,2,1), padding=1),   # 时间先压缩
            nn.SiLU(),
            nn.Conv3d(64, 128, kernel_size=3, stride=(2,2,1), padding=1),
            nn.SiLU(),
            nn.Conv3d(128, 256, kernel_size=3, stride=(1,2,2), padding=1), # 再压缩空间
            nn.SiLU(),
            nn.Conv3d(256, 512, kernel_size=3, stride=(1,1,1), padding=1),
            nn.SiLU(),
            nn.Conv3d(512, dit_inner_channels, kernel_size=1),  # 投射到DiT内部维度
        )
        # 注意: HR分支输出不经过VAE的16-channel bottleneck
        # 直接在dit_inner_channels (3072) 维度上工作, 保留最大保真度
 
    def forward(self, hr_vae_latent, lr_dit_context):
        """
        Args:
            hr_vae_latent: [B, 16, T_h, H_h, W_h] 高分辨率VAE latent
            lr_dit_context: [B, N_lr, C] 低分辨率经DiT patchify后的context
        Returns:
            compressed_context: [B, N_total, C] 压缩后的history context (~5k tokens)
        """
        # 高分辨率残差增强
        hr_features = self.hr_encoder(hr_vae_latent)  # [B, 3072, T', H', W']
        hr_tokens = hr_features.flatten(2).permute(0, 2, 1)  # [B, N_hr, 3072]
 
        # 与低分辨率context相加融合
        compressed_context = lr_dit_context + hr_tokens  # element-wise add
        return compressed_context
 
 
# ============================================================
# Stage 1: 预训练过程 (Frame Retrieval Task)
# ============================================================
 
def pretrain_compression_model(
    compression_model: MemoryCompressionModel,
    dit_with_lora: nn.Module,
    vae: nn.Module,
    dataloader,
    optimizer,
):
    """预训练: 从压缩context中检索任意时间位置的帧"""
    for video_batch, text_prompts in dataloader:
        # video_batch: [B, T_total, 3, H, W], 20+秒视频
 
        # Step 1: 随机选择帧索引 Omega
        T_total = video_batch.shape[1]
        num_keep = torch.randint(1, T_total // 2, (1,)).item()
        omega = torch.randperm(T_total)[:num_keep].sort().values  # 随机帧索引
 
        # Step 2: 构造masked history (noise-as-mask)
        history = vae.encode(video_batch)  # [B, 16, T, H/8, W/8]
        mask_noise_levels = sample_shifted_logit_normal(T_total, low=0.2, high=1.0)
        for t in range(T_total):
            if t not in omega:
                noise = torch.randn_like(history[:, :, t])
                history[:, :, t] = (1 - mask_noise_levels[t]) * history[:, :, t] \
                                   + mask_noise_levels[t] * noise
 
        # Step 3: 压缩history
        hr_latent = vae.encode(video_batch)  # 高分辨率
        lr_video = F.interpolate(video_batch, scale_factor=0.25)  # 降采样
        lr_latent = vae.encode(lr_video)
        lr_context = dit_with_lora.patchify_and_project(lr_latent)
 
        compressed = compression_model(hr_latent, lr_context)  # ~5k tokens
 
        # Step 4: 构造retrieval target (clone选中帧)
        target_frames = history[:, :, omega]  # clean frames at omega positions
 
        # Step 5: Flow matching loss
        t_i = sample_shifted_logit_normal(batch_size=video_batch.shape[0])
        epsilon = torch.randn_like(target_frames)
        x_t = (1 - t_i) * target_frames + t_i * epsilon  # noisy target
 
        # 将noisy target与compressed context concat
        context = torch.cat([compressed, x_t.flatten(2).permute(0, 2, 1)], dim=1)
 
        pred = dit_with_lora(context, t_i, text_prompts)
        loss = F.mse_loss(pred, epsilon - target_frames)  # velocity prediction
 
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
 
 
# ============================================================
# Stage 2: 自回归视频模型微调
# ============================================================
 
def finetune_autoregressive(
    compression_model: MemoryCompressionModel,  # frozen or jointly trained
    dit_with_lora: nn.Module,
    vae: nn.Module,
    dataloader,
    optimizer,
):
    """微调: 用压缩history生成下一段视频"""
    for video_batch, text_prompts in dataloader:
        # 将视频分为 history 和 future
        history_frames = video_batch[:, :-1]  # 前20秒
        future_frames = video_batch[:, -1:]   # 下1秒 (target)
 
        # 压缩history
        compressed_history = compression_model(
            vae.encode(history_frames),
            dit_with_lora.patchify_and_project(vae.encode(downsample(history_frames)))
        )
 
        # Flow matching on future frames
        X_0 = vae.encode(future_frames)
        t_i = sample_shifted_logit_normal(batch_size=video_batch.shape[0])
        epsilon = torch.randn_like(X_0)
        X_t = (1 - t_i) * X_0 + t_i * epsilon
 
        context = torch.cat([compressed_history, X_t.flatten(2).permute(0, 2, 1)], dim=1)
        pred = dit_with_lora(context, t_i, text_prompts)
        loss = F.mse_loss(pred, epsilon - X_0)
 
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
 
 
# ============================================================
# Inference: 自回归长视频生成
# ============================================================
 
@torch.no_grad()
def autoregressive_generate(
    compression_model: MemoryCompressionModel,
    dit_with_lora: nn.Module,
    vae: nn.Module,
    storyboard_prompts: list,  # [(timestamp, prompt), ...]
    num_iterations: int = 20,
    num_diffusion_steps: int = 50,
):
    """自回归生成: 每次生成~1秒, 累计压缩history"""
    history_latents = None
    all_frames = []
 
    for i in range(num_iterations):
        # 获取当前时间戳对应的prompt
        current_prompt = get_prompt_for_timestamp(storyboard_prompts, i)
 
        if history_latents is not None:
            # 压缩累积的history (几乎全卷积, 支持on-the-fly concat)
            compressed = compression_model(history_latents)  # ~5k tokens
        else:
            compressed = None
 
        # 多步去噪生成下一段
        x_t = torch.randn(1, 16, T_segment, H, W)  # pure noise
        for step in reversed(range(num_diffusion_steps)):
            t_i = step / num_diffusion_steps
            context = torch.cat([compressed, x_t.flatten(2).permute(0, 2, 1)], dim=1) \
                      if compressed is not None else x_t.flatten(2).permute(0, 2, 1)
            velocity = dit_with_lora(context, t_i, current_prompt)
            x_t = x_t - velocity * (1.0 / num_diffusion_steps)  # Euler step
 
        # Decode并存储
        new_frames = vae.decode(x_t)
        all_frames.append(new_frames)
 
        # 将新生成的latent加入history
        if history_latents is None:
            history_latents = x_t
        else:
            history_latents = torch.cat([history_latents, x_t], dim=2)  # concat on time
 
    return torch.cat(all_frames, dim=2)
 
 
def sample_shifted_logit_normal(size, low=0.0, high=1.0):
    """Shifted logit-normal distribution for flow matching timesteps"""
    u = torch.rand(size)
    u = u * (high - low) + low
    return torch.sigmoid(torch.erfinv(2 * u - 1) * 1.4142)

9. Code Mapping

PFP 的代码目前集成在 FramePack 仓库中, 尚未单独开源完整的预训练代码。以下是论文概念与 FramePack 代码/架构的映射关系:

论文概念代码对应说明
Memory Compression Model φ(·)FramePack 中的 context packing encoder双分支 3D Conv + patchify 架构
HR 残差增强分支3D Conv encoder → 直接输出 dit_inner_channels绕过 VAE 16-ch bottleneck, 输出 3072/5120 维
LR 语义分支DiT patchify + first projection复用 DiT 现有的 patchify 层处理低分辨率 latent
预训练 (Frame Retrieval)独立训练 compression model随机帧 mask + diffusion 重建目标
微调 (Autoregressive)LoRA finetuning on Wan/HunyunVideo DiTrank=128, 将 φ(H) 作为 context conditioning
Noise-as-maskshifted logit-normal 与 flow matching 的噪声调度一致
Sliding window 扩展3 latent frames sliding window与 compression context concat, 减少镜头切换
Cross-attention 增强类 IP-Adapter, encoder hidden → DiT cross-attn可选, 增强高细节场景一致性
压缩率 4×4×2latent space H=4, W=4, T=2 的 stride最终压缩率 = latent rate × patchify rate × compression rate
Storyboard promptsGemini-2.5-pro / QwenVL 生成带时间戳的多镜头 prompt, 自回归时取最近时间戳