Unifying CNNs and ViTs for Learning-Efficient and Scalable Variational AutoEncoder

会议: ECCV 2026
论文: ECCV 2026
领域: 图像生成
关键词: 变分自编码器, 视觉分词器, Vision Transformer, 分辨率外推, 混合架构

一句话总结

TransVAE 提出一种融合浅层 CNN 局部前向提取与深层 Transformer 全局上下文建模的混合 VAE 架构，通过纯 RoPE 位置编码、多级卷积下采样与卷积前馈网络，实现极高的训练收敛效率、强鲁棒的单分辨率训练高分辨率外推能力以及 44M 到 2.3B 参数的可预测扩展性。

研究背景与动机

连续值变分自编码器（Continuous VAE）作为潜在扩散模型（Latent Diffusion Models, LDM）的核心视觉分词器，承担着将高分辨率像素空间压缩至低维紧凑隐空间的重任。长期以来，视觉分词器主要由纯卷积神经网络（CNN）主导，得益于局部连接与平移等变性归纳偏置，CNN 能够高效捕获细粒度纹理与颜色分布。然而，卷积核有限感受野导致其长程非局部依赖建模能力先天不足，若单纯依靠堆叠网络深度建立全局上下文，极易遭遇计算与优化瓶颈。

与此同时，Vision Transformer（ViT）借助全局自注意力机制展现出卓越的全图上下文捕获能力，但在连续值 VAE 分词器设定下，纯 ViT-VAE 暴露出三大关键缺陷：首先是分辨率外推能力极差，在单一分辨率训练后直接外推至高分辨率推理时，传统绝对位置编码（APE）插值会破坏已习得的空间先验并引发网格伪影；其次是细粒度局部先验缺失，导致色彩与高频纹理收敛极其缓慢且存在高频瑕疵；最后是扩展性收益不一致，单纯增大 ViT 参数量无法稳定提升图像重建质量。

针对纯 CNN 缺乏全局感受野与纯 ViT 缺失局部归纳偏置的结构性断层，本文的核心切入点是寻求二者在宏观分级与微观算子上的精简互补。核心 idea：构建浅层 CNN 局部特征抽取结合深层 Transformer 全局建模的宏观分层结构，并配合纯 RoPE 编码、多级重叠卷积与带残差的卷积前馈网络（ConvFFN），实现收敛快、外推稳、可线性扩展至 2.3B 的高效 VAE。

方法详解

整体框架

TransVAE 遵循对称的编码器-解码器架构，致力于在极简改动下提供兼具局部保真度与全局语义的完备视觉表征。对于输入图像 \(x \in \mathbb{R}^{H \times W \times 3}\)，编码器首先通过基于 ResBlock 的前置 CNN 阶段提取低层细粒度特征并逐级下采样，随后交由多阶段深层 TransVAE Block 提取高层全局语义，最后通过高斯建模与 KL 正则化映射为隐变量 \(z \in \mathbb{R}^{\frac{H}{f} \times \frac{W}{f} \times d}\)；解码器则采用严格对称的反向拓扑结构，先由逆向深层 TransVAE Block 进行长程依赖解码，再经由对称的 CNN 阶段完成逐级上采样并重构高保真像素图像 \(\hat{x}\)。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入图像 x<br/>H × W × 3"] --> B["浅层 CNN 前端<br/>连续重叠卷积与 ResBlock 逐级下采样"]
    B --> C["深层 TransVAE 骨干<br/>多阶段注意力建模全局语义"]
    C --> D["隐空间投影与采样<br/>高斯建模与 KL 正则化生成潜变量 z"]
    D --> E["深层逆向 TransVAE 骨干<br/>多阶段长程上下文逆映射"]
    E --> F["浅层 CNN 解码后端<br/>对称 ResBlock 逐级上采样与精细重建"]
    F --> G["重构输出 x_hat<br/>高保真图像重建"]

关键设计

1. 纯 RoPE 位置编码机制：彻底根除高分辨率外推时的先验失真与网格伪影

传统 ViT 分词器普遍采用加性可学习绝对位置编码（APE）\(X = X + P_{APE}\)，当测试输入分辨率超出训练尺寸（如 \(256 \times 256 \to 512 \times 512 / 1024 \times 1024\)）导致 token 序列拉长时，必须对参数矩阵 \(P_{APE}\) 进行网格插值，直接破坏了局部空间相对距离并造成严重的方块状畸变。TransVAE 完全移除 APE，将二维旋转位置编码（2D RoPE）引入自注意力机制中，通过正交旋转矩阵对查询向量 \(q_m\) 与键向量 \(k_n\) 进行位置绑定：

\[({q'}_m)^\top {k'}_n = g(q_m, m)^\top g(k_n, n) = h(q_m, k_n, m - n)\]

注意力分数显式依赖相对空间位移向量 \(m - n\) 而非绝对坐标索引。由于旋转变换对任意长度序列具有天然的相对距离平移保真性，TransVAE 无需任何位置参数插值即可在任意高分辨率下进行零样本泛化，同时极大加速了像素色彩细节的收敛。

2. 层次化多阶段 CNN-ViT 混合架构：构建局部细粒度与全局语义的互补分工

传统 ViT-VAE 采用单阶段大步长无重叠 Patch Embedding（如直接 \(16 \times 16\) 卷积），将相邻像素生硬割裂为孤立 token，破坏了高频连续性。TransVAE 提出宏观上的多阶段混合策略：在编码器前两个阶段完全采用标准卷积 Residual Block（ResBlock），并在级间通过 \(3 \times 3\) 步长为 2 的连续重叠卷积（SeqConv）进行下采样，使每个特征单元拥有重叠感受野并强化局部边界感知；随后的深层网络才接入 Transformer 模块负责广阔视场与全局拓扑组织，解码器则实施对称的逆级联结构。该设计让 CNN 专注于底层颜色与边缘细节抽取，卸载了 Transformer 学习局部平移归纳偏置的巨大计算负担。

3. 卷积前馈网络 ConvFFN 与局部微观增强：在深层注意力块中恢复空间连续性

标准 Transformer 内部的 FFN 为纯点式逐 token 投影，在深层切断了相邻 token 间的空间交互。TransVAE 在 FFN 中间引入带残差连接的 \(3 \times 3\) 全卷积旁路（FullConv）：在第一层全连接投影扩展通道后，特征通过空间 \(3 \times 3\) 卷积进行局部特征混合，再与点式线性变换结果相加并经过激活函数与输出投影。相较于深度可分离卷积（DWConv），FullConv 在保持紧凑参数量的前提下提供了更强的局部特征表达力，使深层网络兼具全局注意力与局部连续性。

4. 稳定扩展与无损跳跃连接：保障大模型训练平稳性与高倍率压缩保真度

为支撑网络规模从 44M 拓展至 2.3B Giant 级别，TransVAE 在每个注意力块的 \(Q, K, V\) 线性变换之前统一前置 RMSNorm 归一化层，避免自注意力梯度爆炸引发 NaN 发散；同时针对 \(f16\) 和 \(f32\) 等高空间压缩倍率，在上下采样级联中引入并行的无损特征直连分支（Shortcut），确保底层高频几何细节无损回传至深层解码器。

损失函数 / 训练策略

TransVAE 在重构训练阶段采用多任务联合损失函数，涵盖逐像素重构误差 \(L_1\)、感知损失 \(L_{\text{LPIPS}}\) 以及约束隐空间分布的 KL 散度损失 \(L_{\text{KL}}\)：

\[\mathcal{L}_{\text{total}} = \mathcal{L}_1(x, \hat{x}) + \lambda_{\text{lpips}} \mathcal{L}_{\text{LPIPS}}(x, \hat{x}) + \lambda_{\text{kl}} \mathcal{L}_{\text{KL}}(q_\phi(z|x) \parallel \mathcal{N}(0, I))\]

针对下游生成对齐模型（带 VF 后缀变体），进一步引入向量场对齐损失（Vector Field / VF Loss）拉近隐表征与语义先验的距离。训练全程采用 BFloat16 精度与 AdamW 优化器，在单分辨率（\(256 \times 256\)）下以批大小 32 完成训练。

实验关键数据

主实验

在 ImageNet-1k 验证集上跨不同分辨率评估模型的单分辨率训练直接外推能力（模型均仅在 \(256 \times 256\) 分辨率下训练）：

模型 (Tokenizer)	压缩规格	参数量	评估分辨率	rFID ↓	PSNR ↑	LPIPS ↓	SSIM ↑
DeTok-BB-FT	f16d16	171M	256 × 256	0.54	25.24	0.137	0.71
ViTok S-B/16	f16d16	129M	256 × 256	0.50	24.36	-	0.75
VA-VAE	f16d32	70M	256 × 256	0.27	28.57	0.090	0.80
TransVAE-T (Ours)	f16d32	44M	256 × 256	0.38	29.08	0.083	0.83
TransVAE-L (Ours)	f16d32	545M	256 × 256	0.31	29.78	0.075	0.84
DeTok-BB-FT	f16d16	171M	512 × 512	3.43	25.22	0.262	0.70
VA-VAE	f16d32	70M	512 × 512	0.12	31.05	0.095	0.84
TransVAE-T (Ours)	f16d32	44M	512 × 512	0.17	31.10	0.090	0.85
TransVAE-L (Ours)	f16d32	545M	512 × 512	0.13	32.44	0.077	0.88
DeTok-BB-FT	f16d16	171M	1024 × 1024	8.86	25.06	0.371	0.73
VA-VAE	f16d32	70M	1024 × 1024	0.06	36.15	0.071	0.93
TransVAE-T (Ours)	f16d32	44M	1024 × 1024	0.08	36.38	0.063	0.93
TransVAE-L (Ours)	f16d32	545M	1024 × 1024	0.03	38.25	0.051	0.95

消融实验

基于 5 个 Epoch 快速收敛设置与 3 个随机种子平均表现，分析位置编码、分层结构及 FFN 模块设计对重构损失与外推表现的影响：

配置模块	参数设定 / 变体	5 Epoch 验证重构损失 ↓	外推至 512 伪影现象	说明
位置编码机制 (PE)	无位置编码 (w/o PE)	~0.530	严重色彩失真	缺乏坐标参考，色彩与高频细节收敛极慢
位置编码机制 (PE)	APE (绝对位置编码)	~0.490	严重网格状伪影	插值破坏学习到的空间先验
位置编码机制 (PE)	纯 RoPE (Ours)	~0.370	无伪影且平滑保真	相对距离不变性完美支持任意分辨率外推
下采样与分级设计	单阶段 ViT + PixelShuffle	~0.480	存在边缘模糊	缺乏局部归纳偏置
下采样与分级设计	双端连续重叠卷积 (Dual SeqConv)	~0.340	边缘平滑	重叠感受野显著加速局部特征捕捉
下采样与分级设计	两级 CNN + 多级 ViT (Multi-Stage)	~0.250	高质量细粒度结构	前置 ResBlock 奠定局部基础，逼近纯 CNN 收敛速度
前馈网络设计 (FFN)	标准 MLP-FFN (37M)	~0.245	局部纹理略平滑	点式投影切断了邻域交互
前馈网络设计 (FFN)	DWConv-FFN (37M)	~0.210	纹理丰富	深度可分离卷积恢复局部归纳偏置
前馈网络设计 (FFN)	FullConv-FFN (42M, Ours)	~0.185	最佳细节保真	全卷积空间混合全面超越纯 CNN 收敛基线

关键发现

• 训练收敛速度倍增：TransVAE-T 仅需 25 个 Epoch 即达到 28.61 dB PSNR，超越了强大基线 VA-VAE 经 100 个 Epoch 训练的 28.57 dB，训练效率提升近 4 倍。
• 扩展性定律稳步成立：从 44M（Tiny）、140M（Base）、545M（Large）、1.3B（Huge）扩展至 2.3B（Giant），验证集损失与 PSNR 持续单调改善（3 Epoch 下 PSNR 从 27.2 dB 提升至 28.7 dB，50 Epoch 线性探测准确率从 10.91% 跃升至 65.18%），彻底打破了以往 ViT-VAE 在规模扩大时收益停滞的瓶颈。
• 语义对齐与重构的权衡：引入 VF 对齐损失虽使得 PSNR 产生约 0.5–1.0 dB 的微小回落，但大幅提升了下游扩散模型的生成质量，配合 LightningDiT-XL 在 80 个 Epoch 训练下将生成 FID 从 VA-VAE 的 4.51 / 1.96 压低至 3.58 / 1.85。

亮点与洞察

• 大道至简的极简主义设计：没有引入繁琐复杂的多分支融合结构，仅通过将浅层卷积、RoPE 和 ConvFFN 有机拼接，就完美化解了 ViT-VAE 长年存在的收敛慢、外推差、不扩展三大硬伤。
• 外推能力的归宿在于相对距离感知：验证了绝对位置编码插值是破坏高分辨率生成重建的罪魁祸首，纯 RoPE 能够零成本提供跨尺度无限外推能力。
• 可复用的通用生成特征抽取经验：深层 Transformer 配合局部卷积旁路的思路不仅适用于 VAE，同样适用于自回归图像建模和高分辨率连续视频潜空间的构建。

局限与展望

• 推理吞吐尚未极致优化：当前原型实现侧重于网络拓扑的简洁性与有效性，除引入 Flash Attention 外未对卷积与注意力的融合算子进行深度底层重排，推理吞吐略低于纯 CNN 高度优化的实现（如 44M TransVAE-T 编码吞吐量为 193 img/s，低于 70M VA-VAE 的 371 img/s）。
• 单分辨率对齐导致的外推语义衰减：在单一分辨率下进行 VF 语义对齐训练易造成高分辨率下的语义过拟合与外推性能下降，未来需要探索多分辨率联合对齐或更鲁棒的语义约束机制。

相关工作与启发

• vs VA-VAE: VA-VAE 依然采用传统的卷积加注意力设计，尽管在生成对齐上效果良好，但在更大分辨率与大模型扩展上面临瓶颈；TransVAE 参数可扩展至 2.3B 并具备更好的外推和重构上限。
• vs DeTok / ViTok: DeTok 因采用 APE 在分辨率外推时彻底崩溃；ViTok 虽支持 RoPE 外推但缺乏局部归纳偏置导致细节重建较差；TransVAE 兼顾了 RoPE 的外推通用性与 CNN 的局部高保真。
• vs DC-AE: DC-AE 在极深网络中主要依靠纯 CNN 并辅以末端微弱 ViT；TransVAE 则是真正的多级深度 Transformer 架构，赋予了隐表征更强大的全局语义与线性探测能力。

评分

• 新颖性: ⭐⭐⭐⭐ [针对 ViT-VAE 的缺陷进行了清晰的归因，提出了极简且高度有效的混合方案]
• 实验充分度: ⭐⭐⭐⭐⭐ [跨越 44M 到 2.3B 全量规模验证，涵盖重构、外推、线性探测及下游 DiT 生成]
• 写作质量: ⭐⭐⭐⭐⭐ [逻辑严密清晰，动机层层递进，消融对比与图表十分完备]
• 价值: ⭐⭐⭐⭐⭐ [为连续潜在扩散模型提供了极具实用价值的高效大容量分词器架构参考]

相关论文

• [CVPR 2026] SRA 2: Variational Autoencoder Self-Representation Alignment for Efficient Diffusion Training
• [ICLR 2026] Pareto Variational Autoencoder
• [ECCV 2024] Scalable Group Choreography via Variational Phase Manifold Learning
• [ECCV 2026] A Scalable Vector Graphics Latent Space
• [ICLR 2026] Latent Diffusion Model without Variational Autoencoder