File size: 6,022 Bytes
d832dce
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
# notable-gam: LocusGS loop without per-token tables

This branch starts from `grad`. The goal is to replace the learned per-token tables (`gs_tokens`,
`gs_anchor_xyz_raw`, `gs_anchor_radius_raw`) with a formula, so that the token count can later adapt per scene.
The model is trained **from scratch**; only the TokenGS image encoder is loaded.

## What differs from `grad`

All changes are in `locusgs/models/locusgs.py`, `locusgs/models/enc_dec.py`, `locusgs/train.py`,
`locusgs/data/registry.py` and `locusgs/version.py`.

1. **No tables.** Token *i* gets Sobol point `u_i` (3-D, not scrambled, origin skipped).
   - **Anchor start:** a canonical frustum in the normalised frame (first camera at the origin, looking +z).
     The direction is `((2u1-1)*0.5, (2u2-1)*0.5, 1)`; the disparity `1/z` is uniform in `[1/100, 1/0.1]`.
     It is placed by a global learnable shift and scale (3 + 3 numbers).
   - **Radius:** one global learnable scalar, init 1.0.
   - **Token:** `h_i = W * RFF(2u_i - 1) + b`, a linear map of random Fourier features (fixed B, 3x512, sigma 4, seed 1).
   - `gs_tokens` and the two anchor tables are now properties computed from this formula.
2. **Error CNN feedback branch removed.** The G3R gradient feedback stays, into the token and into the anchor/radius steps.
3. **Round-1 stop-gradient removed.** With it, the encoder/kv gradient was about 1e-5 (grad: about 1e-2).
4. **Decoder LayerScale gamma.** It is initialised to grad's trained block-0 channel means: cross 0.0212, self 0.0008,
   mlp 0.0945. The old init was 5e-5.
5. **Recipe.** lr 1e-4 constant after a 125-update ramp. Global batch 32 by gradient accumulation (batch 1 x 32).
   Local data: gs960p subsets 1K-3K.

## How the released checkpoint was trained

- `notable_canon` ran 14 epochs from scratch, with stop-grad.
- `notable_nosg` resumed it without stop-grad, to ep 26.
- `notable_gam` then resumed with gamma set to grad's values. In that checkpoint gamma was overwritten by a constant,
  which flipped the sign of the negative channels; the eval dropped 1 dB and recovered within 1 epoch.

The release asset is at **epoch 396**: eval 20.09 (mean of the last 10 epochs), train PSNR 19.41.
It holds `model.safetensors`, `optimizer.pth` (AdamW moments), `scheduler.pth` and `metadata.json`.

## Train / resume

```bash
# data roots (override the defaults in locusgs/data/registry.py)
export DL3DV_GS960P_ROOT=/path/to/dl3dv_gs960p      # contains 1K/ 2K/ 3K/
export DL3DV_BENCH_ROOT=/path/to/dl3dv_benchmark140
export PYTHONPATH=$PWD

# to RESUME: put the four release files in $WS/checkpoints/checkpoint_latest/ (auto-resume from there)
WS=/path/to/workspace
accelerate launch --config_file acc_configs/gpu1.yaml -m locusgs.train train_dl3dv_base \
  --model_type locusgs --cross_attn_variant geometric_positional --use_anchor_radius --anchor_radius_refinement \
  --anchor_supervision_layers 12 --img_size 256 256 --num_gs_tokens 1024 --enc_depth 12 \
  --workspace $WS --batch_size 1 --lr 1e-4 --num_workers 6 --num_epochs 100000 --checkpoint_save_every 5 \
  --resume /path/to/tokengs_image_encoder_only.safetensors
```

- From scratch, `--resume` must be a file holding **only** the 224 TokenGS image-encoder tensors.
- `locusgs/version.py` pins the configuration and refuses other values: lr, img size, token count and supervision layers.
- `train.py` sets the gradient accumulation so that the global batch is 32.
- LIVE lines to check in the log: `[notable] LIVE`, `[nocnn] LIVE`, `[nosg] LIVE` and `[accum] LIVE update 1`.

## Measured state at ep 376-396, vs grad

grad is `/scratch2/TokenGS-Loop-grad/ckpt`, the released lineage. All rows below were measured on the same scenes.

| quantity | ours | grad |
|---|---|---|
| wrong-scene images (decoder reads another scene's latent) | -1.7 dB | -13 dB |
| all tokens -> their mean | -2 dB | -16 dB |
| block-0 cross-attn content-logit std, round 1 | 2.5 | 24.3 |
| effective rays attended, round 1 | ~600 | 4.7 |
| q_norm / k_proj_norm gain | 1.16 / 1.16 | 2.66 / 3.99 |
| R^2 of the anchor's image position from the query, round 6 (other-scene test) | 0.34 (flat since ep 86) | 0.96 |
| radius r, R1 -> R12 | 0.96 -> 0.3 | 0.21 -> 0.08 |
| live Gaussian size | ~0.025 | 0.008 |

Findings:
- grad's tokens get their image position from **self-attention**. With self-attention zeroed, the R^2 drops from 0.96 to 0.2.
  Ours does not: the position part of the self-attention input is about 0.17x the token-content part (grad about 1.2x).
- The render loss pushes r **larger** in both models; only the visibility term pushes it smaller, and it is 30-50x weaker.
  grad's round-1 radius step is -16 (raw), ours -0.35. In grad's code round 1 is detached, so that step cannot have been
  learned under the current grad code.
- Gradient directions of the shared modules are inconsistent across disjoint 8-scene sets.
  For example the kv_proj cos is -0.09 for ours and +0.50 for grad.
- Renders are coarse and correct (layout and colour) but have no detail.
- Training PSNR has been flat around 19.3 for more than 50 epochs.

Open idea, not tested: init r at 0.3 (LocusGS paper value). The geometric bias scales with 1/r^2.

## Probes (`probes/`, run beside training)

Each probe is called as `CKPT=... python probes/<name>.py train_dl3dv_base <the model flags above>`.
`MEMFRAC` caps the GPU memory fraction.

| probe | what it measures |
|---|---|
| `probe_read.py` | read ablation and wrong-scene test (NB scenes, with SE) |
| `probe_tokshare.py` | token identity |
| `probe_attn.py`, `probe_attn2.py` | cross-attention sharpness and whether the content q.k points near the anchor |
| `probe_posinfo.py` | query -> anchor-position R^2 |
| `probe_sape.py` | self-attention position-vs-content magnitudes |
| `probe_rounds.py` | per-round anchor depth, r and Gaussian size |
| `probe_gs_light.py` | live/dead Gaussians |
| `probe_gradcons2.py` | gradient consistency |
| `probe_gradreach2.py` | per-module gradient norms |
| `render_dump.py` | target-view renders |