Quantifying the Per-layer Contribution of LLMs

Setup设置

The baseline is the MuonH Qwen3 recipe of the Marin speedrun, from 130m to 1.2B. We add a next-token loss after every layer and judge each layer by its final training loss against backbone compute 6ND, where N counts the parameters of the layers up to that one, without the embedding and the heads, and D the training tokens. Code

基线是 Marin speedrun 里的 MuonH Qwen3 配方，规模从 130m 到 1.2B。我们在每一层后面加一个预测下一个词的损失，用每层的最终训练损失来衡量；横轴是骨干算力 6ND，N 是到这一层为止各层的参数量，不算 embedding 和头，D 是训练 token 数。代码

--- baseline loss (simplified)+++ per-layer losses (simplified; the Code link has the real code) def loss_fn(model, tokens):     targets = tokens[1:]     x = model.embed(tokens[:-1])-    for block in model.blocks:+    loss = 0.0+    for k, block in enumerate(model.blocks[:-1]):         x = block(x)-    return cross_entropy(model.head(model.norm(x)), targets)+        # probes: read the layer, but send no gradient into the model+        h = stop_gradient(x) if probe else x+        # every earlier layer has its own norm and output head, and predicts the next token+        logits = model.layer_heads[k](model.layer_norms[k](h))+        loss += cross_entropy(logits, targets)+    x = model.blocks[-1](x)+    # the usual loss at the last layer, at the same weight+    return loss + cross_entropy(model.head(model.norm(x)), targets)

--- 基线损失（简化）+++ 逐层损失（简化；真实代码见“代码”链接） def loss_fn(model, tokens):     targets = tokens[1:]     x = model.embed(tokens[:-1])-    for block in model.blocks:+    loss = 0.0+    for k, block in enumerate(model.blocks[:-1]):         x = block(x)-    return cross_entropy(model.head(model.norm(x)), targets)+        # 探针：读取这一层，但梯度不传回模型+        h = stop_gradient(x) if probe else x+        # 前面每一层都有自己的 norm 和输出头，也预测下一个词+        logits = model.layer_heads[k](model.layer_norms[k](h))+        loss += cross_entropy(logits, targets)+    x = model.blocks[-1](x)+    # 最后一层仍是原来的损失，权重相同+    return loss + cross_entropy(model.head(model.norm(x)), targets)

Experiments实验

Takeaways核心结论

Related work相关工作

• Deeply-Supervised Nets, Lee et al. 2015 add a classifier loss at every hidden layer of an image network and find that it helps. We put the same kind of loss on a small language model: each layer gets closer to a shallower model, and the last layer gets worse.
• The logit lens, nostalgebraist 2020 and the tuned lens, Belrose et al. 2023 read the middle layers of a trained model, through its own output head or a small learned map per layer. Our probes do the same while the model trains, and they agree: the prediction forms in the last layers.
• LayerSkip, Elhoushi et al. 2024 trains early layers to predict through the shared output head, so that inference can stop early. Our shared-head setup is that training signal; here it costs the last layer about 0.13 nats.
• LLAL, Jin et al. 2026 adds a language-model loss to the layers of an ultra-sparse mixture-of-experts model to stop its experts from dying. We test the loss on a dense model, which has no such failure to fix.

• Deeply-Supervised Nets，Lee et al. 2015 在图像网络的每个隐藏层加分类损失，发现有帮助。我们在小语言模型上加同类损失：每层更接近一个更浅的模型，最后一层却变差。
• logit lens，nostalgebraist 2020 和 tuned lens，Belrose et al. 2023 读取训练好的模型的中间层，用模型自己的输出头，或每层学一个小映射。我们的探针在训练过程中做同样的事，结论一致：预测在最后几层才形成。
• LayerSkip，Elhoushi et al. 2024 让前面的层也通过共享输出头预测，这样推理时可以提前停下。我们的共享头做法就是这种训练信号；在这里它让最后一层多付约 0.13 nats。
• LLAL，Jin et al. 2026 在超稀疏混合专家模型的层上加语言模型损失，防止专家失活。我们在 dense 模型上检验这个损失，而 dense 模型没有这个问题可修。

Acknowledgements致谢

The author is pleased to acknowledge that the work reported on in this post was substantially performed using the Princeton Research Computing resources at Princeton University. Princeton Research Computing is a consortium of groups including the Princeton Institute for Computational Science and Engineering (PICSciE) and Research Computing at Princeton University.

本文报告的工作主要使用普林斯顿大学 Princeton Research Computing 的计算资源完成。Princeton Research Computing 是由普林斯顿计算科学与工程研究所（PICSciE）和普林斯顿大学 Research Computing 等团队组成的联合体。

Citation

@misc{chai2026perlayer,
  title        = {Quantifying the Per-layer Contribution of LLMs},
  author       = {Chai, Wenhao},
  year         = {2026},
  howpublished = {Blog post},
  url          = {https://wenhaochai.com/blogs/per-layer-contribution.html}
}