Update README with adapted component stack (DINOv3, Qwen3, SAM3.1, etc.)

README.md

@@ -1,11 +1,11 @@
 ---
-title: ml-intern sandbox
-emoji: 🌍
-colorFrom: gray
-colorTo: blue
-sdk: docker
-app_port: 7860
-pinned: false
+tags:
+  - multimodal
+  - reasoning
+  - jepa
+  - world-model
+  - vision-language
+license: apache-2.0
 ---
 
 # MR-JEPA: Multimodal Reasoning via Joint-Embedding Predictive Architecture
@@ -22,23 +22,184 @@ z₀ (initial evidence) → z₁ (first refinement) → z₂ (deeper reasoning)
 
 This trajectory is supervised by a **JEPA objective**: a target encoder (EMA) generates target latent states, and the online predictor learns to predict them. The JEPA loss encourages the model to learn **meaningful intermediate reasoning states** — not just the final answer.
 
+---
+
 ## Architecture
 
 ```
-┌─────────────┐     ┌──────────────┐     ┌─────────────────┐     ┌──────────┐
-│  DINOv2/v3  │────▶│   Evidence   │────▶│  Latent Rollout │────▶│  Answer  │
-│  (frozen)   │     │   Memory     │     │  z₀→z₁→z₂→z₃   │     │  Heads   │
-└─────────────┘     │  (Perceiver) │     │  (shared block)  │     └──────────┘
-                    └──────┬───────┘     └────────┬────────┘
-┌─────────────┐           │                      │
-│  DeBERTa-v3 │───────────┘              ┌───────┴────────┐
-│  (frozen)   │                          │ Target Encoder  │
-└─────────────┘                          │  (EMA copy)     │
-                                         └────────────────┘
-┌─────────────┐                                 │
-│ OCR/Layout/ │──────────┘               JEPA Loss: L₂ + SIGReg
-│ Chart/SAM   │ (Phase 3)
-└─────────────┘
+┌──────────────┐     ┌────────────────────┐     ┌──────────────────┐     ┌───────────────┐
+│  DINOv3-L/16 │────▶│   Evidence Memory  │────▶│  Latent Rollout  │────▶│ Disc. Head    │
+│   (frozen)   │     │ (Perceiver Resampl)│     │  z₀→z₁→z₂→z₃    │     │ (MC scoring)  │
+└──────────────┘     └────────┬───────────┘     │ (shared block)   │     └───────────────┘
+                              │                 └────────┬─────────┘     ┌───────────────┐
+┌──────────────┐              │                          │          ├───▶│ Gen. Decoder   │
+│ Qwen3-Embed  │──────────────┘                  ┌───────┴────────┐     │ (Qwen3.5-4B)  │
+│   0.6B       │                                 │ Target Encoder │     └───────────────┘
+│   (frozen)   │                                 │  (EMA copy)    │
+└──────────────┘                                 └────────────────┘
+                                                         │
+┌──────────────┐                                  JEPA Loss:
+│ Phase 3 opt: │──────────┘                       SmoothL1/Cosine
+│ PaddleOCR-VL │                                  + SIGReg (purist)
+│ SAM 3.1      │                                  / VICReg (hybrid)
+└──────────────┘
+```
+
+---
+
+## Component Stack
+
+| Module | Primary Choice | Alternative | Notes |
+|--------|---------------|-------------|-------|
+| **Visual backbone** | [`timm/vit_large_patch16_dinov3.lvd1689m`](https://hf.co/timm/vit_large_patch16_dinov3.lvd1689m) — DINOv3-L/16, 1024-dim, 300M | DINOv3-B/16 (purist); DINOv2-L/14 (ablation) | Frozen Phase 1; last 6 layers unfrozen Phase 2 |
+| **Text encoder** | [`Qwen/Qwen3-Embedding-0.6B`](https://hf.co/Qwen/Qwen3-Embedding-0.6B) — 1024-dim, 596M | Qwen3-Embedding-4B (heavier); EmbeddingGemma-300M (lighter) | Frozen Phase 1; last 4 layers unfrozen Phase 2 |
+| **Evidence memory** | Perceiver Resampler, 64 queries, 4 cross-attn layers | Q-Former as baseline | Modality-typed tokens (visual/text/OCR/layout/chart/SAM) |
+| **OCR / doc / charts** | [`PaddlePaddle/PaddleOCR-VL-1.5`](https://hf.co/PaddlePaddle/PaddleOCR-VL-1.5) — 958M | MinerU2.5 for heavy PDF parsing | Phase 3 only, offline token extraction |
+| **Segmentation** | [`facebook/sam3.1`](https://hf.co/facebook/sam3.1) — SAM 3.1, 3.3GB, gated | SAM 2.1-Large (stable) | Phase 3 optional, offline mask extraction |
+| **Latent rollout** | Shared transformer predictor, 6 layers, K=3 | Per-step unshared blocks (ablation) | Weight-tied across steps; sigmoid evidence gates |
+| **Target encoder** | EMA copy (cosine 0.996→1.0) of evidence+rollout | Frozen target (ablation baseline) | From I-JEPA |
+| **JEPA loss** | SmoothL1 + VICReg (hybrid); Cosine + SIGReg (purist) | MSE (ablation) | SIGReg emphasis in purist branch |
+| **Disc. head** | MLP/bilinear scorer | Cross-encoder scorer (ablation) | Attention-pooled z_K × option embeddings |
+| **Gen. decoder** | [`Qwen/Qwen3.5-4B`](https://hf.co/Qwen/Qwen3.5-4B) — 4.7B, multimodal | [`HuggingFaceTB/SmolLM3-3B`](https://hf.co/HuggingFaceTB/SmolLM3-3B) (cheaper); Gemma3-4B | Phase 3+, cross-attends to z_K + evidence |
+| **Teacher/baseline** | InternVL3.5 / Qwen3-VL | External comparison only | NOT used as internal module |
+
+---
+
+## Training Protocol
+
+### Phase 1: Reasoning Core (15–20 epochs)
+- **Freeze** all perception (DINOv3 + Qwen3-Embedding)
+- **Train** evidence memory + latent rollout + discriminative head
+- Full JEPA loss + task loss
+- LR: 3e-4, effective batch: 64
+
+### Phase 2: Perception Fine-tuning (10 epochs)
+- **Unfreeze** last 6 DINOv3 layers + last 4 Qwen3-Embedding layers (1e-5)
+- Continue training reasoning core (1e-4)
+
+### Phase 3: Enriched Evidence + Generative Decoder (10 epochs)
+- **Enable** PaddleOCR-VL tokens, SAM 3.1 masks, layout/chart tokens
+- **Attach** Qwen3.5-4B generative decoder for open-ended answers
+- End-to-end fine-tuning, LR: 5e-5
+
+---
+
+## Target Benchmarks (9)
+
+| Benchmark | Type | Metric | Key Challenge |
+|-----------|------|--------|---------------|
+| MMMU | MC (multi-image) | Accuracy | Multi-discipline, up to 7 images |
+| MathVista | Mixed MC/Open | Accuracy | Mathematical reasoning |
+| ScienceQA | MC | Accuracy | Scientific diagrams, nullable images |
+| AI2D | MC | Accuracy | Science diagram comprehension |
+| MMBench | MC | CircularEval Acc | General visual understanding |
+| MMStar | MC | Accuracy | Vision-dependent questions |
+| DocVQA | Open | ANLS | Document text extraction |
+| TextVQA | Open | VQA Accuracy | Scene text reading |
+| ChartQA | Open | Relaxed Accuracy | Chart data extraction |
+
+---
+
+## Experimental Branches
+
+### Hybrid-main (competitive)
+- DINOv3-L backbone, SmoothL1 + VICReg, K=3
+- Full enriched evidence in Phase 3
+- Target: state-of-the-art on all benchmarks
+
+### Purist-side (scientific validation)
+- DINOv3-B backbone, Cosine + SIGReg, K=5
+- No enriched evidence, pure JEPA reasoning
+- Target: demonstrate JEPA contributes beyond perception
+
+---
+
+## Ablation Experiments
+
+| Experiment | Modification | Purpose |
+|------------|-------------|---------|
+| `hybrid_main` | Full model | Baseline |
+| `no_jepa` | Remove L_JEPA, task loss only | Validate JEPA objective |
+| `no_rollout` | K=0, use z₀ directly | Validate iterative refinement |
+| `no_gate` | Remove evidence gating | Validate adaptive evidence flow |
+| `K1` / `K5` / `K7` | Vary rollout depth | Find optimal depth |
+| `dinov2_ablation` | DINOv2-L/14 backbone | DINOv3 vs DINOv2 |
+| `purist` | DINOv3-B, no enriched ev., SIGReg | Isolate JEPA contribution |
+| `mse_loss` / `cosine_loss` | Alternative JEPA losses | Loss function ablation |
+
+---
+
+## Project Structure
+
+```
+MR-JEPA/
+├── README.md                    # This file
+├── train_mrjepa.py              # Complete training script (CLI, all ablations)
+├── test_architecture.py         # Architecture validation tests (synthetic data)
+│
+├── mr_jepa/
+│   ├── __init__.py
+│   ├── ARCHITECTURE.md          # Detailed architecture specification
+│   │
+│   ├── configs/
+│   │   ├── __init__.py
+│   │   └── model_config.py      # All hyperparameter dataclasses
+│   │
+│   ├── models/
+│   │   ├── __init__.py
+│   │   ├── mr_jepa.py           # Main model (integrates all components)
+│   │   ├── backbones.py         # Visual (DINOv3/v2) + Text (Qwen3-Embedding)
+│   │   ├── evidence_memory.py   # Perceiver Resampler multimodal fusion
+│   │   ├── latent_rollout.py    # K-step shared predictor + evidence gates
+│   │   ├── target_encoder.py    # EMA encoder + JEPA/SIGReg/VICReg losses
+│   │   └── answer_heads.py      # Discriminative (MC) + Generative (open-ended)
+│   │
+│   ├── data/
+│   │   ├── __init__.py
+│   │   ├── unified_dataset.py   # 9-benchmark unified loader with format quirks
+│   │   └── data_utils.py        # Collator, dataloader factory, benchmark configs
+│   │
+│   ├── training/
+│   │   ├── __init__.py
+│   │   ├── trainer.py           # 3-phase training loop
+│   │   └── phase_scheduler.py   # Phase transitions, LR scheduling
+│   ��
+│   ├── evaluation/
+│   │   ├── __init__.py
+│   │   └── metrics.py           # Accuracy, ANLS, VQA Acc, Relaxed Acc
+│   │
+│   └── utils/
+│       ├── __init__.py
+│       ├── visualization.py     # Trajectory PCA, gate analysis
+│       └── ablation.py          # Systematic ablation runner
+│
+├── results/                     # Training results (auto-pushed)
+│   ├── hybrid_main.json
+│   ├── no_jepa.json
+│   ├── no_rollout.json
+│   └── ...
+│
+└── checkpoints/                 # Best model checkpoints (auto-pushed)
+    ├── hybrid_main_best.pt
+    └── ...
 ```
 
-See `mr_jepa/ARCHITECTURE.md` for the complete specification.
+---
+
+## Paper Contribution
+
+> **A world model for multimodal reasoning**: We demonstrate that modeling the evolution of a latent belief state via JEPA-style prediction improves performance on static multimodal benchmarks compared to single-pass baselines. The evidence-gated rollout with K=3 steps learns meaningful intermediate reasoning states, validated through ablation studies across 9 benchmarks. The JEPA objective (not human chain-of-thought) supervises a latent trajectory generated by an EMA target encoder, showing that self-supervised dynamics training transfers to discriminative reasoning tasks.
+
+---
+
+## Key References
+
+1. **I-JEPA** (Assran et al., 2023) — [arxiv:2301.08243](https://arxiv.org/abs/2301.08243): JEPA architecture, EMA target, L2 loss, narrow predictor
+2. **LeWorldModel** (Maes et al., 2025) — [arxiv:2603.19312](https://arxiv.org/abs/2603.19312): SIGReg anti-collapse, end-to-end JEPA
+3. **Coconut** (Yu et al., 2024) — [arxiv:2412.06769](https://arxiv.org/abs/2412.06769): Chain of Continuous Thought, latent reasoning
+4. **DINOv3** (Meta, 2025) — [arxiv:2508.10104](https://arxiv.org/abs/2508.10104): Dense SSL with RoPE + Gram anchoring
+5. **SoftCoT++** (Xu et al., 2025) — [arxiv:2505.11484](https://arxiv.org/abs/2505.11484): Soft chain-of-thought with contrastive learning
+
+## License
+
+Apache-2.0