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. OT adds hashed 2-gram and 3-gram embeddings to each input token, and we test two forms of its output side: OT-H (head), the paper's form, a second head that predicts the token two positions ahead; and OT-2 (2-gram), a real output vocabulary of hashed 2-grams. 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。OT 给每个输入词加上 2-gram 和 3-gram 的 hash embedding；输出端我们测试两种形式：OT-H（head），即论文的形式，用第二个头预测往后第二个位置的词；OT-2（2-gram），真正的 hash 2-gram 输出词表。每种做法看训练速度，以及 Paloma 损失随骨干算力 6ND 的变化，其中 N 是 transformer 各层的参数量，不含 embedding 和输出头，D 是训练 token 数。代码

--- baseline loss (simplified)+++ OT-H + 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-H, 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-H + 采样 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-H 输出端：第二个头预测往后第二个位置的词，权重 0.1+    loss += 0.1 * cross_entropy(model.head2(h[:-1]), tokens[2:])+    return loss

Experiments实验

Answered已有结论Do the sampled softmax and OT-H make pretraining cheaper or better?采样 softmax 和 OT-H 能让预训练更省或更好吗？

The sampled softmax trains faster at every size and ends level with the baseline; OT-H lowers the loss up to 520m and raises it at 1.2B, the one size that repeats most of its data.

采样 softmax 在每个规模上都训练得更快，最终损失与基线持平；OT-H 在 520m 及以下降低损失，在 1.2B 上反而升高损失，而 1.2B 是唯一把大部分数据重复使用的规模。

Answered已有结论Does a change to the training loss help OT-H further?改训练损失能让 OT-H 更好吗？

Not much: focal loss changes the loss of OT-H little, and the sampled softmax helps OT-H only at 130m, not at the larger sizes.

帮助不大：focal loss 对 OT-H 的损失影响很小；采样 softmax 只在 130m 上让 OT-H 更好，在更大的规模上没有帮助。

Answered已有结论Why does OT-H lose at 1.2B?OT-H 为什么在 1.2B 上变差？

Because the data repeats: at 300m with 2.4 passes OT-H falls behind after each new pass, the large n-gram input tables cause it, and tables ten times smaller, or frozen after the first pass, shrink the damage.

因为数据重复：在 300m 上训练 2.4 遍时，每开始新的一遍，OT-H 都会落后；原因是大的 n-gram 输入表；把表缩小 10 倍，或在第一遍后冻结表，都能减轻损失。

Answered已有结论Does a real 2-gram output vocabulary beat the head of OT-H?真正的 2-gram 输出词表能胜过 OT-H 的头吗？

Weakly: at 300m OT-2 ends ahead of OT-H with and without the sampled softmax, by a margin close to the noise of single runs; at 130m the four setups end within noise of each other.

证据较弱：在 300m 上，无论是否用采样 softmax，OT-2 最终都领先 OT-H，但领先的幅度接近单个 run 的噪声；在 130m 上，四种做法的最终结果都在噪声范围内。

Takeaways核心结论

A sampled softmax trains the Marin recipe 1.07–1.30× faster from 130m to 1.2B at no cost in loss; the Over-Tokenized Transformer's input tables lower the loss up to 520m but overfit repeated data, which turns the gain into a loss at 1.2B, and a real 2-gram output head adds at most a small gain.

采样 softmax 让 Marin 配方在 130m 到 1.2B 上训练快 1.07–1.30 倍，损失不变；Over-Tokenized Transformer 的输入表在 520m 及以下降低损失，但会在重复数据上过拟合，使 1.2B 上的增益变成损失；真正的 2-gram 输出头最多只带来很小的增益。

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 its form, OT-H, on the Marin recipe: it helps up to 520m and hurts at 1.2B, the one size that repeats most of its data; OT-2 replaces its second head with a real 2-gram output vocabulary, which the paper avoids for its cost.
• 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.
• Zhang et al. 2022 find that click-through-rate models with sparse ID features lose accuracy sharply at the start of each new epoch when trained with fast optimizers such as Adam, and that hashing the IDs into far fewer rows removes the drop at some cost in accuracy. OT-H shows the same jump at each new pass, and its 10 times smaller tables barely suffer from repetition.
• Jha et al. 2026 find that mixture-of-experts language models degrade more than dense ones under repeated data, as set by total rather than active parameters. The n-gram tables of OT-H at 300m hold 6.5B parameters, more than the 2.5B tokens they repeat.

• Over-Tokenized Transformer，Huang et al. 2025 把输入词表和输出词表分开，发现更大的输入词表在它测试的每个规模上都有帮助。我们在 Marin 配方上训练它的形式 OT-H：在 520m 及以下有帮助，在 1.2B 反而有害，而 1.2B 是唯一把大部分数据重复使用的规模；OT-2 把它的第二个头换成真正的 2-gram 输出词表，论文因为代价高而没有这样做。
• Jean et al. 2015 训练翻译模型时，只对很大目标词表中采样出的一部分算 softmax。我们在预训练里用同样的思路来提速，并在训练结束前换回完整 softmax。
• nanoGPT speedrun 第 92 号记录 在训练的前几个阶段用采样 softmax，同时用一张 hash n-gram 表，和 OT 的输入端很接近。我们沿用它的采样计划，并从 130m 到 1.2B 检验它。
• Zhang et al. 2022 发现，ID 特征稀疏、用 Adam 这类快速优化器训练的点击率模型，在每个新 epoch 开始时准确率都会急剧下降；把 ID hash 到少得多的行可以消除这个下降，但会损失一些准确率。OT-H 在每一遍开始时也出现同样的跳变，而行数少 10 倍的表几乎不受重复数据影响。
• Jha et al. 2026 发现，在重复数据上，mixture-of-experts 语言模型比稠密模型退化得更多，决定退化程度的是总参数量而不是激活参数量。300m 的 OT-H 的 n-gram 表有 6.5B 参数，比被重复的 2.5B token 还多。

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 等团队组成的联合体。

Experiment record实验记录

Citation

@misc{chai2026ssov,
  title        = {Sampled Softmax with Over-Tokenized Transformer},
  author       = {Chai, Wenhao},
  year         = {2026},
  howpublished = {Blog post},
  url          = {https://wenhaochai.com/exps/sampled-softmax-ov.html}
}