Asankhaya Sharma

I think the data mixing should carry over to larger models, there is existing work from others that suggests this, e.g. - https://www.datologyai.com/blog/beyondweb

For evaluation we used lm-evaluation-harness with a custom wrapper to handle diffusion-specific probability calculations for multiple choice tasks.

For inference we used standard Transformers library. The diffusion models use a custom generate() method that handles parallel token generation with configurable diffusion steps. Throughput was measured with batch size 1, generating 100 tokens per prompt averaged over multiple runs.

I ran the numbers on layer-only params (excluding embeddings):

The 48L config has the most layer params (59M) but is in the Low tier, while 12L has fewer (37.7M) and is High tier.

The hidden dimension threshold still dominates. But er-layer representation width seems critical, with hidden=320 or 256, you create an information bottleneck that more layers can't overcome, unless you hit the critical depth thresholds (32 or 64 layers) where something else compensates.

This suggests the finding should be reframed as: at small scale, you need sufficient hidden dimension AND appropriate depth.

(BTW, based on your earlier comment I've added a note to the article clarifying the parameter matching limitations — thanks for the feedback!)

Here's the full breakdown of where parameters come from:

1. Embeddings (scales linearly with d_model)

• Token embeddings: vocab_size × d_model = 50,257 × d
• Position embeddings: 1,024 × d
• Total: ~51,281 × d

2. Per transformer layer (scales quadratically with d_model)

• Attention (Q, K, V, O): 4 × d²
• MLP (up + down, with 4x intermediate): 2 × d × 4d = 8d²
• LayerNorms: ~4d (negligible)
• Total per layer: ~12d²

3. LM Head

• Usually tied with embeddings (free) or d × vocab_size

4L × 768:

• Embeddings: 51,281 × 768 ≈ 39.4M
• Layers: 4 × 12 × 768² ≈ 28.3M
• Total: ~68M

12L × 512:

• Embeddings: 51,281 × 512 ≈ 26.3M
• Layers: 12 × 12 × 512² ≈ 37.7M
• Total: ~64M

Thanks for the references I will take a look.