Three changes combined, attention projector, batching support, and explicit variance regularisation, and the first genuinely positive result after three parts of negative ones.
Part 12 ended with a list of partial explanations, none severe enough on its own to account for the full collapse. Before the session started, a new candidate surfaced worth taking seriously: representation collapse.
The next-token loss only ever asks "does this embedding help predict the next word?" It has no built-in requirement that embeddings from different images actually look different from each other. If the reasoning text's structure, movement words, point-tag formatting, is predictable enough on its own, the loss can be minimised without the vision embeddings carrying much image-specific signal at all. The model learns to produce near-identical embeddings for every image because the labels do not demand otherwise.
This is a known failure mode in self-supervised learning, not unique to this project.
VICReg's approach adds an explicit loss term that is blind to what the labels say and only cares about the embeddings themselves. For every feature dimension, compute the standard deviation across the batch (across different images) and penalise any dimension whose standard deviation falls below a target:
relu(target - std).mean()This directly prevents all images from collapsing to the same embedding region, independent of what the next-token loss is doing.
One important constraint: this can only be computed with more than one example in the batch at a time. Which is exactly why batching was bundled into this session rather than treated as a separate concern.
VLM.from_pretrained now instantiates AttentionProjector again instead of the plain MLP. The existing verify_vlm.py smoke test had a stale assumption baked in: it assumed the projector's output length always equals the raw patch count, which is true for the MLP but not for AttentionProjector, which compresses 196 patches down to 32 learned queries. Fixed the assertion to measure the projector's actual output length rather than assuming it. Clean 32 vision plus 8 text equals 40 token sequence confirmed.
The original concern was that Model.forward's causal mask logic, which had only ever been exercised at batch size 1, would need significant work to support padding. Working through it carefully first: vision tokens are a fixed 32-token block at the very start of every sequence, and padding always sits at the very end of the right-padded text portion.
Causal masking already guarantees a query token can never attend to a key at a later position. Since pad tokens are always later than every real token, causal masking already excludes them automatically, with zero extra code.
The only place padding needed explicit handling was the loss: cross-entropy needed a mask so padded positions do not contribute gradient. make_collate_fn in data/maze_dataset.py pads variable-length reasoning text to the batch's longest example and returns a loss_mask alongside the padded input_ids.
A whole planned category of work turned out to be unneeded.
Hoped for batch_size=8. Reality: an 8GB card cannot hold it.
The real cost driver is not attention or the projector at all. It is the (batch, seq_len, vocab_size=151936) logits tensor, which scales linearly with batch size regardless of anything else. batch_size=8 crashed immediately. batch_size=4 crashed a bit later, precisely when a batch happened to contain two longer-than-usual sequences at once. Landed on batch_size=2 as the practical ceiling.
While chasing this down, found and fixed a genuine inefficiency: the point-token logit splice in qwen3/model.py was doing logits = logits.clone(); logits[..., ids] = point_logits, allocating a second full-vocab tensor just to overwrite two columns. Replaced with in-place logits.index_copy_(-1, ids, point_logits). Small win, but the kind of thing worth catching once vocab size is the dominant memory cost.
First real batched training attempt crashed at step ~180 of 250. Notably not because anything was going wrong: ce_loss had already dropped from ~2.7 to the 0.1 to 0.2 range, exactly the healthy trend expected.
The cause: padding forces every example in a batch up to the longest one's length. A single unusually long maze path can blow up memory for a whole batch that would otherwise fit fine.
Fixed with a length cap in MazeDataset.__getitem__: if a generated path's token count exceeds max_tokens (300), retry with a deterministic fallback seed (idx + attempt * num_examples) rather than truncating the text mid-reasoning. Truncating mid-path would hand the model corrupted supervision, a path that stops partway through with no exit reached. The fallback is still fully reproducible since the seed is a pure function of idx.
Retrained both stages from scratch with the new setup. ce_loss and variance_loss fell together throughout training, variance dropping from 0.96 to 0.01 to 0.15, CE from 3.2 to 0.1 to 0.2. A good sign in itself: a collapse fix that just traded one loss for the other would not actually mean anything.
The real test: re-ran the cross-image cosine-similarity diagnostic on held-out seeds (900 to 904, never in the training set) to remove any doubt about train and test overlap.
| Stage | Mean similarity |
|---|---|
| Raw untrained ViT patches (theoretical floor for this data) | 0.44 to 0.48 |
| Every previous post-training attempt (Parts 9 to 12) | 0.80 to 0.9989 |
| This session: attention projector + batching + variance reg | 0.5423 |
A real, large improvement. Much closer to the raw-pixel floor than any previous post-training measurement.
Generation confirmed the same story from a different angle. Across 6 different maze seeds, the outputs are no longer the byte-identical collapse from Part 9. Four of 6 mazes produced genuinely different first moves, direction sequences, and coordinates, clearly tracking something about each specific image. But 2 of the 15 pairwise comparisons were still exact duplicates.
Read honestly: a shift from "one universal answer regardless of input" to "a small set of memorised attractor paths." Real progress, a different and less severe failure mode, but not yet full per-maze grounding.
[[1294,915]]). Every training label is constructed to stay within that range by construction, so this looks like extrapolation past the training distribution rather than a bug in the coordinate maths. Worth watching.The representation-collapse hypothesis held up under a real test. Attacking it directly, attention pooling, explicit variance regularisation, and batching to make that regularisation meaningful, cut the collapse roughly in half on the 0.44 to 0.99 scale. None of the previous attempts (rank, isolated training, colour, split learning rate) got anywhere close.
Not a finished story. The remaining gap between 0.54 and the ~0.46 floor, and the attractor-path pattern in generation, both suggest there is still real collapse left to address. Possible next directions: a stronger variance-loss weight, more training data and steps now that batching makes larger-scale training more efficient, or examining whether the attractor paths correlate with something structural about the mazes that land in the same bucket.
Decide which single variable to change next, stronger variance-loss weight, more data, or more steps, and test it cleanly rather than stacking changes. The baseline is now 0.5423 on held-out seeds. Any future change can be compared against that number directly.