SPRIG v0.1 — Design & Contracts (single source of truth)

Reference docs: full architecture spec ~/workspace/SPRIG-C-architecture.md; project plan ~/.claude/plans/ok-now-i-ltierally-calm-sifakis.md (data/eval/infra details). v0.1 = finite-lattice Stage A at 64x64. No continuous cuts, no morphogen, no occlusion masks, no DINO channel.

1. Repo layout & file ownership

sprig/data/procgen/{vocab,sampler,captions,render,writer}.py   # owner: PROCGEN agent
sprig/data/{dataset,embed_t5}.py, sprig/data/clevr/prep.py     # owner: DATA agent
sprig/dp/{lattice,inside}.py                                   # owner: DP agent
sprig/model/{gmt,atlas,dl,sprig}.py                            # owner: MODEL agent
sprig/eval/{baseline_pixmix,tree_metrics,color_checks,probe,prompts,report,monitors}.py  # owner: EVAL agent
train.py, configs/*.yaml, infra/*.sh                           # owner: HARNESS agent
tests/test_<module>_*.py                                       # each agent writes tests for its own modules

Python 3.12, torch. No deps beyond: torch, numpy, pillow, transformers, sentencepiece, protobuf, tensorboard, matplotlib, pyyaml, tqdm, einops, pytest. Everything must run CPU-only for tests (no .cuda() hardcoded; device from config).

2. Global config (configs/main64.yaml values)

Canvas 64x64 RGB. Grid stride g=8 px → 8x8 cells.
S=1024 nonterminal symbols, R=64 rule components, T_v=256 texels, d=384 model width.
Caption: precomputed T5-base token embeddings [L,768], L≤64.
Leaf-eligible regions: both sides ≤ 16 px. Regions with any side > 16 px MUST expand. 8x8 regions MUST terminate.
Batch 256. AdamW: lr 3e-4 for embedding tables (E_N, E_T, V, W, P_T, cut-type tables), 1e-4 for GMT + atlas renderer + Φ; betas (0.9, 0.95); cosine decay to 10% over 250k steps, 2k warmup; grad clip 1.0; EMA 0.9999 (eval only).
Rule-logit temperature τ_ann: 2.0 → 1.0 linearly over first 50k steps (divide rule/termination logits by τ_ann).
10% of steps: caption replaced by null embedding (dataloader-level).
bf16 autocast for matmuls; ALL logsumexp/log-domain accumulation in fp32. Emission DL log-scales clamped [-7, 2].

3. Lattice (sprig/dp/lattice.py) — DP agent

Region = axis-aligned rect with corners on the 8-px grid, i.e. cell-interval pair. All (C(9,2))²=1296 rects are regions. Precompute once (pure function of canvas/grid config, cache to .pt):

regions: int32 [N_reg, 4] (x0,y0,x1,y1 in px), region_id lookup dict.
Cuts: for each region, every interior grid line on each axis is a valid cut → children are regions. Global cut-type vocabulary for parameter tying: (axis ∈ {H,V}) × (relative offset bucket ∈ 7 buckets, nearest of {1/8..7/8}) = 14 cut types. Per region: cut_list of (cut_type_id, child_lo_id, child_hi_id). child_lo = lesser-coordinate child (left for V, top for H) — this ordering is what lets the grammar express left/right, above/below.
Level order: by cell area ascending (levels 1..64). Per level: flattened index tensors parent_ids [M], cut_type [M], child_lo [M], child_hi [M] for vectorized DP.
leaf_mask [N_reg] (both sides ≤16), must_terminate [N_reg] (8x8), must_expand [N_reg] (any side >16), phi_geom [N_reg, 64] Fourier features of (log-area, log-aspect, center-x, center-y).
Leaf shape groups: leaf regions grouped by (h,w) ∈ {8,16}² for batched emission scoring.

4. Model (sprig/model/) — MODEL agent

gmt.py — GrammarModulationTransformer. Symbol embeddings E_N [S,d] (queries; no symbol-symbol self-attention). 4 blocks: {cross-attn to projected caption tokens (768→d, key_padding_mask from emb_len), FFN d→4d→d, pre-LN}. Output H [B,S,d]. Heads:

U = H @ W_u → p(k|A,c) logits [B,S,R].
Termination: MLP([H_A ; phi_geom(r)]) → logit per (B, region, S) — computed lazily as MLP over H then dot with a geometry projection: implement as term_logit[b,r,A] = MLP_h(H)[b,A,:] · MLP_g(phi_geom)[r,:] + bias_A (factorized, cheap).
Cut-type distribution p(s|k,c): component embeddings e_k [R,d] cross-attend once to caption → logits [B,R,14]; per region, mask to valid cut types and renormalize (log-softmax over the masked set). NOTE: multiple concrete cuts in a region can share a cut type — split p(s|k,c) mass uniformly across same-type concrete cuts (add -log(count) correction, precomputed per region).
Terminal texel prior: p(T|A,c) = Σ_k p(k|A,c) softmax(P_T[k]) with static P_T [R,T_v].
Children: static V, W [R,S]: p(B|k)=softmax(V[k]), p(C|k)=softmax(W[k]). (v0.1: no q(B) coupling.)
Illumination field Φ(c): mean-pooled caption → MLP → deconv to [B,8,16,16], bilinear-resampled at leaf positions, FiLM (scale,shift per DL-mean channel group) on emission means.

atlas.py — TexelAtlas. E_T [T_v,d]. Renderer: per texel, cross-attn(E_T row → caption, 2 blocks width 256) → seed [256,4,4] → 2× (conv + 2x nearest-upsample) → atlas [B,T_v,40,16,16]. 40 ch = 4 DL components × (1 weight + 3 means + 3 log-scales + 3 channel-coupling coeffs). Per-texel additive bias grid [T_v,40,16,16] (trainable table, no renderer) — the resurrection-writable parameterization (F3.3/M1.2).

Emission scoring: for each leaf shape group (h,w): resample atlas 16x16 → (h,w) via adaptive avg pool; FiLM by Φ at leaf position; score pixels under 4-comp discretized logistic with RGB coupling (means: μ_R; μ_G+α·R; μ_B+β·R+γ·G), quantized to 256 bins → ell [B, n_leaf_in_group, T_v] summed over pixels. fp32 result.

dl.py — discretized logistic mixture log-prob (PixelCNN++ style, vectorized, fp32-safe).

sprig.py — SPRIGModel (implements contracts):

log_marginal(image_u8, emb, emb_len) -> logZ [B] — full inside DP (tempered during training via self.eta; an eta=0 flag for reported numbers).
loss(batch) -> (loss, metrics_dict) — see §6.
map_parse(image, emb, emb_len) -> list[ParseNode] — Viterbi (max-semiring) + backtrace. ParseNode{rect, axis, cut_px, symbol, texel, children} (dataclass in sprig/model/sprig.py; texel/leaf fields None for internal nodes).
posterior_usage(image, emb, emb_len) -> dict(symbol_usage [S], texel_usage [T_v], node_entropy, emit_mag, rule_mag, mean_depth, mean_leaves) — expected counts via the autograd identity: make per-(region,symbol) termination potentials and per-(region,cut) potentials require grad, grad of logZ w.r.t. them = posterior marginals; node entropy = occupancy-weighted entropy of conditional split posteriors (good enough for the PI controller).
sample(emb, emb_len, seed_struct, seed_material, n) -> (images_u8 [n,64,64,3], trees) — ancestral, breadth-parallel over the frontier; two torch.Generators: structural draws (term/k/s/B/C) from seed_struct, material draws (texel choice + pixel sampling; pixels = DL means for v0.1 crispness, texel choice sampled) from seed_material. Best-of-K: sample_bestof(emb, K) reranks K derivations by joint log p(τ, x̂|c).

5. Inside DP (sprig/dp/inside.py) — DP agent

inside(ell_leaf, term_logits, cut_logits, U_logmix, logV, logW, lattice, temper_kappa) -> beta [B,N_reg,S] fp32, logZ [B]

β(r,A) for leaf-eligible r: logaddexp( log p_term + logsumexp_T(log p(T|A,c) + ell(r,T)/κ(r)), log(1−p_term) + expand_term(r,A) ) with must_terminate/must_expand masks (∓inf).
expand_term via level-synchronous sweep, area ascending. Per level, with index tensors: Bhat[m,k] = logbmm(beta[:, child_lo[m], :], logV.T); same for Chat with logW; comb[m,k] = logsumexp over cuts grouped by parent of (log p(s=cut_type|k,c) − log count_correction + Bhat + Chat); expand[parent,A] = logbmm(comb, U_log.T) where U_log = log-softmax(U/τ_ann). Group-by-parent logsumexp via torch.segment_reduce or index_put with scatter-logsumexp helper (write one: max-shift per segment, scatter_add of exp, log). All fp32.
logbmm(x_log [.., K], w_log [K, M]): max-shifted: x_log.max(-1) shift, exp, matmul (bf16 ok), log, add shifts. Provide exact fp32 fallback for tests.
Tempering: κ(r) = max(1, area_px(r)^η), η a module buffer. PI controller (in train.py, every 2k steps): target node entropy band [0.5, 3.0] nats; if H < 0.5: η += 0.05 + 0.1·(0.5−H); if H > 3.0: η −= 0.05; clamp [0, 1.5]. Anneal η→0 linearly over final 20k steps; all REPORTED bpd at η=0.
Viterbi: same sweep, max instead of logsumexp, argmax backtrace tables per level.
CORRECTNESS TEST (blocking): tiny lattice (16x16 canvas, 8px grid → 9 regions), tiny config (S=3, R=2, T_v=2, random weights), brute-force enumerate ALL derivation trees (recursive over cuts/symbols/texels), sum exact joint probs, compare to inside logZ within 1e-4. Also: Viterbi tree's joint prob ≤ logZ; posterior marginals from autograd sum to expected node counts.

6. Loss & regularizers

L = mean_b( −logZ_tempered / (3·64·64) )                      # nats/subpixel
  + 1.0 · Σ_T max(0, 1/(4·T_v) − texel_usage_T)               # texel under-use hinge (F3.2)
  + 0.5 · Σ_A max(0, 1/(4·S)  − symbol_usage_A)               # symbol under-use hinge (M1-style)

usage vectors = batch-mean posterior expected counts, normalized to sum 1 (from posterior_usage computed on the training batch — reuse the same graph's grads: compute via one extra backward-free trick or just take grads of logZ before the optimizer backward; simplest correct: compute usage with torch.autograd.grad(logZ.sum(), potentials, retain_graph=True)). InfoNCE caption-contrast: DEFERRED; add only if caption-swap margin is flat by 50k steps. Dead-texel resurrection (train.py, every 2k steps): texels with usage < 0.1/T_v → overwrite bias grid with (a random training-image 16x16 crop converted to DL-mean params, small noise on other channels), perturb E_T row ±ε. Log resurrection count.

7. Training curriculum (procedural main run)

Steps 0–200k on proc2d with tier reweighting: 0–20k tiers (0:.55,1:.35,2:.10,3:0); 20k–60k (.25,.35,.30,.10); >60k (.10,.30,.40,.20) [target mix].
η per PI controller; τ_ann 2→1 over 50k; report-eta-0 bpd from 180k; final 20k η→0 anneal.
CLEVR fine-tune (separate config): init from main ckpt, lr×0.3, 30–50k steps, 20% proc2d replay.
Smoke config (configs/smoke.yaml): S=128, R=16, T_v=32, d=128, batch 32, 200 steps, val_fast 64 — must run CPU-only on Mac AND on klaus-1 GPU.

8. Contracts (C1–C4) — binding for all agents

C1 batch (from sprig/data/dataset.py collate): {image: u8 [B,64,64,3], emb: f16 [B,L,768], emb_len: i32 [B], tier: i8 [B], idx: i64 [B]}. Null-caption substitution (10%, train only) in the dataset/loader using ~/data/sprig/t5/null.f16. C2 model.log_marginal(image, emb, emb_len) -> logZ [B] (η=0 honored when model.eval() + report_mode=True). C3 model.map_parse(...), model.posterior_usage(...) as §4. C4 model.sample(emb, emb_len, seed_struct, seed_material, n), model.sample_bestof(emb, emb_len, K, seed). Data formats, eval metrics, prompt bank, run infra: per the plan file §Part 1/3/4 (memmaps, meta.jsonl GT-tree schema, 32-prompt bank, scalars.jsonl, atomic checkpoints, run_forever.sh) — copy schemas exactly from there.

9. Milestone gates (in order, each blocking)

G0: full pytest suite green on Mac CPU (incl. DP brute-force test, holdout scan test, determinism hashes, tree-metric self-test).
G1: overfit 1 image (script scripts/overfit1.py): 2k steps tiny config → bpd < 1.0, MAP parse stable across steps.
G2: overfit 100 images with captions: Δ_c > 0 (caption info gain), texels alive > 50%.
G3: klaus-1 smoke (30 min real config): >2 steps/s at batch 256 (else profile before main run), dataloader ≥95% synthetic throughput, ckpt kill/resume works, eval callback produces grids/overlays.
G4: main run launch. Twice-daily monitor checks per plan alarms.