Contextual encoder: what one training step does

Stage 2 as implemented in contextual_encoder/model.py (run s2_v1). One 32 s crop, batch dimension dropped. Read left to right; the amber cells are the only places the loss is computed.

visible latent step (context) masked step (target, loss here) frozen, no gradient
1 · TARGETS (frozen) piano roll, 32 s 4 × 88 × 1560 frames S1-a encoder frozen · stride 6 · RF 0.5 s drop 2 edge steps per side teacher latents z 256 steps × 256 d · 8 fps span 1 s span 1.75 s 0.5 s spans 4–32 steps (0.5–4 s), sampled until 30 % of valid steps are covered; same z serves as input and as target 2 · TRUNK (trained; this is the representation) trunk input Linear 256→512 · masked steps ← learned mask token [M] [M] [M] [M] [M] [M] [M] [M] [M] [M] [M] [M] [M] same z, masked Transformer trunk · 12 layers × 512 · RoPE · SwiGLU · 37.9 M full self-attention over all 256 steps, mask tokens included (so context knows where the gaps are and how long) F = trunk output 256 steps × 512 d unused unused by the loss unused at inference: nothing is masked, F at every step → Stage-3 heads, probes 3 · PREDICTOR (trained, then discarded) keys / values: F at visible steps only (masked and padded steps are excluded) 2 × cross-attention + SwiGLU · 6.4 M queries = mask token + RoPE at the masked position: "what should be at step t, given the visible context?" predicted ẑ Linear 512→256, read at masked steps loss = Smooth-L1( ẑ , z ) at masked ∧ valid steps z standardized per dim (teacher stats); gradient flows into predictor and trunk, never into the teacher logged beside it: interp = same loss for "mean of nearest visible neighbours"; ratio = loss / interp target z at masked steps
One training step. The same teacher latents play two roles: masked and projected they are the trunk's input, untouched they are the target. The trunk sees the whole timeline including the mask tokens, so visible steps know where the gaps are. The predictor asks, for each masked position, what belongs there, but it may only look at trunk outputs at visible steps. Only the amber cells are scored. The trunk's own outputs at masked steps receive no direct loss, which is what keeps F from becoming a reconstruction head.
PREDICTOR DETAIL · one masked step t inside a span F masked span masked visible context (left) visible context (right) — may contain a repeat of the masked material q_t mask token + RoPE(t) carries position, no content cannot attend to its own span (or any masked step) interp baseline: mean of these two neighbours only
Why the ratio is the honest learning curve. A query for step t sees every visible step in the crop, near and far, and none of the masked ones. The interpolation baseline sees only the two visible latents adjacent to the span. Loss ÷ interp below 1 therefore measures exactly what the trunk adds beyond blending neighbours; at step 5k it is 0.72 on scores and 0.69 on performances.

Shapes at each stage (batch of one crop)

tensorshapenote
piano roll4 × 88 × 15601536 frames + 12 overhang per side; channels active, onset-phase, velocity, validity
teacher latents z256 × 256260 steps from the encoder, 2 dropped per side; standardized
mask256 bool≈77 steps masked in 4 to 32-step spans; visible = valid ∧ ¬mask
trunk input256 × 512projection of z; masked rows replaced by the mask token
F256 × 512final LayerNorm output; the representation
predictor queries256 × 512mask token at every step for static shapes; only masked ones are scored
prediction ẑ256 × 256compared with z at masked valid steps

Three things the picture settles

Why the trunk sees mask tokens at all. Without them the trunk would have to be run on a shorter, gap-free sequence, and its visible outputs would not know that a gap exists or how long it is. With them, the trunk is exercised on full-length sequences, which is also how it is used at inference when nothing is masked.

Why the predictor is separate. If the trunk's own output at step t had to equal the teacher latent, the last layers would turn into a reconstruction head and F would carry that specialization into every probe. The predictor absorbs it and is thrown away.

Why spans, not scattered steps. Adjacent teacher latents share input frames (0.5 s receptive field), so a lone masked step is mostly visible through its neighbours. The 4-step minimum equals the receptive field, so the interior of every span is out of any visible latent's sight.