Does a Wider Embedding Space Help More on Diverse Data?

TL;DR

• Increasing the embedding dimension from 32 to 64 improved validation loss in both the low- and high-diversity training conditions.

• Most of the improvement appeared between 32 and 48 dimensions; the gain from 48 to 64 was smaller.

• Training on diverse data produced better performance on a shared, multi-topic validation set at every model width.

• At the aggregate level, width did not help substantially more on diverse data. However, the per-cluster results showed that wider models improved more consistently across topics after diverse training.

• In this setup, the 48-dimensional model looked like a reasonable quality-versus-compute balance.

GitHub: View the complete code and experiment notebooks

In the previous experiment, I explored how model size, depth, and compute affected training.

This time, I wanted to isolate a narrower question:

Does increasing the embedding dimension improve next-token prediction—and does it help more when the training data contains a broader range of topics?

My intuition was that a larger embedding space would give the model more room to represent relationships between words, contexts, and topics.

But there was an important caveat: in my GPT implementation, emb_dim does more than control the initial token embeddings. It also determines the hidden width used throughout the transformer.

So this was really a model-width experiment, not a perfectly isolated test of the embedding table.

Experiment setup

I trained models using three embedding dimensions:

Everything else remained fixed:

• four attention heads;

• four transformer blocks;

• context length of 128;

• the same optimizer and learning-rate schedule;

• 200,192 training tokens per run.

I also created two token-matched WikiText corpora:

• Low diversity: articles selected from one topic cluster.

• High diversity: articles sampled across six topic clusters.

Every model was evaluated on its matched validation set and on a shared, diverse validation set.

Research on model scaling has shown that model capacity, data, and compute interact rather than operating independently. Other work has also explored separating embedding size from hidden-layer capacity and the importance of pretraining-data composition.

Hypothesis

I expected wider models to achieve lower validation loss because they could represent richer relationships.

I also expected the benefit of width to be larger for the high-diversity corpus, since the model needed to represent a wider range of topics and vocabulary patterns.

The first part of the hypothesis held up.

The second was more complicated.

Wider models achieved lower loss

Validation loss improved as embedding dimension increased in both conditions.

The gains were not linear.

For the low-diversity corpus:

• 32 → 48 improved loss by 0.1375

• 48 → 64 improved it by 0.0384

For the high-diversity corpus:

• 32 → 48 improved loss by 0.0502

• 48 → 64 improved it by 0.0273

So width helped, but the returns diminished.

The 48-dimensional model captured much of the improvement while using fewer parameters and less compute than the 64-dimensional model.

Top-1 accuracy also peaked at 48 dimensions in both matched-validation comparisons, although validation loss continued to improve at 64 dimensions.

That is possible because accuracy only checks whether the highest-ranked token was correct, while cross-entropy loss evaluates the full predicted probability distribution.

Diverse training improved generalization

I evaluated every model on the same shared, diverse validation set.

At every width, the models trained on diverse data performed better on the shared diverse evaluation.

This suggests a fairly intuitive pattern:

The low-diversity models became better specialists, while the high-diversity models became better generalists.

The Pile and later data-mixture research similarly emphasize that the composition and coverage of pretraining data can meaningfully affect cross-domain performance.

Did width help more on diverse data?

At the aggregate level, not really.

On the shared validation set, increasing width from 32 to 64 improved loss by:

• 0.0879 after low-diversity training;

• 0.0775 after high-diversity training.

Those improvements were similar. The overall validation score therefore did not support a strong interaction between width and topic diversity.

But the per-cluster results added useful nuance.

Per-cluster results

After high-diversity training, increasing width from 32 to 64 improved loss in all six clusters.

After low-diversity training, width improved four clusters, barely changed one, and slightly worsened another.

The effect varied substantially across topics.

For the high-diversity models:

• Cluster 5 improved by 0.1655

• Cluster 2 improved by 0.1520

• Cluster 1 improved by 0.1345

• Cluster 0 improved by only 0.0280

So additional width did not benefit every topic equally.

The low-diversity model was especially strong on one cluster but considerably weaker on most of the others. The high-diversity model performed better on five of six clusters at 64 dimensions.

That supports a more nuanced version of the original hypothesis:

Wider representations did not create a much larger aggregate benefit on diverse data, but they produced broader and more consistent gains across individual topics.

Quality versus cost

Increasing emb_dim from 32 to 64:

• roughly doubled model parameters;

• roughly doubled estimated training compute;

• reduced training throughput;

• produced progressively smaller improvements in loss.

That makes 48 dimensions look like a reasonable efficiency point in this particular setup.

It captured most of the benefit from increasing width, achieved the highest matched top-1 accuracy, and avoided the full runtime and compute cost of the 64-dimensional model.

What I learned

This experiment produced three main observations:

1. Wider models achieved lower validation loss.

2. Diverse training data improved broad generalization.

3. Width helped more consistently across topic clusters after diverse training, even though the aggregate interaction was weak.

My takeaway is narrower than my original hypothesis:

Wider representations helped, and diverse training helped—but they did not strongly amplify each other in the overall validation result.

The more interesting effect appeared at the topic level.

The low-diversity model behaved more like a specialist. The high-diversity model behaved more like a generalist. Additional width improved that generalist across a broader range of topics.

As with my previous experiments, these findings are specific to my model, clustering method, dataset, learning rate, and training budget.

But they offered a useful reminder:

Aggregate metrics can hide where an architectural change is actually helping.

References

1. Vaswani, A. et al. (2017). Attention Is All You Need

2. Lan, Z. et al. (2019). ALBERT: A Lite BERT for Self-supervised Learning of Language Representations

3. Kaplan, J. et al. (2020). Scaling Laws for Neural Language Models

4. Hoffmann, J. et al. (2022). Training Compute-Optimal Large Language Models

5. Gao, L. et al. (2020). The Pile: An 800GB Dataset of Diverse Text for Language Modeling

6. Xie, S. M. et al. (2023). DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining