# Autoresearch: fast differentiable token clustering for ZipSplat

Pure exploration (all code in this folder). Autoresearch loop, run manually
following `~/.claude/skills/autoresearch/SKILL.md`.

## Goal
The ZipSplat forward pass clusters DA3 backbone tokens into K queries (each ->
32 Gaussians). Today it uses `hard_kmeans_chunked` under `no_grad` (NOT
differentiable; gradient reaches the model only via the gathered token, not the
assignment). Find a clustering that is:
1. **fast + memory-light** (forward AND backward),
2. **differentiable** wrt the input features (likely soft fixed point + IFT),
3. **at least as good as k-means** in feature-space distortion — or auto-K.

Clustering runs on the **original pretrained DA3-Large features** (frozen,
`weights/da3/da3-large.pth`, priors OFF), NOT zipsplat-finetuned/fused features.

## Autoresearch config
- **Scope (modifiable):** `cluster.py` only (`cluster(x, K) -> centroids`, diff wrt x).
- **Metric (minimize):** median fwd+bwd wall-clock ms at `CLUSTER_LAYER`, `PRIMARY_RATIO`.
- **Guard (all pass, else discard):**
  - quality: WCSS <= kmeans WCSS * (1 + 0.03)   [k-means is the WCSS optimum, so "match" is the bar]
  - quality: silhouette >= kmeans silhouette - 0.02
  - differentiable: finite, non-zero grad to input features (`grad_check`)
  - memory: peak (fwd+bwd) <= kmeans peak (fwd) * 2.0
- **Direction:** lower ms is better. Auto-K credited via lower effective cost at equal WCSS.
- **Baseline:** `hard_kmeans_chunked` (production), measured once in `build_cache.py` -> `baseline_nopriors.json`.

## Metric rationale
Metric is INTRINSIC (feature-space) by user choice — no render/head in the loop,
so every iteration is cheap and fully decoupled. Two caveats carried, not blocking:
(1) k-means globally minimises WCSS, so any soft/diff method sits slightly above it
-> guard is "within 3%", not "beat". (2) most-clusterable layer (silhouette) is not
guaranteed to be the most render-useful layer -> optional frozen-head render anchor
can spot-check the layer pick later.

## Setup
```
conda run -n splatfactory --no-capture-output python -m playground.autoresearch_clustering.build_cache
```
Caches 8 DL3DV scenes (6 ctx views, frozen DA3-Large, priors OFF) at study layers
[3,7,11,15,19,23], and writes the k-means baseline + per-layer clusterability table.

## Loop protocol (per iteration)
1. Review `results.tsv` + `git log` of this folder.
2. Make ONE change to `cluster.py` (a CONFIG knob or the algorithm).
3. `git add playground/autoresearch_clustering/ && git commit -m "experiment: ..."`.
4. `conda run -n splatfactory ... python -m playground.autoresearch_clustering.verify "<note>"`.
5. Keep if METRIC < champion AND GUARD=PASS; else surgically revert this folder.
6. (verify auto-appends a row to `results.tsv`.)

GPU hygiene: every timed pass waits until no other process holds the GPU.

## Setup notes / current state
- DA3-Large: 24 blocks (0..23), embed_dim 1024, cat_token -> D=2048. alt_start=8
  (camera token + global cross-view attention begin at block 8). out_layers prod=[11,17,23].
- Priors OFF: backbone uses its built-in learned `camera_token` (standard DA3 inference).
- Token count: 6 views x T patches @252px (square-pad, patch 14).

## Layer study (8 scenes, k-means; normWCSS = WCSS / total-variance, lower = more clusterable)
| layer | normWCSS r0.5 | normWCSS r0.25 | sil r0.5 | note |
|---|---|---|---|---|
| L3  | 0.121 | 0.218 | 0.540 | early, local-only (< alt_start=8) — least clusterable |
| L7  | 0.104 | 0.188 | 0.527 | early |
| L11 | **0.044** | **0.083** | 0.549 | deep; best at high compression |
| L15 | 0.059 | 0.113 | 0.553 | |
| L19 | 0.054 | 0.106 | 0.551 | |
| L23 | 0.046 | 0.092 | **0.554** | **prod clustering_layer=0**; highest sil -> CLUSTER_LAYER |
- **Deep layers (L11/L23) are markedly more clusterable** (~4-5% residual variance) than early (L3/L7, ~12%).
- Silhouette ~flat across layers (~0.55 @ r0.5) — not discriminative. Raw WCSS grows with depth (feature norms grow); normWCSS controls for it.
- Speed/mem identical across layers (same N, D). Loop fixes on **L23** (matches prod). L11 is the runner-up.

## Results (8 DL3DV scenes, layer=23, ratio=0.5). k-means ref: 17ms fwd, 194MB, WCSS 1086, sil 0.554
| iter | method | fwd+bwd ms | WCSS ratio | sil | peak MB | grad cos | guard |
|---|---|---|---|---|---|---|---|
| 0 | soft k-means UNROLLED (it5,T0.1) | 61 | 0.997 | 0.554 | 411 | 1.000 | **FAIL (mem>388)** |
| 1 | + phantom (1-step IFT) | 32 | 0.996 | 0.554 | 284 | 1.000 | PASS |
| 2 | + bf16 | 16.3 | 1.012 | 0.554 | 233 | — | PASS (< kmeans fwd!) |
| 3 | n_iters=2 | 10.3 | 1.017 | 0.552 | 233 | 0.913 | PASS |
| 4 | n_iters=1 (FASTEST) | 8.5 | 1.026 | 0.550 | 234 | 0.842 | PASS (WCSS at edge) |
| 5 | T=0.05 | 9.0 | 1.028 | 0.549 | 234 | — | DISCARD (worse) |
| **6** | **n_iters=3 (EXACT-GRAD champion)** | **11.8** | **1.014** | **0.553** | **234** | **1.000** | **PASS** |

**CHAMPION: soft k-means | phantom-IFT | bf16 | cdist | autograd-bwd | fixed n_iters=3 | T=0.1 | uniform**
- 6v: 12.3ms fwd+bwd (1.37x faster than the 17ms NON-differentiable production kmeans), 234MB (1.2x kmeans).
- WCSS 1.014x kmeans, silhouette ties. EXACT gradient (cosine 1.0 vs fully-unrolled).
- Holds at ratio 0.25 (WCSS 1.019) and scales: 1.75x @12v, 1.70x @24v (advantage GROWS with V).
- Fastest variant: n_iters=1 -> 8.5ms (~1ms/scene) but gradient only 0.84-faithful.

### Speed/mem loop (iters 7-12): sub-ms investigation (all on top of the iter-6 cdist champion)
| iter | change | result |
|---|---|---|
| 7 | matmul-expansion ‖x-c‖², all updates fp32 | 18.1ms DISCARD (fp32 fwd cost > cdist bwd saved) |
| 8 | split: bf16 cdist converge + matmul-fp32 grad | 13.8ms DISCARD (same reason, smaller) |
| 9 | cossim (norm'd matmul assign + raw agg, bf16) | WCSS 18.3x FAIL — DA3 feats in a narrow cone, cosine non-discriminative |
| 10-11 | early-stopping convergence (rel move < es_tol) | works (3-4 adaptive iters) but per-iter GPU sync -> 12-37ms VARIANCE. es_tol=0 default |
| 12 | custom matmul backward (no [N,K,D]) | SAME speed as autograd (5.8 vs 5.9ms) -> cdist bwd is NOT the bottleneck. fast_bwd=False |

### View-count scaling, DA3-LARGE (champion vs kmeans, ratio 0.5, layer 23, D=2048)
| V | N | K | kmeans ms | champ ms | speedup | km MB | champ MB | WCSS |
|---|---|---|---|---|---|---|---|---|
| 6 | 1944 | 972 | 16.9 | 12.3 | 1.37x | 202 | 234 | 1.014 |
| 12 | 3888 | 1944 | 57.8 | 33.0 | 1.75x | 385 | 485 | 1.009 |
| 24 | 7776 | 3888 | 218.8 | 128.7 | 1.70x | 735 | 1017 | 1.020 |

### View-count scaling, DA3-GIANT = production primary (layer 39, D=3072; CL_BACKBONE=giant)
| V | N | K | kmeans ms | champ ms | speedup | km MB | champ MB | WCSS |
|---|---|---|---|---|---|---|---|---|
| 6 | 1944 | 972 | 24.6 | 14.9 | 1.64x | 280 | 316 | **0.998** |
| 12 | 3888 | 1944 | 84.7 | 50.0 | 1.69x | 553 | 666 | **0.996** |
| 24 | 7776 | 3888 | 325.0 | 181.8 | **1.79x** | 1071 | 1383 | **0.995** |
- **On Giant the champion is EVEN BETTER than on Large**: WCSS < 1.0 (BEATS hard-kmeans distortion),
  1.64-1.79x faster (margin grows with V), exact gradient. Bigger D=3072 makes kmeans' 5 Lloyd iters
  relatively costlier -> our fewer-updates advantage grows. Generalizes cleanly to the production backbone.
- **Giant layer study (normWCSS r0.5):** L9 0.176, L19 0.129, **L29 0.070 (most clusterable)**, L39 0.105.
  Production clusters on L39 (deepest); L29 clusters notably better (same mid-depth pattern as Large's L11).
  -> worth testing clustering_layer on L29 for Giant.

## Key findings
- **k-means ref is genuinely memory-light: 194MB fwd, 17ms** (chunked cdist+scatter). The earlier 2084MB
  was an all-layers-resident artifact; verify.py recomputes the ref LIVE on single-layer scenes.
- **Unrolled soft k-means is the wrong shape for backward: 411MB** (stores n_iters x [N,K] soft-assign) ->
  fails the 2x-kmeans mem guard, 61ms. The phantom 1-step IFT fixes both: detached fixed-point forward
  (cheap, no_grad) + ONE differentiable update -> O(1) iteration backward memory.
- **The phantom (Jacobian-free) gradient is EXACT at convergence, not an approximation** (validate.py):
  cosine vs true unrolled grad = 0.84 @2 updates, 0.91 @3, **1.0000 @>=4**. At a converged soft-kmeans
  fixed point df/dC -> 0, so (I - df/dC)^-1 = I and the 1-step gradient = the full IFT gradient. So
  "converge under no_grad (cheap) + 1 grad step" gives the exact gradient for free.
- **bf16 was the single biggest speed win** (32 -> 16ms, 2x) and cut memory 284 -> 233MB. Matches the
  production autocast precision; WCSS cost negligible (0.996 -> 1.012).
- Quality is a non-issue at high K (ratio 0.5/0.25): soft ties hard k-means. The game is SPEED + MEMORY of
  a DIFFERENTIABLE method, plus GRADIENT FIDELITY (the dimension the ms-metric alone hid -> needs n_iters>=3).
- DISCARDED: lower temperature (sharper assignment tracks 1 token, not the 2-token-mean optimum at K=N/2).

## Sub-ms investigation conclusions (iters 7-12)
- **The cost is the cdist soft-updates themselves** (~0.7ms/scene each), NOT the backward and NOT launch
  overhead. Controlled head-to-head: custom matmul backward == autograd cdist backward (5.8 vs 5.9ms);
  batching [B,N,D] gave 0 speedup (GPU already saturated). The only speed knob that works is fewer
  updates (n_iters), which trades gradient fidelity (n_iters=1 0.84-grad 8.5ms .. n_iters=3 exact 12ms).
- **Every algebraic shortcut to a cheap bf16 matmul fails for DA3 features**: fp32 matmul-expansion is
  too slow; bf16 expansion catastrophically cancels (large norms); cosine is non-discriminative (feats
  in a narrow cone). Euclidean + stable precision is required -> cdist (or a fused kernel) is the way.
- **Data-dependent early stopping is a trap in the hot path**: the per-iter `if move<tol` forces a
  GPU->CPU sync -> 12-37ms variance. Sync-free fixed n_iters (tuned to the measured convergence: 3-4
  iters -> exact gradient) is the right call. Convergence is non-monotonic (bf16 + discrete reassign).

## Remaining directions (not yet explored)
- **True sub-ms/scene**: a fused Triton kernel for ONE soft-update (cdist+softmax+aggregate in SRAM,
  recompute in backward) -- the only untried lever, high-effort, uncertain gain (we appear compute-bound
  on the cdist, not bandwidth/launch). The n_iters=1 variant (~1ms/scene, 0.84-grad) is the cheap floor.
- **Auto-K (the stretch goal)**: DP-means (one lambda threshold spawns clusters) -> scene-adaptive K,
  differentiable relaxation. Would need a count-aware metric (current guard assumes matched K).
- **Render anchor**: wire the champion into zipsplat `_cluster` + frozen-head render to confirm the intrinsic
  (WCSS) win transfers to PSNR, and to validate the L23-vs-L11 layer pick against the real task.
