Sampled Softmax with Over-Tokenized Transformer

Setup设置

The baseline is the MuonH Qwen3 recipe of the Marin speedrun, which covers 130m to 1.2B parameters. We replace the softmax of the training loss with a sampled softmax and add the input and output vocabularies of the Over-Tokenized Transformer, or OT; each setup is judged by its training speed and by its Paloma loss against backbone compute 6ND, where N counts the parameters of the transformer layers without the embedding and the heads, and D the training tokens. Code

基线是 Marin speedrun 的 MuonH Qwen3 配方，覆盖 130m 到 1.2B 参数。我们把训练损失中的 softmax 换成采样 softmax，并在输入端和输出端加上 Over-Tokenized Transformer 的输入和输出词表，简称 OT；每种做法看训练速度，以及 Paloma 损失随骨干算力 6ND 的变化，其中 N 是 transformer 各层的参数量，不含 embedding 和输出头，D 是训练 token 数。代码

--- baseline loss (simplified)+++ OT + sampled softmax (simplified; the Code link has the real code) def loss_fn(model, tokens):     targets = tokens[1:]-    x = model.embed(tokens[:-1])+    # OT, input: add the hashed 2-gram and 3-gram embeddings that end at each token+    x = model.embed(tokens[:-1]) + model.ngram_embed(tokens[:-1], n=2) + model.ngram_embed(tokens[:-1], n=3)     h = model.norm(model.blocks(x))-    logits = h @ model.head.weight.T-    return cross_entropy(logits, targets)+    # sampled softmax: score this GPU's targets plus random negatives, P tokens in all+    cand = targets.unique() + random_tokens(P - len(targets.unique()))+    logits = h @ model.head.weight[cand].T+    loss = cross_entropy(logits, position_in(cand, targets))+    # OT, output: a second head predicts the token two positions ahead, at weight 0.1+    loss += 0.1 * cross_entropy(model.head2(h[:-1]), tokens[2:])+    return loss

--- 基线损失（简化）+++ OT + 采样 softmax（简化；真实代码见“代码”链接） def loss_fn(model, tokens):     targets = tokens[1:]-    x = model.embed(tokens[:-1])+    # OT 输入端：加上以每个词结尾的 2-gram 和 3-gram 的 hash embedding+    x = model.embed(tokens[:-1]) + model.ngram_embed(tokens[:-1], n=2) + model.ngram_embed(tokens[:-1], n=3)     h = model.norm(model.blocks(x))-    logits = h @ model.head.weight.T-    return cross_entropy(logits, targets)+    # 采样 softmax：只给本 GPU 的目标词和随机负样本打分，共 P 个词+    cand = targets.unique() + random_tokens(P - len(targets.unique()))+    logits = h @ model.head.weight[cand].T+    loss = cross_entropy(logits, position_in(cand, targets))+    # OT 输出端：第二个头预测往后第二个位置的词，权重 0.1+    loss += 0.1 * cross_entropy(model.head2(h[:-1]), tokens[2:])+    return loss

Experiments实验

Takeaways核心结论

Related work相关工作

• Over-Tokenized Transformer, Huang et al. 2025 separates the input and output vocabularies and finds that a larger input vocabulary helps at every model size it tests. We train both of its vocabularies on the Marin recipe: they help up to 520m and hurt at 1.2B, the one size that repeats its data.
• Jean et al. 2015 train translation models with a softmax over a sampled part of a very large target vocabulary. We use the same idea for speed in pretraining and return to the full softmax before training ends.
• Record #92 of the nanoGPT speedrun uses a sampled softmax in the early stages of training, together with a hashed n-gram table close to the input side of OT. We take its schedule and test it from 130m to 1.2B.
• MiLe loss, Su et al. 2023 reweights each token by the entropy of the prediction and compares with focal loss in language-model pretraining, where focal loss improves downstream accuracy slightly over cross-entropy. In our runs, focal loss raises the Paloma loss of the baseline and lowers that of OT by less than 0.01 nats.
• Scaling Data-Constrained Language Models, Muennighoff et al. 2023 find that up to 4 epochs of repeated data change the loss little compared with unique data. Our 1.2B runs make 2.4 passes, and OT turns worse as each new pass begins, so the n-gram tables may make a model more sensitive to repetition.

• Over-Tokenized Transformer，Huang et al. 2025 把输入词表和输出词表分开，发现更大的输入词表在它测试的每个规模上都有帮助。我们在 Marin 配方上训练它的两种词表：在 520m 及以下有帮助，在 1.2B 反而有害，而 1.2B 是唯一把大部分数据重复使用的规模。
• Jean et al. 2015 训练翻译模型时，只对很大目标词表中采样出的一部分算 softmax。我们在预训练里用同样的思路来提速，并在训练结束前换回完整 softmax。
• nanoGPT speedrun 第 92 号记录 在训练的前几个阶段用采样 softmax，同时用一张 hash n-gram 表，和 OT 的输入端很接近。我们沿用它的采样计划，并从 130m 到 1.2B 检验它。
• MiLe loss，Su et al. 2023 按预测分布的熵给每个词重新加权，并在语言模型预训练中和 focal loss 比较；在那里 focal loss 的下游准确率比交叉熵略高。在我们的实验中，focal loss 让 baseline 的 Paloma 损失升高，让 OT 的损失降低不到 0.01 nats。
• Scaling Data-Constrained Language Models，Muennighoff et al. 2023 发现，重复数据在 4 个 epoch 以内，损失和用不重复数据时相差很小。我们的 1.2B 训练了 2.4 遍，而 OT 在每一遍开始时都变差，说明 n-gram 表可能让模型对重复更敏感。

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{chai2026ssov,
  title        = {Sampled Softmax with Over-Tokenized Transformer},
  author       = {Chai, Wenhao},
  year         = {2026},
  howpublished = {Blog post},
  url          = {https://wenhaochai.com/blogs/sampled-softmax-ov.html}
}