I added the gate experiment to Daystorm’s Stage B training last week to see whether the vision-language backbone was simply repeating learned phrases. The case study had shown that a model could reach high fused_exact scores by memorising the most common answer tokens, so I needed a way to force the model to rely on the actual sensor prefixes.
First I tried copying the training loop from stage_a.py and adding a second loss term that penalised high confidence on shuffled prefixes. I shuffled the modality order inside each batch before feeding it to the model, then computed a cross-entropy loss on the shuffled predictions and added it to the main loss with a weight of 0.5. After three epochs the fused_exact score stayed above 0.93 but the shuffled-exact score also stayed above 0.80, meaning the model was still memorising regardless of order. This approach failed because the shuffled prefixes still contained the same token distribution; the model could learn to ignore the order and rely on the shared vocabulary.
I then implemented the gate check exactly as described in the case study. I created a helper function that runs two forward passes: one with the original fusion prefix and one with a randomly permuted prefix across the batch. The function returns the exact match accuracy for both passes. I added this function to the training script and logged both metrics every fifty steps. The gate check lives in train_stage_b.py and looks like this:
def gate_check(model, batch):
fused = model.forward(batch, prefix="fusion")
shuffled = model.forward(batch, prefix="shuffle")
fused_exact = (fused.logits.argmax(-1) == batch.answer).float().mean()
control_exact = (shuffled.logits.argmax(-1) == batch.answer).float().mean()
return fused_exact.item(), control_exact.item()
I called gate_check inside the training loop after computing the regular loss. If fused_exact dropped below 0.90 or control_exact rose above 0.50 I stepped back the learning rate by a factor of two and continued. The first few runs showed fused_exact hovering around 0.88 and control_exact at 0.62, indicating the model was still leaning on memorisation. I increased the QLoRA rank from 4 to 8 and raised the alpha from 8 to 16, which gave the adapter more capacity to learn modality-specific projections without overfitting the backbone. After that change fused_exact climbed to 0.91 while control_exact fell to 0.48, satisfying the gate condition.
The gate experiment added measurable overhead. Each training step now required two forward passes instead of one, increasing step time from 12 ms to 22 ms on the T4. To keep the effective batch size at 32 I reduced gradient accumulation from four steps to two, which raised VRAM usage from 13.8 GB to 15.2 GB, leaving only 800 MB for activations. I switched the optimizer’s state to fp16 to stay within the 16 GB limit, and I re-enabled torch.compile with mode="reduce-overhead" and attn_implementation="sdpa" to recover some of the lost speed. The compiled step time settled at 18 ms, a 50 % increase over the baseline but still acceptable for nightly runs.
I also updated the checkpoint format to store the gate’s threshold values alongside the model weights. The checkpoint now includes a metadata.json with "gate_fused_threshold": 0.90 and "gate_control_threshold": 0.50. When resuming from a checkpoint I read these values and restore the gate logic; otherwise the training script would use hard-coded defaults and potentially drift from the intended behaviour.
The gate experiment has so far prevented the model from collapsing into a memorised state on three different random seeds. However I still see occasional spikes where control_exact jumps to 0.55 for a few hundred steps before falling back. I suspect the shuffling mechanism occasionally creates batches where the permuted prefix still aligns strongly with a particular answer token due to class imbalance in the synthetic dataset. Next I plan to stratify the shuffle by answer label so that each batch contains a balanced distribution of shuffled prefixes across classes. I also want to log the gate metrics to TensorBoard alongside the loss curves to make the behaviour more visible during debugging. Until those changes are in place I will keep the gate thresholds as they are and monitor the training logs for any sustained deviation from the pass criteria.