Token Averaging for Compute-Efficient Language Model Pretraining

Training language models on more text ordinarily requires proportionally more compute. We introduce token averaging, a parameter-free method that replaces each group of $k$ consecutive token embeddings with their mean before the first attention layer. The transformer therefore processes $k$ times fewer positions while predicting the first token of the next group. This allows a fixed training budget to cover approximately $k$ times more raw text; at $k=2$, the model trains on twice as much text using $5$-$7\%$ fewer FLOPs than the baseline (without token averaging). Since token embeddings are averaged, we randomize group boundaries during training and evaluate using $k$ different boundary offsets, ensuring that the token-averaged model and the baseline are scored on identical target positions. In each offset, a sequence $k$ times shorter is processed, so the cost of all $k$ offset evaluations is approximately equal to one forward pass for the baseline. At inference, each pass stores $1/k$ as many KV-cache entries and incurs $1/k^2$ of the attention cost compared with the baseline, i.e., half the cache and one-quarter of the attention cost at $k=2$. Across models from 50M to 500M parameters, our experiments show that the baseline requires $1.30\times$, $1.28\times$, $1.17\times$, and $1.08\times$ more compute, respectively, to reach the same loss as the $k=2$ model. Full-position evaluation shows $3.4\%$ lower perplexity at 125M with $7\%$ fewer training FLOPs and comparable loss at 500M with approximately $5\%$ fewer FLOPs. Additional experiments show that learned pooling over token content and within-group position outperforms mean pooling by $0.043$ nats at $k=2$ and $0.721$ nats at $k=8$. Together, these results show that token averaging increases both training-data throughput and inference efficiency without compromising model quality.