Motivation

• Provide opt-in, high-reward training / inference techniques (Mixture-of-Depth token routing, gated MLP / SquareGLU, EMA self-distillation during warmdown, and Grokfast gradient low-pass) to improve compute-efficiency and generalization within the existing N-gram backoff + legal TTT recipe.
• Keep reported baseline behavior unchanged by making all new techniques disabled by default and configurable via environment flags so re-runs can opt-in without changing published numbers.

Description

• Added new hyperparameters to Hyperparameters for toggling and configuring the features: MOD_*, SQUAREGLU_ENABLED, EMA_DISTILL_*, and GROKFAST_*, and documented them in the submission README and submission.json.
• Implemented Mixture-of-Depth style token routing inside Block with per-block routers and a _make_mod_mask that produces top-k token masks used to selectively scale attention and MLP outputs while remaining legal/score-first during eval.
• Implemented SquareGLU-style gated MLP as an opt-in path in MLP, added an mlp_gate_bank parameter bank, wired gate bank initialization, included it in bank/optimizer grouping, and passed gate slices into block forward calls.
• Added EMA warmdown self-distillation by maintaining an EMA snapshot and periodically syncing it as a teacher during warmdown to compute a KL distillation term added to the training loss, and added a Grokfast gradient low-pass augmentation that maintains an EMA of gradients and injects a fraction back into current gradients before optimizer steps.
• Updated training loop wiring: bank handling, TTT collect logic to include optional mlp_gate_bank, EMA teacher creation, distillation accumulation/logging, Grokfast gradient augmentation before optimizer phases, and small logging additions for the distillation metric.

Testing

• Performed a syntax/compile check with python -m py_compile records/track_10min_16mb/2026-03-31_LeakyReLU2_LegalTTT_NGramCache_XSA/train_gpt.py which completed successfully and reported no syntax errors.
• Updated README and submission.json to document the new opt-in flags and keep baseline-reported metrics unchanged; those documentation files were validated for consistency via simple file checks.

Codex Task