Revisiting the Shape Convention of Transformer Language Models

Paper: 2602.06471 Authors: MediaTek Research Categories: cs.CL

Abstract

This paper re-examines the tensor shape conventions used in transformer language model architectures. The standard transformer uses specific dimension ratios between model width, FFN hidden size, and attention head dimensions that have been largely inherited from the original “Attention Is All You Need” paper without systematic justification. This work investigates whether alternative shape configurations can yield better efficiency-performance tradeoffs.

Key Contributions

• Systematic analysis: Rigorous examination of how dimension ratios affect transformer performance
• Alternative conventions: Identifies configurations that deviate from standard practice but yield improvements
• Practical implications: Guidance for architecture design decisions in new model development

The Shape Convention Problem

The original Transformer paper established dimension ratios that have been widely adopted: model dimension d_model, FFN hidden size 4×d_model, and attention head dimension d_model/num_heads. These ratios have been inherited across GPT, LLaMA, Qwen, and most modern architectures with minimal questioning.

But are these ratios optimal? Or are they historical artifacts that persist through convention?

Key Findings

The paper systematically varies dimension ratios and evaluates the impact on:

• Model quality: Perplexity and downstream task performance
• Training efficiency: Convergence speed and compute utilization
• Inference efficiency: Throughput and latency characteristics

The results suggest that alternative configurations can yield meaningful improvements, particularly when accounting for modern hardware characteristics that differ from the GPU landscape when transformers were first designed.

Practical Implications

• Architecture designers should treat dimension ratios as hyperparameters to optimize, not fixed constants
• The optimal ratios may depend on the target hardware and deployment scenario
• Small changes in shape conventions can compound across many layers

Takeaways

1. Question conventions: Inherited design choices deserve periodic re-evaluation
2. Hardware co-design: Optimal model shapes depend on the target compute platform
3. Low-hanging fruit: Shape optimization is a simple lever that’s often overlooked