AR-VAE

Adaptive-Length Autoregressive VAE for Image Compression

一个信息密度自适应的图像 tokenizer：将图像编码为变长的连续 latent token 序列。简单图像用少量 token，复杂图像用更多 token。整个系统在单张 RTX 4090 (24GB) 上可训练、可推理。

Status: trained on ImageNet-100 for 200k steps (~31h on RTX 4090). Final val LPIPS 0.23, L1 0.25, avg effective length 5.7 tokens. See Training Results for analysis.

---

核心设计

     image (B, 3, 224, 224)
           │
           ▼
  ┌──────────────────┐
  │  DINOv3 ViT-B/16 │   frozen, no_grad
  │   (patch feats)  │
  └────────┬─────────┘
           │ (B, 196, 768)
           ▼
  ┌──────────────────────────────┐
  │ Mixed-Attention Transformer  │   patch↔patch bidirectional
  │                              │   latent←patch one-way
  │                              │   latent↔latent causal
  └────────┬─────────────────────┘
           │ (B, 32, 768) latent_hidden
           ▼
  ┌───────────────┐   ┌────────────────────┐
  │   z = Wh      │   │  Monotonic Gate    │
  │               │   │  g_t = ∏ σ(α_i)    │
  └──────┬────────┘   └─────────┬──────────┘
         │                       │
         └────────┬──────────────┘
                  ▼
          z̃_t = g_t · z_t
                  │
                  ▼
       ┌────────────────────┐
       │ Transformer Decoder│   cross-attn to z̃
       │ + pixel head       │   L1 + LPIPS (VGG)
       └─────────┬──────────┘
                 ▼
         x̂ (B, 3, 224, 224)

五个关键设计决策：

决策	选择	原因
视觉 backbone	DINOv3 ViT-B/16, frozen	特征质量高，省显存，避免功能重叠
Latent 编码	自回归 (因果 attention)	与 soft gate 协同，保证截断后表示自洽
Gate 形式	单调累积 sigmoid	内嵌"前重要、后次要"先验，完全可微
Decoder	Transformer + L1 + LPIPS	确定性、训练快、4090 友好；不用 FM/diffusion
长度控制	拉格朗日对偶法	用直观的 L_target 替代手调 β

---

数学形式

Monotonic gate (单调保证来自累积乘积):

$$ g_t = \prod_{i=1}^{t} \sigma(\alpha_i), \quad \tilde{z}_t = g_t \cdot z_t $$

实际实现采用 log-space 累加以保证 bf16 下数值稳定：

gates = torch.exp(torch.cumsum(F.logsigmoid(alphas), dim=1))

Main loss:

$$ \mathcal{L} = \mathcal{L}_{\text{L1}} + \lambda_p \mathcal{L}_{\text{LPIPS}} + \gamma \mathcal{L}_{\text{matryoshka}} + \beta \cdot (\bar{L} - L_{\text{target}}) $$

Matryoshka auxiliary — 每个 step 随机采样一个前缀长度 $k \sim \text{Uniform}(1, T_{\max})$，额外做一次前缀重建，鼓励模型把重要信息放到前面：

$$ \mathcal{L}_{\text{matryoshka}} = \mathcal{L}_{\text{recon}}(x, \text{Decode}(\tilde{z}_{1:k})) $$

Lagrangian dual update (β 是对偶变量，不进 optimizer):

$$ \beta \leftarrow \text{clip}\bigl(\beta + \eta_\beta \cdot (\bar{L} - L_{\text{target}}),\ 0,\ \beta_{\max}\bigr) $$

这样不需要手调长度惩罚系数 — 只需要设定目标平均长度 L_target，β 会自动收敛。

Gate warm-up phase — 冷启动时 decoder 还没学会依赖 $z$，可以用位置 query 输出"均值图像"来最小化 L1 + LPIPS。此时 $\partial D / \partial g_t \approx 0$，任何微小长度惩罚都会把 gate 全部关死，Matryoshka 也失效（所有前缀产出相同的退化重建）。实验证实：不加 warm-up 训 100 步内 avg_length 从 10 跌到 0.02，distortion 和 distortion_matryoshka 完全相等。

修复方案：前 2000 步 bypass gate（z_tilde = z），只用 distortion + Matryoshka 训练，强制 decoder 依赖 $z$。结束后释放 gate，此时关闭 gate 会显著抬高 distortion，产生正确的反坍塌梯度。

---

安装

git clone https://github.com/<your-org>/AR-VAE.git
cd AR-VAE
pip install -r requirements.txt

依赖：

• torch >= 2.3, torchvision >= 0.18
• transformers >= 4.40 (DINOv3 weights via HuggingFace)
• lpips >= 0.1.4 (real perceptual loss, 训练必需)
• tensorboard >= 2.14, matplotlib >= 3.7
• einops, pyyaml, pytest

---

Quick Start

1. 跑测试 (不需要 GPU / 不需要下载权重)

python -m pytest tests/ -q

预期输出：21 passed — 包含 mask 因果性、gate 单调性、前缀截断等效性、β 动力学、inference 状态保持等全部 sanity check。

2. Smoke test (合成数据 + mock encoder, 不需要下载任何东西)

python -m ar_vae.train --config ar_vae/configs/default.yaml --mock-encoder --synthetic-data

这条命令用一个 conv 轻量替身代替 DINOv3, 并从合成随机图像中采样。用于验证整条训练管线可跑。

3. 真实训练 (Linux + RTX 4090)

准备 ImageNet-100 到 ./data/imagenet100/，然后：

python -m ar_vae.train --config ar_vae/configs/default.yaml

TensorBoard:

tensorboard --logdir ./outputs/tb

关键监控指标：

• metric/avg_length — 当前 batch 平均有效长度，应向 target_length=8.0 收敛
• metric/beta — 拉格朗日乘子，若一直顶到 beta_max 说明模型压不进 target
• metric/gate_std — gate 多样性，越大说明自适应能力越强
• gate_hist/* — 每个位置 $t$ 的 gate 分布直方图
• train/reconstructions — 每 500 步一个重建样本网格

4. 评估 (生成 CLAUDE.md §评估规范 的 4 张图)

python -m ar_vae.eval \
    --config ar_vae/configs/default.yaml \
    --checkpoint outputs/ckpt_final.pt \
    --output-dir outputs/eval

产出：

• quality_vs_length.png — L1 / PSNR / LPIPS 随 prefix 长度的曲线
• gate_distribution.png — 每个位置平均 $g_t$ ± 1σ 的柱状图
• gate_curves.png — 每张图的 gate 曲线叠加
• reconstruction_grid.png — 简单/复杂图像在不同前缀长度下的重建对比

---

项目结构

AR-VAE/
├── ar_vae/
│   ├── configs/default.yaml        # 所有超参数
│   ├── models/
│   │   ├── encoder.py              # DINOv3 ViT-B wrapper (frozen)
│   │   ├── mixed_attention.py      # 混合注意力 Transformer (核心)
│   │   ├── gate.py                 # 单调 soft gate (log-space)
│   │   ├── decoder.py              # Transformer pixel decoder
│   │   └── ar_vae.py               # 端到端模型
│   ├── losses/
│   │   ├── perceptual.py           # LPIPS wrapper (strict / fallback)
│   │   └── ar_vae_loss.py          # 总损失 + 拉格朗日 β 更新
│   ├── data/imagenet.py            # ImageNet-100 dataloader
│   ├── utils/
│   │   ├── checkpoint.py
│   │   ├── logging.py              # structured stdout logger
│   │   ├── tb_logger.py            # TensorBoard (scalars / histograms / images)
│   │   └── visualize.py            # matplotlib 评估图
│   ├── train.py                    # 训练入口
│   ├── eval.py                     # 评估 + 可视化
│   └── inference.py                # adaptive encode / decode prefix
├── tests/                          # 21 个 sanity check
├── requirements.txt
├── CLAUDE.md                       # 详细设计规范 (中文)
└── README.md

---

Training Results

Setup: ImageNet-100，batch size 32，bf16，RTX 4090（共享 GPU，另一进程占 ~5 GB），200k steps，~31.4 小时。配置见 ar_vae/configs/run_20260418.yaml。

Loss trajectory (train)

step	total	L1	LPIPS	distortion	distortion_matryoshka	avg_length	β
0	2.94	1.16	0.791	1.956	1.967	10.00	0.00
2000	2.20	0.858	0.607	1.465	1.464	12.06	0.00
5000	1.76	0.584	0.569	1.153	1.211	4.07	0.00
10000	1.53	0.454	0.536	0.990	1.073	7.56	0.00
20000	1.26	0.364	0.435	0.798	0.921	7.56	0.00
50000	0.944	0.287	0.315	0.602	0.683	7.75	0.00
100000	0.782	0.244	0.247	0.491	0.581	5.77	0.00
150000	0.707	0.222	0.224	0.446	0.521	5.61	0.00
199950	0.689	0.218	0.220	0.438	0.502	5.72	0.00

Validation (ImageNet-100 val)

step	L1	LPIPS	avg_length	gate_std
2000	0.900	0.603	12.20	3.39
10000	0.481	0.537	7.32	0.78
20000	0.382	0.441	7.40	0.45
50000	0.313	0.317	7.86	0.18
100000	0.271	0.255	5.81	0.098
150000	0.253	0.233	5.61	0.083

观察

工作良好的部分 ✅

• 冷启动坍塌被解决：加入 gate warm-up 后，avg_length 稳定收敛，全程未坍塌。warm-up 前训练试跑过两次，100 步内就跌到 0.02 并卡死，对比明显。
• Matryoshka 提供了有效的反坍塌信号：distortion_matryoshka > distortion 的差值稳定在 ~0.07（0.438 vs 0.502），说明截断确实"贵"，decoder 学会用前缀含义。若此 gap ≈ 0，就是 decoder 忽略 $z$ 的信号。
• 拉格朗日 β 正确表达不等式约束：整个训练中 β 始终 = 0，avg_length 自行收敛到 5.7（< target=8）。模型在 $\bar{L} \le 8$ 的可行域内自动选择了更紧凑的操作点 —— 这正是 Lagrangian 对 $\le$ 约束的预期行为。
• 前重要、后次要先验被学到：gate_pos/g_t 在训练早期就形成单调阶梯，前 ~6 个位置保持在 ~1，$g_{30}$ 接近 0。

已知局限 / 未解决 ⚠️

• Per-sample 自适应性在训练过程中衰减：gate_std（样本间有效长度标准差）从 step 2k 的 3.39 降到 150k 的 0.083。也就是说，最终所有图像被分配几乎相同的有效长度（~5.7 ± 0.08 个 token），并没有如设计目标那样简单图短、复杂图长。Matryoshka + Lagrangian 的组合保证了单调前缀结构，但没有显式奖励"为不同图像分配不同长度"。想要恢复自适应性，需要设计一个新机制（例如基于图像复杂度的长度目标、或在 Matryoshka 中用 per-sample $k$）。
• 重建模糊：最终 LPIPS ~0.22，结构和颜色准确但纹理细节丢失。L1 + LPIPS 的设计决策意味着无法达到 GAN-level 锐度，这是已知 trade-off（见 CLAUDE.md 的架构决策）。
• Train/val 有 ~10% gap：L1 0.218 vs 0.253，LPIPS 0.220 vs 0.233，轻度过拟合 ImageNet-100，符合 100 类 × ~1300 张的数据量规模。

可复现

python -m ar_vae.train --config ar_vae/configs/run_20260418.yaml

---

默认超参数

类别	参数	值
Model	latent_dim / ffn_dim	768 / 3072
Model	num_encoder_layers / num_decoder_layers	8 / 8
Model	max_latent_len	32
Model	gate_init_bias	2.0 (σ(2) ≈ 0.88)
Loss	target_length	8.0
Loss	lpips_net	vgg
Loss	matryoshka_weight	0.5
Loss	beta_lr / beta_max	0.01 / 10.0
Loss	gate_warmup_steps	2000 (bypass gate until decoder depends on $z$)
Optim	lr / weight_decay	1e-4 / 0.05
Optim	warmup_steps / total_steps	2000 / 200000
Optim	grad_clip	1.0
Train	batch_size / precision	32 / bf16

显存不够时按顺序启用：

1. batch_size 32 → 16 → 8
2. gradient_checkpointing: true
3. max_latent_len 32 → 16
4. num_encoder_layers / num_decoder_layers 8 → 6

---

测试套件 (21 项)

文件	覆盖内容
test_mask.py	混合注意力 mask 的因果性 (扰动 $z_{t+1}$ 不应影响 $z_t$)
test_gate.py	Gate 单调性 + log-space 等价性
test_model_end_to_end.py	Forward shape、梯度路由 (encoder 无梯度)、前缀截断等效
test_loss.py	Loss 前反向、β 动力学 (升 / 降 / eval 冻结)
test_inference.py	Adaptive encode、边界情况、训练状态保持
test_tb_logger.py	TensorBoard roundtrip smoke
test_visualize.py	matplotlib 4 张图生成 smoke

运行：python -m pytest tests/ -q

---

推理 API

from ar_vae.models import ARVAE
from ar_vae.models.ar_vae import ARVAEConfig
from ar_vae.inference import encode_adaptive, decode_prefix

model = ARVAE(ARVAEConfig()).eval()

# 自适应编码：根据 gate 阈值自动决定有效长度
z_tilde, gates, t_star = encode_adaptive(model, x, gate_threshold=0.05)
# z_tilde: (B, Tmax, D), gates: (B, Tmax), t_star: (B,)

# 指定前缀长度解码
x_hat = decode_prefix(model, z_tilde, prefix_len=8)

训练和推理的计算图相同 — 因为 latent 间的因果 mask 已经保证了依赖顺序，不需要真的逐步生成，没有 exposure bias。这是这个架构相比传统 AR 的一个实际优势。

---

设计要点 / 不做的事

不做:

• 不要训练 DINOv3 (冻结)
• 不要把 decoder 换成 Flow Matching / Diffusion (违背 4090 友好的决策)
• 不要把 gate 改成离散 stop token (违背端到端可微)
• 不要引入 KL 项 (会引入 KL collapse)
• 不要在 decoder 里重复 mask latent (gate 已经做了 soft mask)

容易踩的坑:

• PyTorch SDPA 的 attn_mask 中 True = 允许 attend，不是屏蔽 (和很多教程相反)
• DINOv3 的 last_hidden_state 包含 CLS + 4 个 register token，要切掉前 5 个
• LPIPS 输入范围是 [-1, 1]，ImageNet normalize 之后范围不同，需要反归一化
• β 不能参与反向，必须在 no_grad() 里 in-place 更新
• Matryoshka 的 $k$ 每步重采样，不要在 epoch 内固定

更多细节见 CLAUDE.md。

---

Acknowledgments

• DINOv3 — Meta AI (facebook/dinov3-vitb16-pretrain-lvd1689m)
• LPIPS — Zhang et al., The Unreasonable Effectiveness of Deep Features as a Perceptual Metric
• 混合注意力 mask 的设计灵感来自 DeepSeek-OCR V2 的 DeepEncoder

---

License

TBD.