A brief recap. In our previous post, we introduced LLAL—Language Modeling Loss as an Auxiliary Loss—to address weak learning signals and expert undertraining in early MoE layers. Our hypothesis was that giving early layers a stable learning signal could help them establish useful functions. LLAL applies an auxiliary language-modeling objective to an early layer for a short initial training window and then removes it. In the settings reported there, the additional training computation was less than 0.3%. We observed:
LLAL applies an auxiliary prediction objective early in training.
What lasting function does LLAL help these experts establish—and how might that explain its broader training benefits? This follow-up investigates the hypothesis behind LLAL through rare-knowledge prediction and layer-wise expert masking. We then use Engram as a complementary architectural intervention to examine whether making a useful knowledge-related function available early supports the development of subsequent experts. The goal throughout is to better understand how LLAL may work.
A possible explanation of LLAL: its early auxiliary signal helps previously underused experts establish a persistent role in supporting rare knowledge. Subsequent layers may then build on that function and learn useful predictive computations without all specializing in acquiring the same knowledge. The Engram placement experiments add support to this interpretation: providing memory before L1 is accompanied by healthier expert development and broadly positive masking effects, whereas providing it before L14 aggravates early-layer undertraining.
Conventions. Layer numbers are zero-based. In the main comparison models, L0 is dense and L1 is the first MoE layer. We probe rare knowledge through low-frequency name continuations and factual answers, examining both knowledge-related dependence and the predictive contributions of the surrounding layers.
Our full-MoE experiments motivated this analysis. When the first layer was also an MoE layer, it did not exhibit the same weak-signal pattern and its weight norms appeared healthy, while other early MoE layers continued to receive weak signals as training progressed. This raised a functional question: what useful role can an early MoE layer acquire, and can establishing that role change the development of the layers that follow?
Earlier reports of bigram-prediction subnetworks and token- or n-gram-sensitive neurons in early feed-forward layers (Bigram Subnetworks and Neurons in Large Language Models) motivated an examination of memory-related predictions. Our aim is not only to locate a rare-knowledge dependence peak, but to understand how an early useful function relates to the broader layer-wise pattern.
The main comparisons shown below are AdamW 2T and Muon models, each with and without LLAL. We also include two Engram models, with the memory module inserted at different depths.
| Task | Dataset | Prediction target |
|---|---|---|
| Entity-name completion | PopQA names: 6,417 | Remaining words of an entity name, given its first word |
| Factual recall | PopQA: 11,822 | Answer span in a question-answering or cloze prompt |
We use subject-name occurrence counts in FineWeb as a proxy for rarity, not as a measurement of exposure in the models’ actual training data. Rare subjects have at most 100 occurrences; common subjects have more than 1,000. Intermediate counts remain visible in the frequency-binned analysis.
For each tested layer, we zeroed the routed-expert output, retaining the shared expert, attention, and residual connection. This is a branch intervention, not deletion of the entire transformer layer. We scored the canonical target using teacher forcing and measured
$$ \Delta \mathrm{NLL} = \mathrm{NLL}{\mathrm{ablated}} - \mathrm{NLL}{\mathrm{original}}. $$
Positive ΔNLL means that the branch was making a useful contribution to the target prediction: removing it made the target less likely. Values are in nats per target span. Negative values mean that removing the branch improved the target score under that condition.
We read these measurements in two ways. Absolute ΔNLL indicates a branch’s predictive contribution on the evaluated task. The rare/common contrast indicates whether that contribution is especially important for rare knowledge. A layer can contribute positively without showing strong rare-knowledge specialization.