Allocation of Parameters in Transformers
This work addresses model efficiency for Transformer-based architectures, providing theoretical insights that could guide parameter allocation in AI applications, though it is incremental as it builds on existing Transformer frameworks.
The paper tackles the problem of how to allocate parameters, specifically attention heads and head dimensions, across Transformer layers to balance expressivity and efficiency, uncovering a saturation behavior in softmax activations that leads to diminishing returns with increased head dimensions, particularly for long sequences, and proposes principled allocation strategies.
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.