| # MAVT (AnTokenizer) — Model & Pipeline |
|
|
| > Unified vision tokenizer cho image / video / 3D, học chung 1 latent space. |
| > Train pipeline 3-stage curriculum (image → +video → +3D). |
| > Code chính: `src/mavt/model/` + `src/mavt/training/lightning_module.py`. |
| |
| --- |
| |
| ## 1. Pipeline tổng quan (7 stages) |
| |
| ``` |
| ┌─────────────────┐ |
| raw input ──▶ │ 1. Patchify │ Conv3d, modality-specific |
| │ (multi-modal) │ → tokens (B,N,D), positions (N,4), plane_ids (N,) |
| └────────┬────────┘ |
| │ |
| ┌────────▼────────┐ |
| │ 2. Hybrid │ Transformer + RGAT, 12 blocks, dim=768 |
| │ Backbone │ → features (B,N,D) |
| └────────┬────────┘ |
| │ |
| ┌────────▼────────┐ |
| │ 3. C-D Split │ GLOBAL content (slot attn) + LOCAL detail (window pool) |
| │ (cd-split) │ → compressed (B, N_c+N_d, D), latent_positions, token_types |
| └────────┬────────┘ |
| │ |
| ┌────────▼────────┐ |
| │ 4. VAE bottleneck│ N×D → N×latent_dim (32), reparametrize z |
| │ + KL loss │ → z, mu, logvar, loss_kl |
| └────────┬────────┘ |
| │ |
| ┌───────────┴───────────┐ |
| ┌──────▼──────┐ ┌──────▼──────┐ |
| │ 5a. Recon │ │ 5b. Underst.│ |
| │ Decoder │ │ Decoder │ |
| │ (z → pixel) │ │ (z → semantic, distill từ SigLIP2) |
| └──────┬──────┘ └──────┬──────┘ |
| │ │ |
| ▼ ▼ |
| recon (B,3,H,W) semantic (B, 768) |
| │ │ |
| └─────────┬──────────────┘ |
| │ |
| ┌─────────▼─────────┐ |
| │ 6. Loss (MAVTLoss)│ l1 + lpips + KL + sem distill + temporal (video) + slot div |
| │ per-modality EMA│ proportional weighting |
| └─────────┬─────────┘ |
| │ |
| ▼ |
| 7. Outputs (`MAVTOutput` dataclass) |
| ``` |
| |
| --- |
|
|
| ## 2. Stage chi tiết |
|
|
| ### Stage 1 — Patchify (`patchify.py`) |
| - **Image** `(B, 3, H, W)`: Conv3d với causal-pad 1 frame ảo → `(B, D, 1, Hp, Wp)` → flatten `(B, Hp·Wp, D)`. Position `(0, i, j, 0)`, plane_id=−1. |
| - **Video** `(B, 3, T, H, W)`: Conv3d giảm temporal `t_patch=2`, spatial `patch=16` → `(B, D, Tp, Hp, Wp)` → `(B, Tp·Hp·Wp, D)`. Position `(t, i, j, 0)`, plane_id=−1. |
| - **Threed** `(B, 3, 3, S, S)` (3 planes XY/XZ/YZ): Conv3d riêng cho mỗi plane → concat `(B, 3·Hp·Wp, D)`. Position encode plane-specific axes; plane_id ∈ {0,1,2}. |
| - Tất cả share Conv3d weight → unified patch embedding qua modality. |
| - Kèm 4D `pos_embed` (Fourier features, dim D) cộng vào tokens. |
|
|
| ### Stage 2 — Hybrid Backbone (`backbone.py`) |
| - 12 layers xen kẽ Transformer blocks + RGAT (Relational Graph Attention). |
| - RGAT là attention dạng graph với: |
| - `r_s=2` neighborhoods spatial (window-based) |
| - `r_t=1` neighborhood temporal |
| - `use_gradient_checkpointing=true`: tradeoff compute/memory (đang OFF cho run tuned để tăng tốc). |
| - SigLIP2 weights được load vào last 4 transformer blocks (Stage 1: frozen, Stage 2: unfrozen). |
|
|
| ### Stage 3 — Content-Detail Split ⭐ **(deep dive section 3)** |
|
|
| ### Stage 4 — VAE bottleneck (`latent_heads.py`) |
| - `VAEHead`: Linear(D=768 → 2·latent_dim=64) → split mu, logvar → reparametrize trick. |
| - KL loss: `KL(N(mu, σ²) || N(0, I))`, scale bằng `kl_weight=1e-4` (built-in). |
| - Output `z` shape `(B, N_c+N_d_local, latent_dim=32)`. |
| - Note: post `42c622b update`, `w_kl=1.0` ở MAVTLoss = passthrough — đã pre-scale 1e-4 trong VAEHead → tránh double-scale. |
|
|
| ### Stage 5 — Decoders (`decoder.py`) |
|
|
| **5a. Reconstruction `AsymmetricDecoder`:** |
| - `UnifiedDetailExpander` cross-attention từ target grid positions vào latent z (dim 32). |
| - 4 self-attention blocks dim 768. |
| - Pixel projection → `(B, 3, H, W)` cho image, frame-by-frame cho video. |
| - **Mới (cd-split)**: nhận `latent_positions` + `latent_token_types` để áp distance-bias attention (xem section 3). |
|
|
| **5b. Understanding `UnderstandingDecoder`:** |
| - 2 cross-attn layers + linear proj → `(B, semantic_dim=768)`. |
| - Trained để khớp với SigLIP2 teacher's `pooler_output` qua cosine loss. |
|
|
| ### Stage 6 — Loss (`losses.py`) |
|
|
| ``` |
| L_total = w_l1 · L1(pred, target) |
| + w_lpips · LPIPS(pred, target) # AlexNet/VGG perceptual |
| + w_kl · L_KL # đã pre-scaled ở VAEHead |
| + w_sem · (1 - cos(MAVT.semantic, teacher.pooler)) |
| + w_temp · L1(Δ_t pred, Δ_t target) # chỉ video, T>1 |
| + w_aux · slot_diversity_penalty |
| ``` |
| - Mỗi modality scale bằng `ModalityEMAWeighter.weight(modality)`: |
| - `weight(m) = ema_m / mean(ema_active)` — modality có loss CAO được boost (cross-stage gradient flow vào branch chưa train) |
| - Sau commit `42c622b update`. Trước đó là 1/ema (logic ngược). |
|
|
| --- |
|
|
| ## 3. ⭐ C-D Split (sau commit `6368dfb cd-split`) |
|
|
| ### Ý tưởng cốt lõi |
|
|
| Phân chia input tokens thành 2 kênh có **đặc tính khác nhau**, encode bằng **cơ chế khác nhau**: |
|
|
| | Kênh | Bản chất | Cơ chế | Position info | |
| |---|---|---|---| |
| | **Content** | Semantic / low-freq / global | Slot cross-attention (toàn ảnh) | Không có (slot là global summary) | |
| | **Detail** | Residual / high-freq / local | Coordinate window pooling | **Có** (window center) | |
|
|
| ### Tại sao Detail cần local + position? |
|
|
| **Trước cd-split** (slot pooler global cho cả detail): |
| ``` |
| detail = SlotPooler(N_d=25 slots)(Residual) |
| (slots tự học pool ở đâu, không có vị trí) |
| ``` |
| - Decoder cross-attend vào detail slots không biết slot này ứng với patch nào → phải reconstruct texture từ "positionless global slots" → khó. |
| - High-freq (texture, edges) cần spatial precision → mất khi pool global. |
|
|
| **Sau cd-split** (windowed pool + position): |
| ```python |
| # Group residual tokens theo coordinate window |
| group_key = (plane_id, t // t_win, x // s_win, y // s_win, z // s_win) |
| # Mean-pool tokens trong cùng window |
| detail_token[g] = mean(residual[token] for token in window g) |
| detail_position[g] = mean(positions[token] for token in window g) + 0.5 # window center |
| ``` |
| - **Mỗi detail token có toạ độ rõ ràng** → decoder biết detail thuộc patch nào. |
| - Decoder dùng **distance bias** (Manhattan) trong cross-attn để mỗi pixel ưu tiên detail token gần. |
| - Compression vẫn tốt: window 2×2 → 4 token → 1 token (75% giảm), tổng compression vẫn ~50% (bằng N_c + N_d_local). |
| |
| ### Architecture sau update |
| |
| ``` |
| ┌─ slot attn ──▶ C (B, N_c, D) [global, positionless] |
| features ──┬─▶ ─────┤ |
| (B,N,D) │ └─ approx via inverse softmax weights: |
| │ x_approx = softmax(C @ xᵀ / √D)ᵀ @ C |
| │ |
| └─▶ R = x - x_approx (residual) |
| │ |
| ▼ |
| ┌──────────────────────────────────────┐ |
| │ _local_detail_pool(R, pos, plane) │ |
| │ │ |
| │ group_key = (plane, t/1, i/2, j/2, k/2) |
| │ D_tokens[g] = mean(R[t] for t in g) │ |
| │ D_pos[g] = mean(pos[t]) + 0.5 │ |
| │ D_tokens ← detail_proj(detail_norm(.)) |
| └──────────┬───────────────────────────┘ |
| │ |
| ▼ |
| detail tokens (B, N_d_local, D) + detail_positions (N_d_local, 4) |
| |
| compressed = concat([C, detail_tokens]) (B, N_c + N_d_local, D) |
| latent_positions = concat([zeros(N_c, 4), detail_positions]) |
| latent_token_types = concat([zeros(N_c), ones(N_d_local)]) |
| ``` |
| |
| ### Decoder sử dụng metadata thế nào? |
| |
| ```python |
| # UnifiedDetailExpander forward |
| kv = z + kv_pos_scale * kv_pos_enc(latent_positions) # add 4D Fourier pos |
| + token_type_scale * token_type_embed(latent_token_types) # +0/+1 embedding |
| |
| # Distance bias chỉ apply cho detail keys |
| dist_manhattan = |query_pos - kv_pos| |
| attn_bias[detail_keys] = -local_detail_bias * dist # local_detail_bias=0.25 |
|
|
| cross_attn(query, kv, kv, attn_mask=attn_bias) |
| ``` |
| → Pixel position xa detail position thì attention bị penalize logarithmically (softmax-scale). |
| → Content tokens KHÔNG bị penalize → decoder vẫn dùng được toàn bộ semantic info. |
| |
| ### Worked example: image 256×256 |
| |
| ``` |
| Input: x shape = (B, 3, 256, 256) |
| Patch_size=16 → Hp=Wp=16 → N=256 tokens |
| positions = [(0,i,j,0) for i,j in 16×16] → (256, 4) |
| plane_ids = [-1] * 256 |
| |
| content_ratio=0.25, detail_ratio=0.25 → N_c=64, N_d_key=64 (key naming only) |
| ``` |
| |
| **Stage 3a: Content slot pool** |
| ``` |
| slot_pooler = SlotPooler(num_slots=64, dim=768, num_heads=8, num_layers=2) |
| C = slot_pooler(features) # 2 cross-attn layers |
| shape (B, 64, 768) |
| ``` |
| - 64 learned slots cross-attend toàn 256 tokens → mỗi slot là weighted summary toàn ảnh. |
| |
| **Stage 3b: Approximate + residual** |
| ``` |
| weights = softmax(C @ features.T / sqrt(768), dim=-1) # (B, 64, 256) |
| x_approx = weights.T @ C # (B, 256, 768) |
| R = features - x_approx # (B, 256, 768) — high-freq |
| ``` |
| |
| **Stage 3c: Local detail pool (window=2)** |
| ``` |
| group_key[token_n] = (plane_id=-1, t=0, i//2, j//2, z=0) |
| = (-1, 0, i//2, j//2, 0) |
| |
| i=0,j=0 → key (-1,0,0,0,0) group 0 |
| i=0,j=1 → key (-1,0,0,0,0) group 0 (same window 2×2) |
| i=0,j=2 → key (-1,0,0,1,0) group 1 |
| i=0,j=3 → key (-1,0,0,1,0) group 1 |
| ... |
| i=1,j=0 → key (-1,0,0,0,0) group 0 |
| i=1,j=1 → key (-1,0,0,0,0) group 0 |
| ... |
| ``` |
| → 4 token (vd i=0..1, j=0..1) gộp vào group 0. |
| → Tổng số group = 8×8 = **64 detail tokens**. |
| |
| ``` |
| counts[0] = 4 (i=0,1; j=0,1) |
| D_token[0] = mean(R[0], R[1], R[16], R[17]) # 4 token trong window 2×2 |
| D_token[0] = detail_proj(detail_norm(D_token[0])) |
| |
| D_pos[0] = mean([(0,0,0,0), (0,0,1,0), (0,1,0,0), (0,1,1,0)]) + 0.5 = (0, 1, 1, 0) |
| → window center floor |
| ``` |
| |
| **Output**: |
| - `compressed` = concat(C, D_tokens) shape `(B, 128, 768)` |
| - `latent_positions` shape `(128, 4)`: |
| - First 64 rows: `(0,0,0,0)` (content, positionless) |
| - Last 64 rows: window centers like `(0, 1, 1, 0)`, `(0, 1, 3, 0)`, ... |
| - `latent_token_types` shape `(128,)`: `[0]*64 + [1]*64` |
|
|
| ### Worked example: video 256×256, 16 frames |
|
|
| ``` |
| T=16, t_patch=2 → Tp=8 |
| N = 8 × 16 × 16 = 2048 tokens |
| |
| content_ratio=0.25 → N_c=512 |
| detail_ratio=0.25 → N_d_key=512 (naming) |
| |
| Detail windows (s_win=2, t_win=1): |
| group_key = (plane=-1, t//1=t, i//2, j//2, 0) |
| t in [0..7]: 8 unique |
| i//2 in [0..7]: 8 unique |
| j//2 in [0..7]: 8 unique |
| → 8 × 8 × 8 = 512 detail windows |
| ``` |
| - Compressed: 512 + 512 = **1024 tokens** (vs 2048 raw → 2× compression) |
| - Mỗi detail token gồm 1 temporal × 4 spatial residual (tổng 4 raw tokens). |
|
|
| ### Compare số token: image 256² (3 modality khác nhau) |
|
|
| | Modality | N raw | N_c (content) | N_d_local (detail, win=2) | Total | Compression | |
| |---|---:|---:|---:|---:|---:| |
| | Image | 256 | 64 | **64** | **128** | 2× | |
| | Video | 2048 | 512 | **512** | **1024** | 2× | |
| | Threed | 768 | 268 | **192** | **460** | 1.67× | |
| |
| ### Hyperparameters (configurable qua CLI hoặc yaml) |
| |
| | Param | Default | Tác động | |
| |---|---:|---| |
| | `local_detail_window_size` | 1 (sau user update) | Kích thước window spatial. 1 = không pool (mỗi token 1 group), 2 = 2×2 windows | |
| | `local_detail_temporal_window_size` | 1 | Window temporal. 1 = mỗi frame riêng | |
| | `content_ratio` (modality-specific) | 0.25 (img/vid), 0.35 (3D) | N_c = N × ratio | |
| | `detail_ratio` (key naming only) | 0.25 | Không ảnh hưởng số detail token thực | |
| | `local_detail_bias` (decoder) | 0.25 | Hệ số distance bias trong cross-attn. Lớn = ép detail mạnh hơn | |
| | `kv_pos_scale` (decoder, learnable) | init 0.1 | Trọng số position encoding cộng vào KV | |
| | `token_type_scale` (decoder, learnable) | init 0.1 | Trọng số token type embed cộng vào KV | |
|
|
| ### Monitoring metrics |
|
|
| | Metric | Ý nghĩa | Target | |
| |---|---|---| |
| | `slot_diversity` | mean pairwise cos sim giữa các content slots | ≤ 0.5 (slots khác nhau) | |
| | `residual_ratio` | `‖R‖ / ‖x‖` | 0.3–0.5 (content giữ phần lớn signal) | |
| | `detail_token_count` | số detail token thực | = số window distinct | |
| | `detail_avg_window_tokens` | trung bình tokens per window | ≈ s_win² · t_win nếu density đều | |
|
|
| ### Tại sao đổi từ slot pooler global → window pool cho detail? |
|
|
| | Khía cạnh | Slot pooler (cũ) | Window pool (mới) | |
| |---|---|---| |
| | Position info | ❌ (global) | ✅ (window center) | |
| | Compression | tốt (25 slots cho 256 token) | tốt (64 windows cho 256 token với win=2) | |
| | High-freq detail | hạn chế (slot abstract) | tốt (mean trong window nhỏ giữ texture) | |
| | Inductive bias | không có spatial prior | có (locality assumption) | |
| | Tham số | trainable slot params + 2 layer cross-attn | non-trainable scatter_add + 1 LayerNorm + 1 Linear | |
| | Compute | cao (slot attn O(N·N_d)) | thấp (scatter O(N)) | |
|
|
| → Chuyển sang windowed pool cho detail là **lossless về expressive power** với inductive bias hợp lý cho high-freq, lại **rẻ hơn** về tham số/compute. |
|
|
| --- |
|
|
| ## 4. Curriculum 3 stage |
|
|
| | Stage | Modalities | SigLIP2 unfreeze | LR | Purpose | |
| |---|---|---|---:|---| |
| | 1 | image only | hoàn toàn frozen | 1e-4 | Học image tokenizer + distill semantic | |
| | 2 | image + video | last 4 blocks unfrozen | 5e-5 | Thêm video poolers, fine-tune backbone cho temporal | |
| | 3 | + threed | toàn bộ unfrozen | 2e-5 | Thêm 3D, polish toàn bộ | |
|
|
| Cross-stage transfer: `--model.init_from_ckpt <prev_stage_ckpt>` (strict=False, weights only). Lightning module `setup('fit')`: |
| 1. `_prepare_cd_split_poolers()` — eagerly tạo content pooler cho mỗi modality active (bắt buộc trước `configure_optimizers`) |
| 2. `_sync_ema_modalities()` — đồng bộ active_modalities từ DataModule vào EMA weighter |
| 3. `load_siglip2_weights()` — load HF weights nếu `init_siglip2=true` |
| 4. `_load_semantic_teacher()` — load frozen SigLIP2 vision tower nếu `use_semantic_distill=true` |
| 5. `_load_weights_from_ckpt(init_from_ckpt)` — load prev-stage weights cuối cùng để override init |
|
|
| --- |
|
|
| ## 5. Latent space chính thức |
|
|
| Sau VAE bottleneck: |
| - Image: 128 token × 32 dim = **4096 floats** (vs raw 196,608 → 48× compression) |
| - Video: 1024 token × 32 dim = **32,768 floats** (vs raw 3,145,728 → 96× compression) |
| - Threed: 460 token × 32 dim = **14,720 floats** (vs raw 589,824 → 40× compression) |
|
|
| So với baselines (theo `results.md`): |
| - AToken-So/C Stage 1: (1, 16, 16) = 256 token × 32 ch = 8192 floats / image (chúng tôi 4096 — gấp đôi compression nhờ slot ratio 0.25) |
| - Cosmos-CI16×16: 256 × 16 = 4096 (same as ours per token volume, nhưng arch khác) |
|
|
| --- |
|
|
| ## 6. Reference files |
|
|
| | File | Content | |
| |---|---| |
| | `src/mavt/model/patchify.py` | Stage 1 | |
| | `src/mavt/model/backbone.py` | Stage 2 (Transformer + RGAT) | |
| | `src/mavt/model/content_detail_split.py` | Stage 3 ⭐ | |
| | `src/mavt/model/latent_heads.py` | Stage 4 (VAEHead) | |
| | `src/mavt/model/decoder.py` | Stage 5 (AsymmetricDecoder, UnderstandingDecoder, UnifiedDetailExpander) | |
| | `src/mavt/losses/losses.py` | Stage 6 (MAVTLoss, ModalityEMAWeighter, temporal_consistency_loss) | |
| | `src/mavt/training/lightning_module.py` | Curriculum, optimizer, logging | |
| | `src/mavt/model/mavt.py` | End-to-end MAVT module | |
| | `configs/model/mavt_base.yaml` | Hyperparams arch | |
| | `configs/train/universal_data/stage{1,2,3}_universal.yaml` | Stage curriculum | |
|
|
