Allocation of Parameters in Transformers

2025-10-04

AI Summary

In Transformers, uneven allocation of parameters—specifically attention head count and head dimension—across layers leads to suboptimal computational efficiency. Method: We propose a cross-layer differentiated parameterization strategy. Theoretically, we analyze information flow via function approximation and first reveal Softmax saturation in deeper layers, proving diminishing returns from increasing head dimension; this yields a principled trade-off model between head count and head dimension under fixed parameter budgets. Contribution/Results: Experiments show early layers benefit from more heads for enhanced feature extraction, while later layers tolerate substantial reductions in both head count and head dimension without performance degradation. Our strategy reduces computational cost by 15–25% while preserving accuracy, providing an interpretable theoretical foundation and practical configuration guidelines for efficient large-scale Transformer design.

Technology Category

• Computer Vision: Large Vision Models
• Machine Learning: Mixture of Experts (MoE)
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: Practical large-scale studies of user experience
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

Transformers have achieved remarkable successes across a wide range of applications, yet the theoretical foundation of their model efficiency remains underexplored. In this work, we investigate how the model parameters -- mainly attention heads and head dimensions -- should be allocated across layers to balance expressivity and efficiency. We first provide mathematical analysis on the role of early layers in information extraction from an approximation perspective, with a theoretical characterization on the trade-off between the number of heads and head dimension under a fixed parameter budget. In addition, we uncover and prove the emph{saturation} behavior of softmax activations: Continuously increasing head dimensions can lead to diminishing returns in learning errors, particularly for long sequences. Supported by both theory and experiments, this saturation pattern suggests that later layers can operate more efficiently with reduced parameters. Combining these insights, we propose principled strategies for allocating attention heads and dimensions across Transformers' layers, shedding light on theoretically-grounded model efficiency of Transformer-based architectures.

Problem

• Optimizing attention head allocation across Transformer layers
• Analyzing head dimension saturation effects on sequence learning
• Balancing parameter efficiency with model expressivity theoretically

Innovation

• Mathematical analysis of early layer information extraction
• Proved saturation behavior of softmax activations
• Proposed principled parameter allocation strategies across layers