erasedwalt/gsm-gpt2-encoder

GSM LatentMAE — CE-only encoder (dmae_ceonly)

A frozen feature extractor for the MMD term in ELF/flow-map distillation on GSM. It maps a gpt2-large latent to features that the MMD is computed in, instead of using the teacher's mid-trunk features.

Trained with a learned token head and zero latent MSE — the objective is discrete (predict the token) rather than continuous reconstruction. That is what makes it the best of the encoders we tried.

• arch plain, 112M params
• cfg: latent_dim=1280, width=768, depth=8, max_len=640, vocab_size=50257
• expects gpt2-large latents, normalised by the run's latent_mean/latent_std

Measured results

Encoder-quality probe (higher is better):

encoder	paired	centered AUC
dmae_ceonly	0.8250	0.9157
dmae_ce_disc	0.8300	0.9139
dmae_lr1e4 (MSE)	–	0.903
dmae_disc	–	0.8743

Downstream, encoder-only MMD (mmd_weight=0), GSM8K test 1319 rows, 3 seeds, N=32 / shift 128:

arm	acc
ceonly, taps (1,3)	29.19
disc	28.99
ceonly, tap (3)	28.76
ce+disc	28.69
ceonly, tap (1)	28.67
MSE baseline	28.61

Caveats worth knowing before you spend runs on it:

• The whole family sits in a 28.6–29.2 band (sd 0.4–1.3), so the ranking above is within noise. Four quite different encoder objectives produced statistically indistinguishable students.
• Encoder quality does not predict downstream quality. dmae_disc detects corrupted programs at 0.989 AUC (vs 0.67 for this one) and still lands mid-pack downstream.
• Taps matter as much as the encoder: (1,3) > (3) > (1), and deeper taps are clearly worse — (5,6) scored 25.70 vs 28.61 for (1,3).
• On our time-grid recipe this encoder was ~3.7 pts below teacher features. On a fixed-point/resample recipe the gap nearly closed (32.55–33.16). So it is recipe-sensitive.

Use

See usage.py. Two entry points, and picking the wrong one is the most common mistake:

• forward_features(...) -> (B, D) pooled, one sample per example. Use with a paired/linear kernel.
• token_features(...) -> (B, S, D) per token. Use with an RBF kernel where samples are a row's tokens — pooled features leave 1 sample per row and the RBF estimator needs >= 2.