Empirical Capacity Model for Self-Attention Neural Networks
Aki Härmä, Marcin Pietrasik, Anna Wilbik
TL;DR
The paper addresses how much information a self-attention transformer can memorize under common optimization with synthetic data. It introduces an interpretable Empirical Capacity Model (ECM) that predicts memorization capacity $C$ as $C = \max( f(H,N)\cdot B, \alpha H + \beta )$ with $f(N,H) = a /(N^{bH+c}+d) + e$, capturing a presaturation linear rise in $B$ and a saturation bound governed by $H$. Through large-scale experiments varying token-vector size $B$, heads $H$, and sequence length $N$, the authors show that the ECM fits measured capacity better than a high-order polynomial and provides a practical, low-parameter tool for a priori hyperparameter selection and architecture design, including a memory-per-parameter metric. The work has implications for designing memory-efficient transformers and Retrieval Augmented Generation systems by linking architecture choices to attainable memorization, with future work extending to denser hyperparameter regimes and natural language data.
Abstract
Large pretrained self-attention neural networks, or transformers, have been very successful in various tasks recently. The performance of a model on a given task depends on its ability to memorize and generalize the training data. Large transformer models, which may have billions of parameters, in theory have a huge capacity to memorize content. However, the current algorithms for the optimization fall short of the theoretical capacity, and the capacity is also highly dependent on the content. In this paper, we focus on the memory capacity of these models obtained using common training algorithms and synthetic training data. Based on the results, we derive an empirical capacity model (ECM) for a generic transformer. The ECM can be used to design task-specific transformer models with an optimal number of parameters in cases where the target memorization capability of the task can be defined.