Update REPA metrics and experiment documentation

EXPERIMENTS_REPA.md

@@ -0,0 +1,170 @@
+# PixelDiT Representation-Alignment (REPA) — experiments & handoff
+
+Study of **what to align PixelDiT's intermediate features to** (the "teacher" in REPA-style
+training) on ImageNet-256, plus a per-layer feature analysis. Variants A–E keep PixelDiT's
+flow-matching loss and alignment point (patch **block 8**) fixed and only change the teacher;
+variant F is a follow-up hybrid that adds SRA guidance at block 4 while retaining block-8 REPA.
+This doc is the authoritative summary + handoff (results, infra, current state, how to continue).
+
+---
+
+## 1. Setup / infra
+
+- **Model**: PixelDiT-XL, 797M — pixel-space DiT, 26 patch blocks + 4 pixel blocks, patch 16, 256×256.
+- **Data**: ImageNet-1K, REPA-E 256px H5 at `imagenet256_data/images.h5` (253 GB, restored from S3
+  `s3://dit-scale-up/fe-experiments/xinlin/pixeldit_imagenet256_xl_lx_vb/data/`).
+- **Recipe** (identical across variants unless noted): 8×80GB, global batch 256 (8×32), AdamW lr 1e-4,
+  EMA 0.9999 (SRA/Staged use **0.999**, see §3D), `bf16-mixed`, 70k steps, checkpoint every 10k
+  (`save_top_k=-1`, all kept, ~12.8 GB each).
+- **Eval**: `main.py predict` → 50k samples (50/class, FlowDPMSolver 100-step, CFG 2.75, interval
+  [0.1,0.9]) → OpenAI guided-diffusion **ADM evaluator** vs `eval_refs/VIRTUAL_imagenet256_labeled.npz`
+  → FID / sFID / IS. FID worker computes one point per 10k checkpoint.
+- **Envs**: `/mnt/localssd/pixeldit-venv` (torch 2.5.0+cu124, training/predict); ADM evaluator in
+  `/mnt/localssd/adm-eval-venv` (tf 2.15) + `/mnt/localssd/guided-diffusion`. Static `tmux` at
+  `~/.local/bin/tmux`.
+- **W&B**: project `pixeldit-imagenet-repro` on `adobesensei.wandb.io` (API key needed;
+  `wandb==0.19.11` required — 40-char check rejects the self-hosted key on ≥0.20).
+- **Run a variant** (8-GPU): `scripts/train_pixeldit_imagenet256_xl.sh <config.yaml> [--ckpt_path=...]`
+  in one tmux; `scripts/<variant>_fid_eval.sh` in another (`WANDB_API_KEY`, `PIXELDIT_SKIP_S3_RESTORE=1`
+  set). FID/predict shares the 8 GPUs with training (~0.06 it/s, ~1.5–2h per 50k point).
+
+## 2. Per-layer feature analysis (which layer is semantic)
+
+On the official converged PixelDiT-XL (`nvidia/PixelDiT-ImageNet` epoch320), hooking each patch
+block and PCA→RGB + linear-CKA-vs-DINOv2: **semantic content peaks sharply at patch block 8**
+(CKA 0.85), exactly the REPA anchor; early blocks = detail, late blocks specialize toward the
+denoising output. DiT-XL/2-256's own features are semantic mid-late (blk 12–21). Scripts:
+`scripts/viz_layer_features.py`, `viz_layer_similarity.py`, `viz_dit_layers.py`
+(figures in `viz_layer_features/`, `viz_dit_layers/` — PNGs are gitignored, regenerate from scripts).
+
+## 3. Experiments (teacher = the only thing that changes)
+
+- **A. baseline (DINOv2 REPA)** — project block-8 tokens (MLP→768), cosine-align to frozen DINOv2
+  patch tokens of the clean image, weight 0.5. `pix256_xl_orig_10k.yaml`. W&B loss `y4489tgd`, FID `xo84hpbc`.
+- **B. DiT-REPA (blk12)** — teacher = frozen DiT-XL/2-256: VAE-encode → noise at ddpm t=100 → DiT
+  forward (class-cond) → block-12 features. **Gotcha:** raw DiT features are DC-dominated (inter-token
+  cosine 0.965) so cosine aligns trivially → fixed by per-image mean-centering. `src/dit_repa.py`,
+  `pix256_xl_ditrepa12.yaml`. W&B `s075fe0g` / `o113h6qc`.
+- **C. dino+VAE** — keep DINOv2 REPA **and** add a 2nd projector on block 8 aligned to the frozen
+  **SD-VAE encoder mid-block** feature (DC-centered), λ_gen 0.05. `src/dino_vae_repa.py`,
+  `pix256_xl_dinovae.yaml`. W&B `0e4hcesw` / `pixeldit_dinovae_fid`.
+- **D. SRA (self-distillation, EMA teacher)** — no external teacher: student block-8 (through a
+  projector) aligns to the **EMA net's block-16** at a lower-noise timestep, smooth-L1 β0.05, weight
+  0.04. Port of [vvvvvjdy/SRA](https://github.com/vvvvvjdy/SRA). **Two fixes were required:**
+  (1) `use_proj=True` — reference projects the STUDENT and aligns to the RAW teacher (asymmetric);
+  (2) **EMA decay 0.9999 → 0.999** — SRA reads the EMA forward every step and drives the online
+  weight-norm up; the wide 0.9999 window lags that norm so the EMA (not the online model) samples as
+  noise by ~30k. A tighter 0.999 window tracks it → clean EMA. `SRATrainer` in `src/diffusion.py`
+  (`use_proj` flag), `pix256_xl_sra.yaml`. W&B `77xtrzcq` / `pixeldit_sra_fid`.
+- **E. StagedTeacher (self-distillation, frozen-snapshot teacher)** — like SRA but the teacher is a
+  **frozen hard copy of the online model, refreshed every 10k steps** (a periodically-updated target
+  network) instead of an EMA. Resumed from SRA's 10k checkpoint. `src/staged_teacher.py`,
+  `pix256_xl_stageteach.yaml` (`refresh_interval: 10000`). W&B `5ywf06q4` / `pixeldit_stageteach_fid`.
+- **F. SRA4 + DINO-REPA8 (hybrid)** — one online forward supplies two complementary targets:
+  online **block 4** (student projector) aligns by smooth-L1 to the EMA net's raw **block 16** at a
+  lower-noise timestep, while online **block 8** independently aligns by cosine to frozen DINOv2
+  clean-image patch tokens. Weights remain SRA 0.04 and REPA 0.5; EMA decay is 0.999.
+  `src/hybrid_sra_repa.py`, `pix256_xl_sra_repa4.yaml`. W&B training run `49p8qfc5`;
+  FID curve pending.
+
+## 4. Results — 50k-sample ADM eval
+
+**FID** (↓ better)
+
+| step | baseline | DiT-REPA | dino+VAE | SRA (EMA) | StagedTeacher |
+|---|---|---|---|---|---|
+| 10k | 389.3 | 368.7 | 383.1 | 187.5 | 187.5¹ |
+| 20k | 363.6 | 350.6 | 387.5 | **102.8** | 116.0 |
+| 30k | 106.1 | 108.1 | 127.9 | **82.2** | 86.9 |
+| 40k | 66.8 | 81.7 | **52.6** | 67.1 | 68.7 |
+| 50k | **31.1** | 57.2 | 39.2 | 55.8 | 58.4 |
+| 60k | **21.3** | 39.6 | 25.9 | — | 49.5 |
+| 70k | — | — | **17.0** | — | 43.1 |
+
+**sFID** (↓ better)
+
+| step | baseline | DiT-REPA | dino+VAE | SRA | StagedTeacher |
+|---|---|---|---|---|---|
+| 10k | 414.7 | 417.3 | 398.4 | 50.4 | 50.4¹ |
+| 20k | 304.7 | 197.4 | 403.2 | **13.9** | 17.5 |
+| 30k | 78.2 | 45.6 | 73.2 | **10.0** | 13.0 |
+| 40k | 46.0 | 42.7 | 30.5 | **8.5** | 8.9 |
+| 50k | 20.3 | 32.2 | 27.0 | 7.9 | **7.8** |
+| 60k | 14.7 | 20.3 | 18.0 | — | **7.3** |
+| 70k | — | — | 14.1 | — | **6.8** |
+
+**IS** (↑ better)
+
+| step | baseline | DiT-REPA | dino+VAE | SRA | StagedTeacher |
+|---|---|---|---|---|---|
+| 10k | 1.07 | 1.04 | 1.06 | **5.40** | 5.40¹ |
+| 20k | 1.40 | 1.64 | 1.00 | **9.59** | 8.43 |
+| 30k | **11.71** | 10.48 | 9.94 | 11.15 | 10.39 |
+| 40k | 17.51 | 10.82 | **19.19** | 12.59 | 12.40 |
+| 50k | **26.65** | 13.84 | 24.78 | 14.19 | 13.73 |
+| 60k | **32.94** | 18.56 | 30.52 | — | 15.36 |
+| 70k | — | — | **37.30** | — | 17.12 |
+
+¹ StagedTeacher resumed from SRA's 10k checkpoint, so its 10k = SRA's 10k. Chart:
+`slides_assets/fid_comparison_5way.png` (regenerate; PNGs gitignored).
+
+## 5. Findings
+
+1. **Self-distillation (SRA, StagedTeacher) converges dramatically faster early and dominates sFID
+   at every step.** At 10k SRA is FID 187 / IS 5.4 while the external-teacher variants are still
+   ~380 / IS ~1; SRA's 20k FID (103) ≈ the DINOv2 baseline's 30k. sFID stays 2–5× better throughout.
+2. **But self-distillation plateaus on FID/IS.** By 50–60k the external-teacher variants overtake on
+   FID (60k: baseline 21.3, dino+VAE 25.9 vs StagedTeacher 49.5) and on IS. Net: self-distillation =
+   fast + best structure (sFID), external semantic teacher = better final FID/diversity (IS).
+   dino+VAE reaches the best FID by 70k (17.0).
+3. **EMA teacher (SRA) ≈ frozen-snapshot teacher (StagedTeacher), EMA slightly better early.** Gap
+   20k 116 vs 103, narrowing to ~identical by 40k. The gain comes from the self-distillation mechanism
+   (block-8→block-16 alignment), not from how the teacher is updated.
+4. **DiT-blk12 as teacher (B) < DINOv2 (A)**; **dino+VAE (C) ≈ baseline** (slightly better late).
+5. **Diffusion-model features need DC removal** to be a usable cosine-alignment teacher (B, C).
+6. Debugging lesson (SRA): the EMA "noise" was **not** EMA corruption — the EMA tracked online
+   (cosine 0.985, drift 0) but lagged in *norm* under the wide 0.9999 window; several wrong fixes
+   (stream, compile, clobber) were tried before the real cause (norm-lag) → decay 0.999.
+
+## 6. Current state (as of 2026-07-28)
+
+- **Running**: hybrid SRA4 + DINO-REPA8 (`49p8qfc5`), newly started from scratch. tmux sessions
+  `hybrid_train`, `hybrid_fid`; the FID worker targets every 10k checkpoint from 10k through 70k.
+- **Completed**: StagedTeacher reached 70k; its final point is FID 43.11 / sFID 6.78 / IS 17.12.
+- **Checkpoints on local SSD** (`train_logs/exp_pixeldit_imagenet256_xl_<variant>/`, every 10k, ~12.8GB):
+  baseline 10–60k, DiT-REPA 10–70k, dino+VAE 10–70k, SRA 10–60k, StagedTeacher 20–70k.
+  **NOT backed up to S3** — a box wipe loses them (`scripts/sync_pixeldit_to_s3_loop.sh` can back up).
+- baseline & SRA runs were stopped early (~62k) to free GPUs for later variants, hence FID only to 60k/50k.
+
+## 7. Handoff notes (for continuing)
+
+- To finish a curve: relaunch its `scripts/<variant>_fid_eval.sh` (idempotent — skips computed points;
+  set `STEPS=...` to target specific steps). Checkpoints are all saved, so any FID point can be
+  (re)computed offline without retraining.
+- To resume/continue training: `train_pixeldit_imagenet256_xl.sh <config> --ckpt_path=<...ckpt>`.
+- Open directions: (a) push SRA/Staged past 70k or with a **stronger λ / IS-oriented tweak** to fix
+  the FID/IS plateau while keeping the sFID lead; (b) **hybrid**: DINOv2 REPA + SRA self-distillation
+  (semantic teacher for late FID/IS + self-distillation for early speed & sFID); (c) tune SRA
+  student/teacher blocks & `sra_t_max`. (d) back up checkpoints to S3 first.
+- Gotchas: EMA decay must be ≤0.999 for any variant that reads the EMA during training (norm-lag);
+  keep DC-centering for diffusion/VAE-feature teachers; `nn.Module.compile()` is lazy (first-forward);
+  the CLI list-index override `--trainer.callbacks.N...` is rejected by jsonargparse (edit the YAML).
+
+## 8. File map
+
+```
+pixeldit_configs/pix256_xl_orig_10k.yaml    # A baseline (DINOv2)
+pixeldit_configs/pix256_xl_ditrepa12.yaml   # B DiT-REPA blk12 (DC-centered)
+pixeldit_configs/pix256_xl_dinovae.yaml     # C DINOv2 + SD-VAE detail head
+pixeldit_configs/pix256_xl_sra.yaml         # D SRA (use_proj=true, EMA decay 0.999)
+pixeldit_configs/pix256_xl_stageteach.yaml  # E StagedTeacher (refresh_interval 10000)
+pixeldit_configs/pix256_xl_sra_repa4.yaml   # F hybrid: SRA 4->16 + DINOv2 REPA on block 8
+PixelDiT/c2i/src/diffusion.py               # REPATrainer, DINOv2, SRATrainer (use_proj flag)
+PixelDiT/c2i/src/dit_repa.py                # B
+PixelDiT/c2i/src/dino_vae_repa.py           # C
+PixelDiT/c2i/src/staged_teacher.py          # E
+PixelDiT/c2i/src/hybrid_sra_repa.py          # F
+scripts/*_fid_eval.sh                       # per-variant FID-curve workers
+scripts/log_fid_wandb.py                    # append an FID point to a W&B run
+scripts/viz_*_layers*.py, build_slides.py   # analysis + slide deck
+```

HYBRID_SRA_REPA.md

@@ -0,0 +1,122 @@
+# PixelDiT Hybrid SRA + REPA
+
+## Motivation
+
+The previous controlled experiments showed complementary behavior:
+
+- DINOv2 REPA converges slowly at the beginning, but gives stronger late-stage FID and IS.
+- SRA converges much faster and gives substantially better sFID, but its FID/IS improvement slows
+  after roughly 40k steps.
+
+The hybrid experiment keeps the original DINOv2 REPA target on patch block 8 and adds an earlier
+SRA target on patch block 4. The goal is to retain SRA's early structural acceleration without
+removing the fixed external semantic anchor that benefits late-stage generation.
+
+## Method
+
+One online PixelDiT forward captures two intermediate features:
+
+```text
+noisy image x_t
+    |
+    +-- online block 4 -- SRA projector -- Smooth-L1 --> EMA block 16
+    |                                                 (same image and noise,
+    |                                                  lower-noise timestep)
+    |
+    +-- online block 8 -- REPA projector -- cosine --> frozen DINOv2 patch tokens
+                                                      (clean image)
+```
+
+The training objective is:
+
+\[
+\mathcal{L}
+=
+\mathcal{L}_{FM}
++ 0.04\,\mathcal{L}_{SRA}^{4\rightarrow16}
++ 0.5\,\mathcal{L}_{REPA}^{8\rightarrow DINOv2}.
+\]
+
+### SRA branch
+
+- Student: online PixelDiT patch block 4.
+- Teacher: EMA PixelDiT patch block 16.
+- The teacher receives a lower-noise point no more than `0.2` away from the student timestep.
+- Student and teacher use the same clean image, class label, and sampled noise.
+- A two-layer MLP projects only the student feature.
+- Alignment uses patch-wise Smooth-L1 with `beta=0.05`.
+- Loss weight is `0.04`.
+- EMA decay is `0.999`, which avoids the feature-norm lag observed with `0.9999` in PixelDiT.
+
+### REPA branch
+
+- Student: online PixelDiT patch block 8.
+- Teacher: frozen DINOv2 ViT-B/14 patch tokens from the clean image.
+- A three-layer MLP projects the PixelDiT feature from 1152 to 768 dimensions.
+- Corresponding spatial patch tokens are aligned with cosine distance.
+- Loss weight remains `0.5`, unchanged from the original PixelDiT DINOv2 REPA baseline.
+
+Both online features are captured during the same forward pass. The EMA and DINOv2 teacher
+branches run without gradients. Both projection heads are training-only and are omitted at
+inference.
+
+## Implementation
+
+- `PixelDiT/c2i/src/hybrid_sra_repa.py`
+  - Defines `HybridSRAREPATrainer`.
+  - Captures online blocks 4 and 8 in one forward.
+  - Reuses the existing SRA timestep construction and Smooth-L1 alignment.
+  - Preserves the existing DINOv2 REPA cosine alignment and spatial-grid matching behavior.
+  - Logs the unweighted and weighted SRA/REPA losses separately.
+- `pixeldit_configs/pix256_xl_sra_repa4.yaml`
+  - Configures SRA `4 -> 16`, REPA on block 8, EMA decay `0.999`, and 70k training steps.
+- `scripts/sra_repa4_fid_eval.sh`
+  - Builds the 10k-to-70k FID curve with 50k ADM samples per checkpoint.
+
+## Initial result
+
+ImageNet-256, PixelDiT-XL, global batch 256, evaluated with the same 50k-sample ADM protocol:
+
+| Experiment at 10k | FID ↓ | sFID ↓ | IS ↑ |
+|---|---:|---:|---:|
+| DINOv2 REPA baseline | 389.3 | 414.7 | 1.07 |
+| DINOv2 + VAE | 383.1 | 398.4 | 1.06 |
+| SRA, block 8 -> 16 | 187.5 | 50.4 | 5.40 |
+| StagedTeacher | 187.5 | 50.4 | 5.40 |
+| **Hybrid SRA 4 -> 16 + REPA block 8** | **86.67** | **18.28** | **11.30** |
+
+Relative to the original SRA run at 10k, the hybrid lowers FID by about 54%, lowers sFID by about
+64%, and increases IS by about 2.1x. Its 10k FID is also already better than the original SRA
+20k FID of 102.8.
+
+These are early results; later checkpoints are required to determine whether the hybrid preserves
+the original REPA run's late-stage FID/IS advantage.
+
+## Running the experiment
+
+```bash
+PIXELDIT_SKIP_S3_RESTORE=1 \
+scripts/train_pixeldit_imagenet256_xl.sh \
+pixeldit_configs/pix256_xl_sra_repa4.yaml
+```
+
+Run the FID worker separately:
+
+```bash
+PIXELDIT_SKIP_S3_RESTORE=1 scripts/sra_repa4_fid_eval.sh
+```
+
+The active W&B training run is `49p8qfc5` (`pixeldit_sra4_repa8_loss`) in the
+`pixeldit-imagenet-repro` project.
+
+## Validation performed
+
+- Python syntax compilation for the hybrid trainer.
+- Shell syntax validation for the FID worker.
+- LightningCLI config parsing with the complete YAML.
+- Synthetic forward/backward test verifying:
+  - gradients reach the online network;
+  - gradients reach both projection heads;
+  - no gradients reach the EMA teacher.
+- Real 8-GPU launch and sustained training beyond 20k steps.
+- Successful 10k checkpoint sampling and 50k ADM evaluation.

METRICS.md

@@ -0,0 +1,22 @@
+# Baseline vs hybrid metric summary
+
+All values use the same ImageNet-256 50k-sample ADM evaluation protocol.
+
+## DINOv2 REPA baseline
+
+| Step | FID ↓ | sFID ↓ | IS ↑ |
+|---:|---:|---:|---:|
+| 10,000 | 389.2824 | 414.7020 | 1.0710 |
+| 20,000 | 363.6215 | 304.7034 | 1.4038 |
+| 30,000 | 106.0903 | 78.1710 | 11.7050 |
+| 40,000 | 66.7541 | 45.9615 | 17.5108 |
+| 50,000 | 31.1164 | 20.3315 | 26.6462 |
+| 60,000 | 21.3407 | 14.7369 | 32.9395 |
+
+## Hybrid SRA4 + REPA8
+
+| Step | FID ↓ | sFID ↓ | IS ↑ |
+|---:|---:|---:|---:|
+| 10,000 | 86.6690 | 18.2777 | 11.2964 |
+| 20,000 | 44.5866 | 11.0986 | 19.4438 |
+| 30,000 | 28.4069 | 9.0362 | 26.7527 |

dinov2_repa_baseline/README.md

@@ -1,3 +1,17 @@
-# DINOv2 REPA baseline checkpoints
+# DINOv2 REPA baseline
 
-PixelDiT-XL ImageNet-256 checkpoints at 10k-step intervals.
+PixelDiT-XL on ImageNet-256. Online patch block 8 is projected and aligned
+patch-wise to frozen DINOv2 clean-image tokens with cosine distance.
+
+## 50k-sample ADM metrics
+
+| Step | FID ↓ | sFID ↓ | IS ↑ |
+|---:|---:|---:|---:|
+| 10,000 | 389.2824 | 414.7020 | 1.0710 |
+| 20,000 | 363.6215 | 304.7034 | 1.4038 |
+| 30,000 | 106.0903 | 78.1710 | 11.7050 |
+| 40,000 | 66.7541 | 45.9615 | 17.5108 |
+| 50,000 | 31.1164 | 20.3315 | 26.6462 |
+| 60,000 | 21.3407 | 14.7369 | 32.9395 |
+
+Raw evaluator outputs are stored under `metrics/`.

dinov2_repa_baseline/metrics.csv

@@ -0,0 +1,7 @@
+step,FID,sFID,IS,Precision,Recall
+10000,389.2823967413824,414.7020167260448,1.0709521770477295,,
+20000,363.62150542991213,304.70335112320214,1.4037563800811768,,
+30000,106.09033483097272,78.1709826505786,11.705048561096191,,
+40000,66.75411054088266,45.96150425066094,17.510835647583008,,
+50000,31.116447437985357,20.331490638313994,26.646221160888672,,
+60000,21.34065440024318,14.736901695812094,32.93950653076172,,

dinov2_repa_baseline/metrics.json

@@ -0,0 +1,38 @@
+[
+  {
+    "FID": 389.2823967413824,
+    "IS": 1.0709521770477295,
+    "sFID": 414.7020167260448,
+    "step": 10000
+  },
+  {
+    "FID": 363.62150542991213,
+    "IS": 1.4037563800811768,
+    "sFID": 304.70335112320214,
+    "step": 20000
+  },
+  {
+    "FID": 106.09033483097272,
+    "IS": 11.705048561096191,
+    "sFID": 78.1709826505786,
+    "step": 30000
+  },
+  {
+    "FID": 66.75411054088266,
+    "IS": 17.510835647583008,
+    "sFID": 45.96150425066094,
+    "step": 40000
+  },
+  {
+    "FID": 31.116447437985357,
+    "IS": 26.646221160888672,
+    "sFID": 20.331490638313994,
+    "step": 50000
+  },
+  {
+    "FID": 21.34065440024318,
+    "IS": 32.93950653076172,
+    "sFID": 14.736901695812094,
+    "step": 60000
+  }
+]

dinov2_repa_baseline/metrics/step-10000-adm_eval.txt

@@ -0,0 +1,3 @@
+Inception Score: 1.0709521770477295
+FID: 389.2823967413824
+sFID: 414.7020167260448

dinov2_repa_baseline/metrics/step-20000-adm_eval.txt

@@ -0,0 +1,3 @@
+Inception Score: 1.4037563800811768
+FID: 363.62150542991213
+sFID: 304.70335112320214

dinov2_repa_baseline/metrics/step-30000-adm_eval.txt

@@ -0,0 +1,3 @@
+Inception Score: 11.705048561096191
+FID: 106.09033483097272
+sFID: 78.1709826505786

dinov2_repa_baseline/metrics/step-40000-adm_eval.txt

@@ -0,0 +1,3 @@
+Inception Score: 17.510835647583008
+FID: 66.75411054088266
+sFID: 45.96150425066094

dinov2_repa_baseline/metrics/step-50000-adm_eval.txt

@@ -0,0 +1,3 @@
+Inception Score: 26.646221160888672
+FID: 31.116447437985357
+sFID: 20.331490638313994

dinov2_repa_baseline/metrics/step-60000-adm_eval.txt

@@ -0,0 +1,3 @@
+Inception Score: 32.93950653076172
+FID: 21.34065440024318
+sFID: 14.736901695812094

hybrid_sra4_repa8/METHOD.md

@@ -0,0 +1,122 @@
+# PixelDiT Hybrid SRA + REPA
+
+## Motivation
+
+The previous controlled experiments showed complementary behavior:
+
+- DINOv2 REPA converges slowly at the beginning, but gives stronger late-stage FID and IS.
+- SRA converges much faster and gives substantially better sFID, but its FID/IS improvement slows
+  after roughly 40k steps.
+
+The hybrid experiment keeps the original DINOv2 REPA target on patch block 8 and adds an earlier
+SRA target on patch block 4. The goal is to retain SRA's early structural acceleration without
+removing the fixed external semantic anchor that benefits late-stage generation.
+
+## Method
+
+One online PixelDiT forward captures two intermediate features:
+
+```text
+noisy image x_t
+    |
+    +-- online block 4 -- SRA projector -- Smooth-L1 --> EMA block 16
+    |                                                 (same image and noise,
+    |                                                  lower-noise timestep)
+    |
+    +-- online block 8 -- REPA projector -- cosine --> frozen DINOv2 patch tokens
+                                                      (clean image)
+```
+
+The training objective is:
+
+\[
+\mathcal{L}
+=
+\mathcal{L}_{FM}
++ 0.04\,\mathcal{L}_{SRA}^{4\rightarrow16}
++ 0.5\,\mathcal{L}_{REPA}^{8\rightarrow DINOv2}.
+\]
+
+### SRA branch
+
+- Student: online PixelDiT patch block 4.
+- Teacher: EMA PixelDiT patch block 16.
+- The teacher receives a lower-noise point no more than `0.2` away from the student timestep.
+- Student and teacher use the same clean image, class label, and sampled noise.
+- A two-layer MLP projects only the student feature.
+- Alignment uses patch-wise Smooth-L1 with `beta=0.05`.
+- Loss weight is `0.04`.
+- EMA decay is `0.999`, which avoids the feature-norm lag observed with `0.9999` in PixelDiT.
+
+### REPA branch
+
+- Student: online PixelDiT patch block 8.
+- Teacher: frozen DINOv2 ViT-B/14 patch tokens from the clean image.
+- A three-layer MLP projects the PixelDiT feature from 1152 to 768 dimensions.
+- Corresponding spatial patch tokens are aligned with cosine distance.
+- Loss weight remains `0.5`, unchanged from the original PixelDiT DINOv2 REPA baseline.
+
+Both online features are captured during the same forward pass. The EMA and DINOv2 teacher
+branches run without gradients. Both projection heads are training-only and are omitted at
+inference.
+
+## Implementation
+
+- `PixelDiT/c2i/src/hybrid_sra_repa.py`
+  - Defines `HybridSRAREPATrainer`.
+  - Captures online blocks 4 and 8 in one forward.
+  - Reuses the existing SRA timestep construction and Smooth-L1 alignment.
+  - Preserves the existing DINOv2 REPA cosine alignment and spatial-grid matching behavior.
+  - Logs the unweighted and weighted SRA/REPA losses separately.
+- `pixeldit_configs/pix256_xl_sra_repa4.yaml`
+  - Configures SRA `4 -> 16`, REPA on block 8, EMA decay `0.999`, and 70k training steps.
+- `scripts/sra_repa4_fid_eval.sh`
+  - Builds the 10k-to-70k FID curve with 50k ADM samples per checkpoint.
+
+## Initial result
+
+ImageNet-256, PixelDiT-XL, global batch 256, evaluated with the same 50k-sample ADM protocol:
+
+| Experiment at 10k | FID ↓ | sFID ↓ | IS ↑ |
+|---|---:|---:|---:|
+| DINOv2 REPA baseline | 389.3 | 414.7 | 1.07 |
+| DINOv2 + VAE | 383.1 | 398.4 | 1.06 |
+| SRA, block 8 -> 16 | 187.5 | 50.4 | 5.40 |
+| StagedTeacher | 187.5 | 50.4 | 5.40 |
+| **Hybrid SRA 4 -> 16 + REPA block 8** | **86.67** | **18.28** | **11.30** |
+
+Relative to the original SRA run at 10k, the hybrid lowers FID by about 54%, lowers sFID by about
+64%, and increases IS by about 2.1x. Its 10k FID is also already better than the original SRA
+20k FID of 102.8.
+
+These are early results; later checkpoints are required to determine whether the hybrid preserves
+the original REPA run's late-stage FID/IS advantage.
+
+## Running the experiment
+
+```bash
+PIXELDIT_SKIP_S3_RESTORE=1 \
+scripts/train_pixeldit_imagenet256_xl.sh \
+pixeldit_configs/pix256_xl_sra_repa4.yaml
+```
+
+Run the FID worker separately:
+
+```bash
+PIXELDIT_SKIP_S3_RESTORE=1 scripts/sra_repa4_fid_eval.sh
+```
+
+The active W&B training run is `49p8qfc5` (`pixeldit_sra4_repa8_loss`) in the
+`pixeldit-imagenet-repro` project.
+
+## Validation performed
+
+- Python syntax compilation for the hybrid trainer.
+- Shell syntax validation for the FID worker.
+- LightningCLI config parsing with the complete YAML.
+- Synthetic forward/backward test verifying:
+  - gradients reach the online network;
+  - gradients reach both projection heads;
+  - no gradients reach the EMA teacher.
+- Real 8-GPU launch and sustained training beyond 20k steps.
+- Successful 10k checkpoint sampling and 50k ADM evaluation.

hybrid_sra4_repa8/README.md

@@ -1,3 +1,18 @@
-# Hybrid SRA4 + REPA8 checkpoints
+# Hybrid SRA4 + REPA8
 
-SRA online block 4 to EMA block 16 plus DINOv2 REPA on online block 8.
+PixelDiT-XL on ImageNet-256 with two representation targets:
+
+- SRA: online block 4 to EMA block 16 at a lower-noise timestep, Smooth-L1 weight 0.04.
+- REPA: online block 8 to frozen DINOv2 clean-image tokens, cosine weight 0.5.
+
+See [`METHOD.md`](./METHOD.md) for implementation details and validation.
+
+## 50k-sample ADM metrics
+
+| Step | FID ↓ | sFID ↓ | IS ↑ |
+|---:|---:|---:|---:|
+| 10,000 | 86.6690 | 18.2777 | 11.2964 |
+| 20,000 | 44.5866 | 11.0986 | 19.4438 |
+| 30,000 | 28.4069 | 9.0362 | 26.7527 |
+
+Raw evaluator outputs are stored under `metrics/`.

hybrid_sra4_repa8/metrics.csv

@@ -0,0 +1,4 @@
+step,FID,sFID,IS,Precision,Recall
+10000,86.66904452942163,18.277675999740154,11.29637622833252,,
+20000,44.58660002402064,11.098587139252231,19.443777084350586,,
+30000,28.406907534417826,9.036190922661376,26.752666473388672,,

hybrid_sra4_repa8/metrics.json

@@ -0,0 +1,20 @@
+[
+  {
+    "FID": 86.66904452942163,
+    "IS": 11.29637622833252,
+    "sFID": 18.277675999740154,
+    "step": 10000
+  },
+  {
+    "FID": 44.58660002402064,
+    "IS": 19.443777084350586,
+    "sFID": 11.098587139252231,
+    "step": 20000
+  },
+  {
+    "FID": 28.406907534417826,
+    "IS": 26.752666473388672,
+    "sFID": 9.036190922661376,
+    "step": 30000
+  }
+]

hybrid_sra4_repa8/metrics/step-10000-adm_eval.txt

@@ -0,0 +1,3 @@
+Inception Score: 11.29637622833252
+FID: 86.66904452942163
+sFID: 18.277675999740154

hybrid_sra4_repa8/metrics/step-20000-adm_eval.txt

@@ -0,0 +1,3 @@
+Inception Score: 19.443777084350586
+FID: 44.58660002402064
+sFID: 11.098587139252231

hybrid_sra4_repa8/metrics/step-30000-adm_eval.txt

@@ -0,0 +1,3 @@
+Inception Score: 26.752666473388672
+FID: 28.406907534417826
+sFID: 9.036190922661376