TokenGS-Loop-notable-gam / NOTABLE_GAM.md
Yusin's picture
Upload folder using huggingface_hub
d832dce verified
|
Raw History Blame Contribute Delete
6.02 kB

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

# 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