TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters

## Abstract

Abstract Transformers have become the predominant architecture in foundation models due to their excellent performance across various domains. However, the substantial cost of scaling these models remains a significant concern. This problem arises primarily from their dependence on a fixed number of parameters within linear projections. When architectural modifications ( e.g. , channel dimensions) are introduced, the entire model typically requires retraining from scratch. As model sizes continue growing, this strategy results in increasingly high computational costs and becomes unsustainable. To overcome this problem, we introduce \ourmethod , a natively scalable architecture that leverages the attention mechanism not only for computations among input tokens but also for interactions between tokens and model parameters, thereby enhancing architectural flexibility. By treating model parameters as tokens, we replace all the linear projections in Transformers with our token-parameter attention layer, where input tokens act as queries and model parameters as keys and values. This reformulation allows for progressive and efficient scaling without necessitating retraining from scratch. Our model scales from 124M to 1.4B parameters by incrementally adding new key-value parameter pairs, achieving performance comparable to Transformers trained from scratch while greatly reducing training costs. Code and models are available at https://github.com/Haiyang-W/TokenFormer .

## 1 Introduction

Designing a powerful neural network architecture is a long-standing goal in machine learning. Recent developments in foundation models (FMs) have shown the potential of Transformers (Vaswani et al., ) as a universal computational architecture. Thanks to their flexibility and scalability, Transformers have achieved state-of-the-art performance across various domains, including natural language processing (NLP) (Radford et al., ; Alec et al., ; Brown et al., ), visual modeling (Dosovitskiy et al., ; Liu et al., ), vision-language (Liu et al., ; Wang et al., ), graph representation (Ying et al., ), and 3D vision (Wang et al., ; ).

Figure: Figure 1: Traditionally, large transformer architectures are trained from scratch without reusing previous smaller-scale models (represented by blue dots on the left). In this paper, we propose a novel fully attention-based architecture that allows scaling model incrementally, thus greatly reducing the overall cost of training large transformer architectures (depicted by red dots on the left). The right panel delineates a comparison between conventional Transformer and our Tokenformer.
Refer to caption: x1.png

Transformers typically divide the computation required to process a single token into two distinct parts: interactions with other input tokens (token-token interaction) and computations involving the model’s parameters (token-parameter interaction). The attention mechanism (Vaswani et al., ) facilitates token-token interactions, allowing modern general-purpose foundation models to encode multi-modal data into a unified token sequence and effectively capture complex dependencies among them (Liu et al., ; Zhu et al., ; Wang et al., ).
Conversely, token-parameter computations rely heavily on linear projections (Dunford & Schwartz, ), where input tokens are multiplied by a fixed set of parameters.
This prescribed design limits scalability because increasing the model size requires altering core architectural components, often necessitating retraining the entire model from scratch.
As models grow larger, this results in excessive resource consumption, making it increasingly impractical.
In this paper, we introduce a novel architecture that enhances the flexibility of token-parameter interactions, allowing for incremental scaling of model parameters and effectively reusing previously trained models, thus significantly reducing the training burden.

To achieve this objective, we introduce \ourmethod, a novel architecture that unifies the computations of token-token and token-parameter interactions by entirely employing the attention mechanism. The flexibility of our token-parameter attention layer, along with its ability to handle a variable number of parameters, inherently enhances the model’s scalability, facilitating progressively efficient scaling.

As shown in Figure 1, we extend the Transformer architecture by preserving the computational patterns between input tokens while reformulating all the linear projections using a cross-attention mechanism. Specifically, to project features with input and output dimensions $D_{1}$ and $D_{2}$, we employ two sets of parameters, each comprising $N$ learnable tokens with channel dimensions of $D_{1}$ and $D_{2}$, respectively. In this formulation, input tokens serve as queries, and model parameters as keys and values. This flexibility renders our model’s parameters inherently scalable with variable $N$, allowing for efficient expansion by continuously adding new key-value parameter pairs. Figure 1 shows that our model can be scaled incrementally from 124M to 1.4B parameters, achieving performance similar to training from scratch while saving more than half of the training cost.

The key contributions of this work are summarized as 1) As shown in Figure 1, we propose \ourmethod, a fully attention-driven neural network that treats model parameters as tokens, maximizing the flexibility of token-parameter computations while achieving competitive performance on standard benchmarks across both language and vision domains. 2) Thanks to this design, our model can be naturally scaled by progressively adding new key-value parameter pairs. Compared with the train-from-scratch approach (Biderman et al., ; Kaplan et al., ), our method achieves nearly the same performance while greatly reducing training costs.

## 2 Related Work

Transformer (Vaswani et al., ) has emerged as a foundational architecture in deep learning due to its versatile attention mechanism, enabling it to process any tokenized data and adapt to numerous domains, including language modeling (Radford et al., ; Touvron et al., ), image processing (Dosovitskiy et al., ), multi-modal understanding (Liu et al., ; Wang et al., ; ; ), decision making (Chen et al., ), graph learning (Yun et al., ), among others. While the Transformer effectively handles interactions among input tokens with flexibility, this property does not extend to computations involving model parameters, which are conducted via prescribed linear projections. In this work, we seek to restructure token-parameter interactions by developing a fully attention-based network that unifies both token-token and token-parameter computations through attention mechanisms, thus further extending the network’s flexibility.

Large Scale Training has proven to be an effective approach for developing powerful foundation models. As demonstrated by models like the GPT series (Radford et al., ; Alec et al., ; Brown et al., ), simple architectures—when supported by larger training datasets and increased model sizes (measured in parameters)—often outperform more complex algorithms. Scaling up data is generally more cost-effective because it is independent of the model’s architecture and allows for the continuous integration of new data through fine-tuning existing models (Kaplan et al., ). In contrast, increasing the model size often incurs extremely high costs, as it alters architectural details and usually requires retraining the entire dataset from scratch at each scaling step (Biderman et al., ). This significantly raises the expenses for building progressively larger models in the industry.

Model Reusing. Previous methods for reusing models have typically involved initializing larger models with pre-trained smaller models by duplicating (Chen et al., ; ), stacking (Gong et al., ), or combining (Wang et al., ) model weights. While these approaches can be effective, they often disturb the pre-established distribution of the smaller model, increasing the risk of losing pre-trained knowledge and slowing convergence. In contrast, our model allows for parameter scaling in a natural and seamless manner and preserves the integrity of the existing model.

## 3 Methodology

In this section, we first revisits the conventional attention mechanism in Section 3.1.
Then, Section 3.2 introduces \ourmethod, a natively scalable architecture centered around a flexible token-parameter attention layer.
Finally, incremental model scaling of \ourmethodis detailed in Section 3.3.

### 3.1 Preliminaries

Transformer models (Vaswani et al., ) have established themselves as fundamental architectures in deep learning, demonstrating outstanding performance across a wide range of tasks. The cornerstone of their success is the self-attention mechanism, which allows the model to dynamically assess the importance of each token, efficiently modeling complex dependencies among them.

Given a set of $T$ input tokens $X\in\mathbb{R}^{T\times d}$ with channel dimension $d$, the self-attention block first derives input-dependent query $Q$, key $K$, and value $V$, with three distinct linear projections as

$$ $Q=X\cdot W^{Q},~{}~{}~{}~{}~{}~{}K=X\cdot W^{K},~{}~{}~{}~{}~{}~{}V=X\cdot W^{ V},$ (1) $$

where the $W^{Q},W^{K}\in\mathbb{R}^{d\times d_{k}}$ and $W^{V}\in\mathbb{R}^{d\times d_{v}}$ are learnable weight matrices. The attention scores are calculated by measuring the similarity between query and key vectors, followed by a softmax function to obtain normalized weights. These scores are subsequently used to compute the output of the scaled dot-product attention as,

$$ $\text{Attention}(Q,K,V)=\text{softmax}[\frac{Q\cdot K^{\top}}{\sqrt{d}}]\cdot V,$ (2) $$

where $\sqrt{d}$ is a scale factor for alleviating small gradients caused by softmax. Finally, the output is,

$$ $O=X_{\text{att}}\cdot W^{O},$ (3) $$

with $X_{\text{att}}$ being the attention output and $W^{O}\in\mathbb{R}^{d_{v}\times d}$ as the output projection matrix.

The above architectural design enables the model to flexibly manage interactions between tokens of varying lengths, thereby allowing modern general models to concurrently process any form and quantity of tokenized multi-modal data. This capability markedly enhances the development of current AI domain and is fundamental to the success of transformer-based systems.

Figure: Figure 2: \ourmethod is a fully attention-driven architecture featuring a new token-Parameter attention (Pattention) layer. The Pattention uses a set of learnable tokens to represent model parameters and lets the input tokens attend to them. As the model scales, \ourmethodadds new learnable tokens to expand the existing key-value parameter sets, while keeping the feature dimension constant and leaving the rest of the computation unaffected.
Refer to caption: x2.png

### 3.2 \ourmethod

Although transformers excel across various domains, their scalability is limited by high computational overheads resulting from prescribed token-parameter interactions (i.e., linear projections). As a result, scaling strategies that adjust architectural components (e.g., channel dimensions) typically require retraining the entire model from the beginning, leading to inefficient use of computational resources.

To overcome this challenge, we propose \ourmethod, an architecture entirely based on attention mechanisms. The central innovation of \ourmethodis token-Parameter attention (Pattention) layer, which incorporates a set of trainable tokens functioning as model parameters and then employs cross-attention to manage interactions between input tokens and these parameter tokens. In this way, the Pattention layer introduces an additional dimension—the number of parameter tokens—which operates independently of the input and output channel dimensions. This decoupling enables input data to dynamically interact with a variable number of parameters, providing the flexibility required for incremental model scaling by reusing pre-trained models. Consequently, training larger models is greatly accelerated while achieving performance on par with transformers trained from scratch.

Pattention Layer. Let the input tokens and output tokens be represented as $\mathcal{I}\in\mathbb{R}^{T\times d_{1}}$ and $\mathcal{O}\in\mathbb{R}^{T\times d_{2}}$, where $T$ is the sequence length, and $d_{1}$ and $d_{2}$ are the input and output dimensions, respectively. To implement our Pattention mechanism, we introduce two sets of $n$ learnable parameter tokens: $K_{P}\in\mathbb{R}^{n\times d_{1}}$ representing the keys, and $V_{P}\in\mathbb{R}^{n\times d_{2}}$ representing the values.
The output $\mathcal{O}$ from the scaled dot-product Pattention layer is computed as:

$$ $\text{Pattention}(X,K_{P},V_{P})=\Theta\left(X\cdot K_{P}^{\top}\right)\cdot V _{P},$ (4) $$

where $\Theta$ is a modified softmax operation for stable optimization of Pattention layer. The output Pattention scores, $S\in\mathbb{R}^{n\times n}$, are formulated as,

$$ $S_{ij}=f\left(\frac{A_{ij}\times\tau}{\sqrt{\sum_{k=1}^{n}|A_{ik}|^{2}}}\right ),~{}~{}\forall~{}i,j\in 1...n,$ (5) $$

where $A$ is the score derived from $(X\cdot K_{P}^{\top})$, $\tau$ is the scale factor, which is set to $\sqrt{n}$ by default, and $f$ is a non-linearity function, which in our formulation is set to the GeLU function (Hendrycks & Gimpel, ). This design improves gradient stability in our architecture and results in better performance compared to the standard softmax operation (see Appendix A and Table 5 for details).

Our Pattention layer employs a cross-attention mechanism to manage interactions between tokens and parameters, thereby fully preserving the adaptability characteristic of attention mechanisms. Similar to how self-attention in Transformer models handles sequences with variable lengths, our Pattention layer is designed to process a flexible number of parameters independently of the input and output channel dimensions used in feature projection. This allows network parameters to be expanded seamlessly along the parameter token axis, enabling the effective reuse of pre-trained weights and offering a naturally incremental manner for model scaling.

Overall Architecture.
Figure 2 illustrates the architecture of \ourmethod. Given the input tokens $X_{\text{in}}\in\mathbb{R}^{T\times d}$, we follow the design of the pre-norm transformer, the computation for the output of a \ourmethodlayer is
represented as follows:

$$ $X_{\text{inter}}=X_{\text{in}}+\text{MHA}(\text{LN}(X_{\text{in}})),$ (6) $$

$$ $X_{\text{out}}=X_{\text{inter}}+\text{FFN}(\text{LN}(X_{\text{inter}})),$ (7) $$

where LN denotes the layer normalization (Ba, ; Zhang & Sennrich, ), and MHA and FFN refer to our modified multi-head self-attention and feed-forward layer, respectively.

In the multi-head self-attention block, for simplicity, we consider a single-head variant and set both $d_{k}$ and $d_{v}$ equal to $d$. Then we replace all the linear projections with our Pattention layers. Let $\text{LN}(X_{\text{in}})$ be denoted as $X$, this block is formulated as follows:

$$ $Q=\text{Pattention}(X,K_{P}^{Q},V_{P}^{Q}),~{}~{}K=\text{Pattention}(X,K_{P}^{ K},V_{P}^{K}),~{}~{}V=\text{Pattention}(X,K_{P}^{V},V_{P}^{V}),$ (8) $$

$$ $X_{\text{att}}=\text{softmax}\left[\frac{Q\cdot K^{\top}}{\sqrt{d}}\right] \cdot V,$ (9) $$

$$ $O_{\text{att}}=\text{Pattention}\left(X_{\text{att}},K_{P}^{O},V_{P}^{O}\right),$ (10) $$

where Eq. 8 and 10 represent token-parameter attention while Eq. 9 represents token-token attention. The key-value parameter tokens for the QKV projections are $(K_{P}^{Q},V_{P}^{Q})\in\mathbb{R}^{n_{q}\times d}$, $(K_{P}^{K},V_{P}^{K})\in\mathbb{R}^{n_{k}\times d}$, $(K_{P}^{V},V_{P}^{V})\in\mathbb{R}^{n_{v}\times d}$, while $(K_{P}^{O},V_{P}^{O})\in\mathbb{R}^{n_{o}\times d}$ is used for the output projection layer.

For consistency and simplicity, the feed-forward block in \ourmethod utilizes a single Pattention Layer. Denote $\text{LN}(X_{\text{inter}})$ as $X_{\text{ffn}}$, and the FFN computation is given by:

$$ $O_{\text{ffn}}=\text{Pattention}\left(X_{\text{ffn}},K_{P}^{\text{ffn}},V_{P}^ {\text{ffn}}\right),$ (11) $$

where $(K_{P}^{\text{ffn}},V_{P}^{\text{ffn}})\in\mathbb{R}^{n_{\text{ffn}}\times d}$ are learnable key-value pairs for FFN block.

By designing the architecture in this manner, we represent all fundamental components-including both input data and model parameters—as tokens within the computational framework. This token-centric perspective allows the utilization of successful attention mechanisms to unify two primary computations within the transformer, token-token and token-parameter interactions, thereby establishing a fully attention-based neural network characterized by exceptional flexibility.

Architecture Configurations. Our model meticulously mirrors the hyper-parameter configuration of the standard Transformer architecture. Taking GPT-2 (Radford et al., ) as an exemplar, which features 12 Transformer layers and a hidden dimension of 768, our model replicates this configuration with identical layer counts and dimensionality. The number of key-value parameter pairs in both the query-key-value and output projections corresponds directly to the hidden dimension. In contrast, the FFN module utilizes four times the number of parameter pairs relative to the hidden size. This architectural alignment facilitates the initialization of our model’s parameters using a pre-trained Transformer, thereby ensuring seamless integration into the Transformer pre-training ecosystem.

### 3.3 Progressive Model Scaling

Our model demonstrates strong suitability for large-scale model training along the parameter axis, attributable to the versatile design of the Pattention layer, which allows for the incremental development of larger models by reusing parameters from smaller, pre-trained counterparts.

To facilitate understanding without compromising generality, we employ a single Pattention layer to exemplify the intricacies of model scaling. Consider an existing \ourmethodmodel equipped with a set of pre-trained key-value parameter tokens, denoted as $K_{P}^{\text{old}},V_{P}^{\text{old}}\in\mathbb{R}^{n\times d}$. As shown in Figure 2, to scale the model, we augment this set by appending new key-value parameter tokens $K_{P}^{\text{new}},V_{P}^{\text{new}}\in\mathbb{R}^{m\times d}$ as

$$ $K_{P}^{\text{scale}}=\left[K_{P}^{\text{old}},K_{P}^{\text{new}}\right],~{}~{} ~{}~{}~{}V_{P}^{\text{scale}}=\left[V_{P}^{\text{old}},V_{P}^{\text{new}} \right],$ (12) $$

where $[\cdot^{\text{old}},\cdot^{\text{new}}]$ means the concatenation operation along the token dimension and $K_{P}^{\text{scale}}$, $V_{P}^{\text{scale}}\in\mathbb{R}^{(m+n)\times d}$ are scaled parameter sets. The forward pass of the scaled model is then defined as

$$ $O=\text{Pattention}\left(X,K_{P}^{\text{scale}},V_{P}^{\text{scale}}\right).$ (13) $$

This scaling scheme permits the integration of an arbitrary number of parameters without altering the input or output dimensions. As demonstrated in Figure 4, this approach notably enhances training efficiency for models at greater scales without degrading performance. Importantly, by initializing $K^{\text{new}}_{P}$ with zero, similar to LoRA technique (Hu et al., ), our model can perfectly resume the model state from the pre-training phase without losing the well-learned knowledge, facilitating faster convergence and accelerating the overall scaling process.

Figure: Figure 3: Evaluating model scaling costs through cumulative computational budgets. The Transformer baseline incurs expenses for each individual scaling step performed independently from scratch, whereas \ourmethodaggregates costs across all scaling stages, including training a 124M model initially, progressively scaling to 354M, 757M, and 1.4B parameters.
Refer to caption: x3.png

## 4 Experiments

In this section, we present experimental results for the techniques described above. Section 4.1 validates the continual expansion capability of our model. Section 4.2 highlights the model’s efficacy in handling tasks within both language and vision domains. Section 4.3 offers an in-depth comparison, highlighting our model’s advantages over standard Transformer models. Finally, Section 4.4 details the ablation experiments conducted to assess the significance of each module in \ourmethod.

### 4.1 Progressive Model Scaling

Datasets.
Our models are trained using the OpenWebText Corpus described in (Gokaslan & Cohen, ). This corpus serves as a widely recognized open-source approximation of OpenAI’s proprietary WebText dataset, which was employed in the development of GPT-2 (Alec et al., ). The dataset comprises textual data extracted from 8,013,769 Reddit-shared documents. During training, we randomly sample segments from these documents.

Baseline Transformer Training from Scratch. To evaluate the effectiveness of our progressive model scaling strategy, we established a baseline by training a Transformer model from scratch. Following the training procedures outlined in Karpathy (); Kaplan et al. (), we employed the AdamW optimizer (Loshchilov & Hutter, ) with a batch size of 512 sequences, each containing 1024 tokens. For a fair comparison with our incremental scaling approach, we configured two training variants based on the total number of training tokens. The first variant underwent $6\times 10^{5}$ steps (approximately 300B tokens), consistent with the training steps utilized by Karpathy () to replicate GPT-2 performance. The second variant was limited to $6\times 10^{4}$ steps (approximately 30B tokens) to ensure comparability with each stage of our progressive scaling. In all trainings included in our analysis, unless otherwise indicated, a learning rate of $6\times 10^{-4}$ was employed, featuring a 2000-step warmup followed by a cosine decay to zero.

with Progressive Model Scaling. Building upon the above training protocols, we testify the performance of our model scaling with parameter sizes ranging from 124M to 1.4B. Unlike the aforementioned scratch-training approach, each scaling iteration leverages a pre-trained smaller \ourmethodto partially initialize the weights of the larger one described in Section 3.3. The scaling procedure begins with training the initial source model from scratch on approximately 300B tokens, mirroring the Transformer baseline. For scaling, we select the pre-trained model closest in parameter count to the target size for weight initialization. For example, to train a model with 354M parameters, we employ the 124M model as a partial initializer and retrain the entire model using a reduced computational budget (e.g., 15B, 30B, or 60B tokens). This iterative process continues for scaling to 757M and then to 1.4B parameters. Notably, to simplify the scaling procedure, both new and existing parameters are trained equivalently with identical training hyperparameters throughout the process.

Our training optimizes the autoregressive log-likelihood (i.e., cross-entropy loss) averaged over a 1024-token context and the log perplexity evaluated on the test set as the test score.

Experimental Analysis. As illustrated in Figure 4, our progressive scaling methodology employing \ourmethodachieves performance comparable to that of a Transformer model trained from scratch, while substantially reducing the training budget. Specifically, Starting with a 124M parameter model trained on 300B tokens, we progressively scaled to 354M, 757M, and 1.4B parameters, requiring only an additional 30B tokens—just one-tenth of the computational budget compared to the scratch-trained Transformer. This scaling process achieved a test perplexity of 11.77 at the 1.4B parameter level. In comparison, a Transformer model of the same size trained from scratch achieved a similar perplexity of 11.63 but with $3\times$ the training cost. Importantly, our approach reports cumulative training costs, encompassing all scaling stages, unlike the Transformer baseline that only accounts for individual stages. Even with this comparison, our method demonstrates a substantially lower computational cost than training a Transformer from scratch, thereby validating the effectiveness of our approach.

Figure 4 presents the training costs at each scaling stage for both our model and the standard Transformer. When compared to Figure 4, the cost savings are even more significant. Specifically, our model requires only one-tenth of the training costs associated with Transformer baselines. To mitigate the effects of varying training data, we also included the performance curve of a Transformer trained from scratch using an equivalent computational budget of 30B tokens. Under the same computational constraints, our progressively scaled model achieves a lower perplexity of 11.77 compared to the Transformer’s 13.34, thereby highlighting the superior efficiency and scalability of our approach.

**Table 1: (Zero-shot Evaluations.) The best performance for each model size is highlighted in bold. Our comparisons are made with publicly available transformer-based LMs with various tokenizers. Following Pythia (Biderman et al., ), our model is trained for up to 300B tokens on pile dataset.**
|  |  | Pile | LAMBADA | LAMBADA | HellaSwag | PIQA | Arc-E | Arc-C | WinoGrande | Average |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | #Param | ppl $\downarrow$ | ppl $\downarrow$ | acc $\uparrow$ | acc $\uparrow$ | acc $\uparrow$ | acc $\uparrow$ | acc $\uparrow$ | acc $\uparrow$ | acc $\uparrow$ |
| Pythia-160M (Biderman et al., ) | 160M | 29.64 | 37.25 | 35.4 | 30.3 | 62.3 | 43.6 | 23.6 | 51.3 | 40.1 |
| Ours (TokenFormer-150M) | 150M | 10.45 | 16.38 | 45.0 | 35.5 | 64.9 | 47.3 | 24.9 | 50.4 | 44.7 |
| Pythia-410M (Biderman et al., ) | 410M | 9.95 | 10.84 | 51.4 | 40.6 | 66.9 | 52.1 | 24.6 | 53.8 | 48.2 |
| Ours (TokenFormer-450M) | 450M | 8.28 | 7.69 | 57.3 | 47.5 | 69.5 | 56.2 | 26.7 | 54.6 | 52.0 |
| Pythia-1B (Biderman et al., ) | 1B | 7.82 | 7.92 | 56.1 | 47.2 | 70.7 | 57.0 | 27.1 | 53.5 | 51.9 |
| Ours (TokenFormer-900M) | 900M | 7.38 | 5.46 | 64.0 | 55.3 | 72.4 | 59.9 | 30.6 | 56.4 | 56.4 |
| GPT-Neo 1.3B (Black et al., ) | 1.3B | - | 7.50 | 57.2 | 48.9 | 71.1 | 56.2 | 25.9 | 54.9 | 52.4 |
| OPT-1.3B (Zhang et al., ) | 1.3B | - | 6.64 | 58.0 | 53.7 | 72.4 | 56.7 | 29.6 | 59.5 | 55.0 |
| Pythia-1.3B (Biderman et al., ) | 1.3B | 7.51 | 6.08 | 61.7 | 52.1 | 71.0 | 60.5 | 28.5 | 57.2 | 55.2 |
| GPT-Neo 2.7B (Black et al., ) | 2.7B | - | 5.63 | 62.2 | 55.8 | 71.1 | 61.1 | 30.2 | 57.6 | 56.5 |
| OPT-2.7B (Zhang et al., ) | 2.7B | - | 5.12 | 63.6 | 60.6 | 74.8 | 60.8 | 31.3 | 61.0 | 58.7 |
| Pythia-2.8B (Biderman et al., ) | 2.8B | - | 5.04 | 64.7 | 59.3 | 74.0 | 64.1 | 32.9 | 59.7 | 59.1 |
| Ours (TokenFormer-1.5B) | 1.5B | 6.91 | 5.24 | 64.7 | 60.0 | 74.8 | 64.8 | 32.0 | 59.7 | 59.3 |

**Table 2: (Image Classification.) Comparison of standard vision transformer on ImageNet-1K. The training hyperparameters are completely consistent (batch size, learning rate, etc.) with He et al. (). ${\dagger}$ denotes models where the parameter size has been matched to that of the standard ViT.**
| Method | Image Size | #Param | Top-1 acc |
| --- | --- | --- | --- |
| ViT-B/16 (Dosovitskiy et al., ) | 384^2 | 86M | 77.9 |
| DeiT-B/16 (Touvron et al., ) | 224^2 | 86M | 81.8 |
| ViT-B/16 (MAE) (He et al., ) | 224^2 | 86M | 82.3 |
| Ours (TokenFormer-B/16^† ) | 224^2 | 86M | 82.1 |
| Ours (TokenFormer-B/16) | 224^2 | 109M | 82.5 |
| ViT-L/16 (Dosovitskiy et al., ) | 384^2 | 307M | 76.5 |
| ViT-L/16 (MAE) (He et al., ) | 224^2 | 307M | 82.6 |
| Ours (TokenFormer-L/16^†) | 224^2 | 307M | 83.0 |
| Ours (TokenFormer-L/16) | 224^2 | 407M | 83.1 |

**Table 3: Parameter counts and training compute estimates for Transformer and our \ourmethod. Sub-leading terms such as nonlinearities, biases, and layer normalization are omitted.**
|  | Parameter | Training FLOPs |  |  |
| --- | --- | --- | --- | --- |
| Operation | Transformer | Ours | Transformer | Ours |
| Embed | $n_{\text{vocab}}d_{\text{model}}$ | $n_{\text{vocab}}d_{\text{model}}$ | - | - |
| Attention: QKV Project | $3n_{\text{layer}}d_{\text{model}}^{2}$ | $n_{\text{layer}}d_{\text{token}}(n_{\text{q}}+n_{\text{k}}+n_{\text{v}}$) | $6n_{\text{layer}}d_{\text{model}}^{2}T$ | $2n_{\text{layer}}d_{\text{token}}(n_{\text{q}}+n_{\text{k}}+n_{\text{v}})T$ |
| Attention: Token-Token | - | - | $4n_{\text{layer}}d_{\text{model}}T^{2}$ | $4n_{\text{layer}}d_{\text{token}}T^{2}$ |
| Attention: Output Project | $n_{\text{layer}}d_{\text{model}}^{2}$ | $n_{\text{layer}}d_{\text{token}}n_{\text{o}}$ | $2n_{\text{layer}}d_{\text{model}}^{2}T$ | $2n_{\text{layer}}d_{\text{token}}n_{\text{o}}T$ |
| Feedforward | $8n_{\text{layer}}d_{\text{model}}^{2}$ | $2n_{\text{layer}}d_{\text{token}}n_{\text{ff}}$ | $16n_{\text{layer}}d_{\text{model}}^{2}T$ | $4n_{\text{layer}}d_{\text{token}}n_{\text{ff}}T$ |
| De-embed | - | - | 2$n_{\text{vocab}}d_{\text{model}}$ | 2$n_{\text{vocab}}d_{\text{model}}$ |
| Total (Non-Embedding) | $N=12n_{\text{layer}}d_{\text{model}}^{2}$ | $N=n_{\text{layer}}d_{\text{token}}(n_{\text{q}}+n_{\text{k}}+n_{\text{v}}+n_{<br>\text{o}}+2n_{\text{ff}})$ | $2NT+4n_{\text{layer}}d_{\text{model}}T^{2}$ | $2NT+4n_{\text{layer}}d_{\text{token}}T^{2}$ |

Figure: Figure 5: The relationship between FLOPs and text length for both Transformer and \ourmethod. As shown in Table 3, Transformer exhibits an increase in computational cost for token-token interactions as $d_{\text{model}}$ scales upwards. Our \ourmethodmodel, however, offers a flexible parameter scaling mechanism that maintains $d_{\text{token}}$ at a constant value. This strategy results in controllable computational costs for token-token interactions and markedly enhances the efficiency of long-text modeling.
Refer to caption: x5.png

### 4.2 Benchmarking of Model Expressiveness

Language Modeling. We assess the efficacy of our proposed architecture through standard autoregressive language modeling tasks, benchmarking against existing transformer-based models. Evaluations are conducted using both pre-training metrics, specifically perplexity, and zero-shot performance measures. Training is performed on the Pile dataset (Gao et al., ), following the training protocol described in Biderman et al. (). Detailed training procedures and model sizes (depth and width) are provided in the Appendix F.

Table 1 presents the performance of \ourmethodacross various widely-recognized zero-shot downstream tasks. Comparisons are drawn against leading open-source transformer models of equivalent scale, notably Pythia (Biderman et al., ), which utilizes the same tokenizer, dataset, and training duration (300B tokens) as our models. As shown in this table, our model achieves competitive performance compared to the standard Transformer, demonstrating the potential of our architecture in terms of expressive power as a foundation model.

Visual Modeling. Table 2 validates the expressiveness of our model in visual tasks. We compare our approach against the standard Vision Transformer (ViT) (Dosovitskiy et al., ) trained with supervised learning on the ImageNet-1K dataset (Deng et al., ). For a fair comparison, we used the MMDetection code base (MMDetection Contributors, ) and followed the hyperparameters and training strategy used in He et al. (). As shown in the table, our model achieves the same performance as ViT in visual modeling, confirming its expressiveness in visual tasks.

### 4.3 Comparison with standard transformer

Transformer can also achieve model reuse to a certain extent. Net2Net (Chen et al., ), a classical model growth method, proposes a technique to expand the width of neural networks by duplicating neurons. In this method, the pre-trained weight matrix of a transformer layer in the smaller model denoted $W_{s}^{\text{old}}\in\mathbb{R}^{d_{s}\times d_{s}}$, is used to create a larger weight matrix $W_{l}^{\text{new}}\in\mathbb{R}^{d_{l}\times d_{l}}$ ($d_{l}>d_{s}$) to fill the larger model. This expansion is formulated as follows,

$$ $W_{l}^{\text{new}}=\left[\begin{array}[]{cc}W_{s}^{\text{old}}&W_{l(12)}^{ \text{new}}\\ W_{l(21)}^{\text{new}}&W_{l(22)}^{\text{new}}\\ \end{array}\right],$ (14) $$

where $W_{l(12)}^{\text{new}}\in\mathbb{R}^{(d_{l}-d_{s})\times d_{s}}$, $W_{l(21)}^{\text{new}}\in\mathbb{R}^{d_{s}\times(d_{l}-d_{s})}$, and $W_{l(22)}^{\text{new}}\in\mathbb{R}^{(d_{l}-d_{s})\times(d_{l}-d_{s})}$ are new parameters for expansion. The scaling procedures are the same as schemes introduced in Section 4.1.

Controllable cost of token-token interaction for long-context modeling. Recent advancements in Chain-of-Thought (CoT) modeling (Wei et al., ) have emphasized the critical importance of efficiently processing lengthy textual sequences (Tay et al., ) within Large Language Models (LLMs). As delineated in Section 1, the training costs of transformer architectures are primarily divided into two components: interactions involving model parameters and interactions among input sequences. Table 3 demonstrates that the computational complexity of transformer-based models exhibits a quadratic dependence on text length, scaling linearly with token-parameter interactions and quadratically with token-token interactions. Consequently, it is imperative to expand model parameters while controlling the computational burden of token-token interaction part.

Conventionally, scaling transformer models involves increasing the channel dimension. For a fixed text length, this results in higher computational costs, mainly because dominant token-token interactions become more intensive, which hampers the model’s performance with long texts. Our proposed model takes a different approach by decoupling the computation cost of token-token interactions from model scaling. We increase the parameter size without changing the token channel dimension, thereby maintaining the computational cost associated with token-token interactions. As shown in Figure 5, our model exhibits increasingly significant computational advantages over Transformers as the number of parameters grows, especially when processing longer sequences.

Figure: Figure 6: Loss curves comparing pre-trained Transformer and \ourmethodas their parameters are scaled during continued training on enwik8.
Refer to caption: x6.png

Scaling without losing the well-learned distribution. Our \ourmethodcan maintain the existing output distribution when new key parameters are initialized to zero. This characteristic is beneficial for continuously scaling models to incorporate additional data, as it facilitates an increase in model capacity without disrupting the ongoing training process, thereby promoting rapid convergence.

To evaluate \ourmethod’s scaling efficacy, we compare the loss curves of Net2Net-based transformer scaling against \ourmethodscaling. Both models, initially with 354M parameters, were pre-trained on the OpenWebText dataset. We then introduced the EnWik8 dataset and continued training with one epoch, expanding the models to 757M parameters to accommodate new data. Figure 7 demonstrates \ourmethodnot only converges more rapidly but also reaches a lower final loss, attributable to its ability to preserve the output distribution during the resumption of training.

Performance benchmarking on incremental scaling. In this study, we progressively scale the standard Transformer using Net2Net approach detailed earlier. For a fair comparison, we aligned all hyperparameters, including the parameter size, learning rate, dataset, and so on. As shown in Figure 7, our model performs better in scaling compared to the standard Transformer.

### 4.4 Ablation Study

**Table 4: Ablation of Softmax part on ImageNet classification with base model.**
| Nonlinear Function | Normalization | Top-1 acc |
| --- | --- | --- |
| $e^{x}$ | $L_{1}$ Norm | 79.6 |
| GeLU | $L_{1}$ Norm | 81.7 |
| GeLU | $L_{2}$ Norm | 82.5 |

Optimized Softmax Function in Pattention Layer. Within the token-parameter attention layer, we address training instabilities arising from the diminished gradients associated with the traditional softmax function. The conventional softmax operation comprises two primary steps: computing the exponential of attention scores followed by $L_{1}$ normalization. As shown in Table 5, to mitigate the issue of small gradients, we substitute the exponential non-linearity with the GeLU function (Hendrycks & Gimpel, ), resulting in a performance enhancement of +2.1 points on the ImageNet classification benchmark. Subsequently, we replace the $L_{1}$ normalization with an $L_{2}$ normalization, yielding an additional improvement of +0.8 points. These modifications collectively allow our model to achieve performance parity with the standard Vision Transformer.

Non-Parametric Layer Normalization. In pursuit of enabling model expansion and the merging of two separately trained parameter token sets for subsequent studies, we modified the Transformer’s layer normalization to a non-parametric variant by removing its trainable weights and biases. This adjustment guarantees that only the key-value parameters are subject to learning within the model. Empirical results presented in Table 5 demonstrate that the model maintains comparable performance after discarding the learnable weights and biases.

## 5 Future Work

Extending the Mixture-of-Experts Paradigm. We interpret \ourmethodas an extreme instantiation of the Mixture of Experts (MoE) framework, where each key-value parameter pair functions as an individual expert. This innovative MoE-like architecture has the potential to significantly reduce the computational costs associated with token-parameter interactions. Additionally, \ourmethod’s adjustable computational load for token-token interactions complements the MoE feature, facilitating the development of more resource-effective foundational models.

Advancing Parameter-Efficient Tuning. The scaling approach of \ourmethod, which involves integrating additional key-value parameter pairs, exemplifies a strategy for parameter-efficient tuning. When confronted with new tasks or datasets, the model can augment its pre-trained parameters by incorporating these new parameter tokens, thereby adapting to specific task requirements quickly.

Integrating Vision and Language Models. Leveraging the parameter-efficient tuning capabilities of Tokeformer, we can achieve seamless integration of visual and linguistic modalities. This can be accomplished by unifying the key-value parameter tokens derived from pre-trained visual \ourmethodand language \ourmethodinto a single parameter set. Then, the new learnable tokens are introduced to perform vision-language alignment and instruction tuning.

Device-Cloud Collaboration. \ourmethodcan serve as the cloud-side knowledge base in device-cloud collaboration of on-device LLMs, with each pair of key-value parameter tokens representing a learnable pattern, leveraging the device for real-time processing and the cloud for intensive tasks.

Enhancing Model Interpretability. As \ourmethodis entirely based on attention mechanisms, it inherently benefits from the interpretability associated with attention in token-parameter interactions. This characteristic enhances the model’s explainability, contributing to the AI community’s efforts to develop more transparent and understandable models.

## 6 Conclusion

This paper introduces \ourmethod, a naturally scalable architecture that leverages the attention mechanism to facilitate not only inter-token computations but also interactions between tokens and model parameters, thereby enhancing architectural flexibility. By representing model parameters as tokens, we replace all linear projection layers in the Transformer with our Pattention layers, allowing for seamless and efficient incremental scaling without the need for retraining from scratch. We believe that this architecture, offering greater flexibility than traditional Transformers, will further contribute to the development of foundation models.

## Appendix A Gradient of Pattention Layer

Our token-parameter attention mechanism employs $L_{2}$-normalization followed by the GeLU (Hendrycks & Gimpel, ) activation function, in contrast to the conventional SoftMax function in standard token-token attention layers, which utilizes an exponential transformation followed by $L_{1}$-normalization.
This design choice is motivated by our experimental observation that SoftMax tends to increase the magnitude of outputs, often pushing them into regions where the gradients become extremely small, leading to an inferior overall performance (see Table 5).

Specifically, given a query token and $n$ key-value pairs with dimension $d$, let the similarity scores between the query and key tokens be represented as $A\in\mathbb{R}^{1\times n}$.
In standard SoftMax attention, the attention scores $S\in\mathbb{R}^{1\times n}$ are computed as follows:

$$ $S_{i}=\frac{\text{exp}(A_{i}/\sqrt{d})}{\sum_{j=1}^{n}\text{exp}(A_{j}/\sqrt{d })},~{}~{}\forall~{}i\in 1...n,$ (15) $$

The derivative of the SoftMax function with respect to $A_{i}$ is given by:

$$ $\frac{\partial S_{i}}{\partial A_{j}}=\frac{1}{\sqrt{d}}S_{i}(\mathbbm{1}_{i=j }-S_{j})=\left\{\begin{aligned} &\frac{1}{\sqrt{d}}S_{i}(1-S_{j})~{}~{}~{}~{}~ {}~{}i=j\\ &-\frac{1}{\sqrt{d}}S_{i}S_{j}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}i\neq j.\end{ aligned}\right.$ (16) $$

In contrast, our activation function uses L2-normalization followed by the GeLU function.
Denoting GeLU as $f$, the attention scores $Z\in\mathbb{R}^{1\times n}$ are computed as:

$$ $\hat{S_{i}}=f(Z_{i})=f(\frac{A_{i}\times\sqrt{n}}{\sqrt{\sum_{j=1}^{n}|A_{j}|^ {2}}}),~{}~{}\forall~{}i\in 1...m,$ (17) $$

For the derivative of our attention function when $i=j$, we have:

$$ $\displaystyle\frac{\partial\hat{S_{i}}}{\partial A_{i}}$ $\displaystyle=f^{\prime}\sqrt{m}\frac{||A||_{2}-A_{i}\frac{A_{i}}{||A||_{2}}}{ ||A||_{2}^{2}}$ (18) $\displaystyle=f^{\prime}\sqrt{n}\frac{||A||_{2}^{2}-A_{i}^{2}\frac{1}{||A||_{2 }}}{||A||_{2}^{2}}$ (19) $\displaystyle=f^{\prime}\sqrt{n}\frac{1}{||A||_{2}}\frac{||A||_{2}^{2}-A_{i}^{ 2}}{||A||_{2}^{2}}$ (20) $\displaystyle=f^{\prime}\sqrt{n}\frac{1}{||A||_{2}}(1-Z_{i}^{2}d)$ (21) $\displaystyle=f^{\prime}\frac{1}{\sqrt{n}}\frac{1}{||A||_{2}}(d-Z_{i}^{2})$ (22) $$

When $i\neq j$, the derivative becomes:

$$ $\displaystyle\frac{\partial\hat{S_{i}}}{\partial A_{i}}$ $\displaystyle=-f^{\prime}\sqrt{n}\frac{A_{i}\frac{A_{j}}{||A||_{2}}}{||A||_{2} ^{2}}$ (23) $\displaystyle=-f^{\prime}\sqrt{n}\frac{1}{||A||_{2}}\frac{A_{i}A_{j}}{||A||_{2 }^{2}}$ (24) $\displaystyle=-f^{\prime}\frac{1}{\sqrt{n}}\frac{1}{||A||_{2}}Z_{i}Z_{j}$ (25) $$

Thus, the derivative of our attention function is:

$$ $\frac{\partial\hat{S_{i}}}{\partial a_{j}}=\left\{\begin{aligned} &f^{\prime} \frac{1}{\sqrt{n}}\frac{1}{||A||_{2}}(n-Z_{i}Z_{j})~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}i=j\\ &-f^{\prime}\frac{1}{\sqrt{n}}\frac{1}{||A||_{2}}Z_{i}Z_{j}~{}~{}~{}~{}~{}~{}~ {}~{}~{}~{}~{}~{}~{}~{}~{}~{}i\neq j.\end{aligned}\right.$ (26) $$

Comparing the gradients of SoftMax (Eq. 16) and our method (Eq. 26), the key difference is that our gradient depends on the product $Z_{i}Z_{j}$, whereas SoftMax relies on $S_{i}S_{j}$. Due to the exponential nature of SoftMax, the distribution of $S_{i}$ tends to be sharper and more concentrated, which often drives the gradients toward zero. Conversely, our activation function produces a smoother distribution of $Z$, mitigating the vanishing gradient problem and enabling more stable training dynamics.

## Appendix B Zero initialization in \ourmethod

As shown in Section 3.3, during model scaling, initializing new key parameters to zero allows the model to continue training with minimal disruption.
This is because zero initialization preserves the model’s original output distribution, preventing significant interference to the learned representations.

In this section, we demonstrate that the Pattention layer is invariant to newly added parameters when they are zero-initialized.
Let $X\in\mathbb{R}^{d}$ be the input vector, and let the Pattention layer have $n$ key-value pairs, represented as $K_{P},V_{P}\in\mathbb{R}^{n\times d}$. The output of the Pattention layer is computed as:

$$ $A=K_{P}\cdot X,$ (27) $$

$$ $S=f(\frac{A}{\sqrt{\sum_{j}A_{j}^{2}}}),$ (28) $$

$$ $O=V_{P}^{\top}\cdot S.$ (29) $$

When scaling the model by adding $m$ new key-value pairs with zero initialization, the output becomes:

$$ $\hat{A}=\begin{bmatrix}K_{P}\\ 0,..,0\\ \vdots\\ 0,..,0\end{bmatrix}\cdot X=\begin{bmatrix}A\\ 0\\ \vdots\\ 0\end{bmatrix},$ (30) $$

$$ $\hat{S}=f(\frac{\hat{A}}{\sqrt{\sum_{j}\hat{A}_{j}^{2}}})=\begin{bmatrix}S\\ 0\\ \vdots\\ 0\end{bmatrix},$ (31) $$

$$ $\hat{O}=\begin{bmatrix}V_{P}^{\top},V_{P}^{\text{new}\top}\\ \end{bmatrix}\cdot\hat{S}=O.$ (32) $$

Since the newly added key parameters are initialized to zero, the attention mechanism does not modify the original output. Therefore, the output $\hat{O}$ remains identical to $O$.
This property is advantageous for scaling models because it increases the model capacity without interrupting the well-learned distribution and the ongoing training process, leading to faster convergence.

## Appendix C Tabular Main Results

Here, we provide the tabulated results corresponding to Figure 4 from the main paper.
Table 7 presents the perplexity on the validation set of OpenWebText. The Transformer models are trained from scratch, while \ourmethodmodels leverage parameter reuse from smaller models (except for the first \ourmethodmodel with 124M parameters, which is also trained from scratch).
We observe that \ourmethodachieves on-par performance to Transformers trained from scratch, but with significantly reduced training costs due to parameter reuse. Notably, the largest \ourmethod(1.4B) achieves even a slightly better perplexity than its Transformer counterpart (11.60 vs. 11.63).

**Table 6: Tabular results of Figure 4. The perplexity of models trained with different numbers of schemes is compared. Transformers are trained from scratch, while \ourmethod are progressively scaled up via parameter resuing. When trained with the same number of tokens (30B), \ourmethod demonstrates superior performance.**
| Perplexity | #Param |  |  |
| --- | --- | --- | --- |
| 354M | 757M | 1.4B |  |
| Transformer<br>Train from scratch 300B | 13.02 | 11.99 | 11.63 |
| Transformer<br>Train from scratch 30B | 15.97 | 13.78 | 13.34 |
| \ourmethod<br>Parameter Reusing 30B | 14.02 | 12.59 | 11.77 |

Table 7 compares Transformer models trained from scratch and \ourmethodtrained by parameter reusing with varying amounts of seen tokens during training. It is evident that Transformers trained with the same number of seen tokens do not reach the performance level of \ourmethodwith parameter reusing. This demonstrates that parameter reusing successfully transfers knowledge from smaller models, reducing training time without sacrificing performance.

## Appendix D Experimental Details on Progressive model scaling

In this experiment, we utilize the OpenWebText dataset (Gokaslan & Cohen, ) to evaluate the model scaling capability.
The dataset comprises 8,013,769 documents, from which we randomly select 5% to serve as the validation set and report perplexity on this subset.
We investigate four model sizes: 124M, 354M, 757M, and 1.4B parameters.
Please find model specifications in Table 8.
During parameter reusing of \ourmethod, we partially resume the old model parameters and add new key-value pairs to the Pattention layers and do not alter the number of layers or feature dimensions.

The training recipe is the same for both Transformers and \ourmethod:
We do not implement dropout in either model and the logits are computed at the final layer using the embedding layer.
The tokenizer is from GPT-NeoX-20B (Black et al., ).
We employ the AdamW optimizer (Loshchilov, ) with $\beta_{1}=0.9$ and $\beta_{2}=0.95$.
A learning rate of $6\times 10^{-4}$ is employed, with a linear warmup over 2000 steps followed by a cosine decay to zero.
The training is conducted with a batch size of 512 and a sequence length of 1024.

## Appendix E Experimental Details on Scaling without losing the well-learned distribution

In this experiment, we utilize the EnWik8 (Mahoney, ) dataset to evaluate the model’s capacity for continued adaptation to new data. The EnWik8 dataset comprises the first 100 million bytes of English Wikipedia in XML format.

We begin with a model size of 354M parameters, which has been pre-trained on OpenWebText (Gokaslan & Cohen, ), and then scale it to 757M parameters for both the Transformer and \ourmethodmodels. In the case of the Transformer, the 354M and 757M models differ solely in feature dimension and the number of heads in the multi-head attention mechanism, and we follow Net2Net(Chen et al., ) methodology for parameter expansion. For \ourmethod, we increase the number of key-value pairs from 2140 to 4850.

We employ a constant learning rate of $6\times 10^{-4}$ and process the dataset for a single pass, resulting in a total of 2204 training steps. The AdamW optimizer (Loshchilov, ) is utilized, and we do not resume the optimizer’s internal state from any previous runs. The batch size is set to 512, consistent with the batch size used during pre-training.

**Table 8: Details of Tokenformer model variants used in Section 4. ${\dagger}$ indicates models whose key-value pairs are chosen to match the parameter numbers of Transformer of equivalent sizes.**
| Model | Layers | Hidden size | Attention<br>KV Pairs | FFN<br>KV Pairs | Heads | #Params |
| --- | --- | --- | --- | --- | --- | --- |
| Language Modeling |  |  |  |  |  |  |
| \ourmethod-150M | 12 | 768 | 768 | 3072 | 12 | 150M |
| \ourmethod-450M | 24 | 1024 | 1024 | 4096 | 16 | 450M |
| \ourmethod-900M | 32 | 1280 | 1280 | 5120 | 16 | 900M |
| \ourmethod-1.5B | 40 | 1536 | 1536 | 6144 | 16 | 1.5B |
| Visual Modeling |  |  |  |  |  |  |
| \ourmethod-Base^† | 12 | 768 | 576 | 2304 | 12 | 86M |
| \ourmethod-Base | 12 | 768 | 768 | 3072 | 12 | 109M |
| \ourmethod-Large^† | 24 | 1024 | 768 | 768 | 16 | 307M |
| \ourmethod-Large | 24 | 1024 | 1024 | 4096 | 16 | 407M |
| Parameter Reusing |  |  |  |  |  |  |
| \ourmethod-124M | 12 | 768 | 576 | 2304 | 12 | 124M |
| \ourmethod-354M | 12 | 768 | 2140 | 8560 | 12 | 354M |
| \ourmethod-757M | 12 | 768 | 4850 | 19400 | 12 | 757M |
| \ourmethod-1.4B | 12 | 768 | 8620 | 34480 | 12 | 1.4B |

## Appendix F Experiments on Language modeling Benchmarking

We evaluate the expressiveness and performance of our proposed architecture using standard autoregressive language modeling benchmarks, comparing its results to those of existing open-source LLMs, including RNN-based methods (Peng et al., ; Gu & Dao, ), in Table 9.
Evaluations are conducted using both pre-training metrics, specifically perplexity, and zero-shot performance measures. For training, we used the Pile dataset (Gao et al., ) over a single epoch, adhering to the training strategy used by Biderman et al. (). The Adam optimizer is applied with a $6\times 10^{-4}$ learning rate, and each batch contains 1024 samples with a sequence length of 2048 tokens, mirroring the setup of Mamba. The training process consists of 14,300 steps, with a 1430-step warmup, equivalent to 1% of total training steps. Detailed model specifications are in Table 8.

**Table 9: (Zero-shot Evaluations.) The best results for each model size are highlighted in bold. We compare our models against open-source language models (LMs), including RNN-based methods, with various tokenizers, trained for up to 300B tokens.**
|  |  | Pile | LAMBADA | LAMBADA | HellaSwag | PIQA | Arc-E | Arc-C | WinoGrande | Average |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | #Param | ppl $\downarrow$ | ppl $\downarrow$ | acc $\uparrow$ | acc $\uparrow$ | acc $\uparrow$ | acc $\uparrow$ | acc $\uparrow$ | acc $\uparrow$ | acc $\uparrow$ |
| Hybrid H3-130M (Fu et al., ) | 130M | - | 89.48 | 25.8 | 31.7 | 64.2 | 44.4 | 24.2 | 50.6 | 40.1 |
| Pythia-160M (Biderman et al., ) | 160M | 29.64 | 37.25 | 35.4 | 30.3 | 62.3 | 43.6 | 23.6 | 51.9 | 40.6 |
| Mamba-130M (Gu & Dao, ) | 130M | 10.56 | 16.07 | 44.3 | 35.3 | 64.5 | 48.0 | 24.3 | 51.9 | 44.7 |
| Ours (TokenFormer-150M) | 150M | 10.45 | 16.38 | 45.0 | 35.5 | 64.9 | 47.3 | 24.9 | 50.4 | 44.7 |
| Hybrid H3-360M (Fu et al., ) | 360M | - | 12.58 | 48.0 | 41.5 | 68.1 | 51.4 | 24.7 | 54.1 | 48.0 |
| Pythia-410M (Biderman et al., ) | 410M | 9.95 | 10.84 | 51.4 | 40.6 | 66.9 | 52.1 | 24.6 | 53.8 | 48.2 |
| Mamba-370M (Gu & Dao, ) | 370M | 8.28 | 8.14 | 55.6 | 46.5 | 69.5 | 55.1 | 28.0 | 55.3 | 50.0 |
| Ours (TokenFormer-450M) | 450M | 8.28 | 7.69 | 57.3 | 47.5 | 69.5 | 56.2 | 26.7 | 54.6 | 52.0 |
| Pythia-1B (Biderman et al., ) | 1B | 7.82 | 7.92 | 56.1 | 47.2 | 70.7 | 57.0 | 27.1 | 53.5 | 51.9 |
| Mamba-790M (Gu & Dao, ) | 790M | 7.33 | 6.02 | 62.7 | 55.1 | 72.1 | 61.2 | 29.5 | 56.1 | 56.1 |
| Ours (TokenFormer-900M) | 900M | 7.38 | 5.46 | 64.0 | 55.3 | 72.4 | 59.9 | 30.6 | 56.4 | 56.4 |
| GPT-Neo 1.3B (Black et al., ) | 1.3B | - | 7.50 | 57.2 | 48.9 | 71.1 | 56.2 | 25.9 | 54.9 | 52.4 |
| Hybrid H3-1.3B (Fu et al., ) | 1.3B | - | 11.25 | 49.6 | 52.6 | 71.3 | 59.2 | 28.1 | 56.9 | 53.0 |
| OPT-1.3B (Zhang et al., ) | 1.3B | - | 6.64 | 58.0 | 53.7 | 72.4 | 56.7 | 29.6 | 59.5 | 55.0 |
| Pythia-1.3B (Biderman et al., ) | 1.3B | 7.51 | 6.08 | 61.7 | 52.1 | 71.0 | 60.5 | 28.5 | 57.2 | 55.2 |
| RWKV-1.5B (Peng et al., ) | 1.5B | 7.70 | 7.04 | 56.4 | 52.5 | 72.4 | 60.5 | 29.4 | 54.6 | 54.3 |
| Mamba-1.4B (Gu & Dao, ) | 1.4B | 6.80 | 5.04 | 64.9 | 59.1 | 74.2 | 65.5 | 32.8 | 61.5 | 59.7 |
| GPT-Neo 2.7B (Black et al., ) | 2.7B | - | 5.63 | 62.2 | 55.8 | 71.1 | 61.1 | 30.2 | 57.6 | 56.5 |
| Hybrid H3-1.3B (Fu et al., ) | 1.3B | - | 11.25 | 49.6 | 52.6 | 71.3 | 59.2 | 28.1 | 61.4 | 58.0 |
| OPT-2.7B (Zhang et al., ) | 2.7B | - | 5.12 | 63.6 | 60.6 | 74.8 | 60.8 | 31.3 | 61.0 | 58.7 |
| Pythia-2.8B (Biderman et al., ) | 2.8B | - | 5.04 | 64.7 | 59.3 | 74.0 | 64.1 | 32.9 | 59.7 | 59.1 |
| Ours (TokenFormer-1.5B) | 1.5B | 6.91 | 5.24 | 64.7 | 60.0 | 74.8 | 64.8 | 32.0 | 59.7 | 59.3 |

## Peer Reviews

confidence: 4
contribution: 4
initial_rating: 8
presentation: 4
questions: - Table 1: "The best performance highlighted for each model" seems to be misleading. It seems it is highlighted in bold for each dataset and model size what the best value is, and even then it is unclear why e.g. LAMBADA acc highlights ours-1.5B and not Pythia 1.8B (this is also the case for several other columns)
- Figure 6: Would you say the (Token)Transformer converged on the test loss respectively (left)?
- It would be nice if the practical relation between parameters, dataset size, scaling stepsize and performance was further investigated. Is there a minimum model size for which this approach only starts being useful? What were your practical considerations in choosing the steps and initializations you did? 
- Is there a different effect of scaling the channels and the dimensions?  
- Are there artifacts of the scaling steps in the final model? E.g. do the norms of the new channels (new tokens) that were initialized at zero have a different magnitude than those from previous time steps?
soundness: 3
strengths: - Innovative and flexible idea of tokenizing model parameters, that will probably inspire further experimental work in this direction 
- Clearly written, well-motivated and contextualized
- Demonstrated lower training cost in comparison of the proposed approach compared to a number of different open-source architectures and datasets
summary: This work proposes a new architectural component that can be used in transformers to replace linear projections present in the attention and feed-forward layer. A Pattention layer executes an attention mechanism between the input and some learned key and query tokens to output a transformed token. Both the channel dimensions and the number of key-query tokens can be scaled flexibly, in the sense that smaller models can be used as initial configurations for larger models. Iteratively enlarging the model results in significantly smaller training costs and similar performance to comparable architectures.
weaknesses: - The computational complexity of inference and the model itself is only dicussed late in the work. 
- The work would benefit from an explicit discussion of the architectures practical limitations.
- See questions below

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Your main argument for scaling is based on adding parameter tokens to the Pattention layers. However, this necessarily requires that you fix A) the hidden dimension, and B) the model depth. 
    - 1a. Fixing the hidden dimension can lead to representation crowding, which typically does not occur when increasing it (e.g. larger vector space for representations). Could you investigate this with your scaled models? As a limitation it should result in an upper limit to capacity, where adding new parameter tokens will no longer improve performance.
    - 1b. Fixing the model depth can have a significant impact on the range of functions the transformer can learn. It’s possible that this was not apparent due to the model size and benchmarks, but would severely limit feasibility for more complex problems, which often require larger scales (>70B). Could you address this limitation?
    - 1c. By fixing the hidden dimension, achieving improved performance would then depend on expanding the representation span within a fixed vector space (i.e. approaching 100% channels to capture the variance). This will likely have a negative impact on quantization methods, which may offer substantial computational improvements beyond what Pattention can provide alone. Could you address this?
2. The parameter token scaling method raises several questions that should be addressed.
    - 2a. When scaling the number of parameter tokens, do you observe representation collapse, or do you apply a regularization term to prevent this? Adding fixed parameters without enforced symmetry breaking often leads to vector consolidation into tight groupings. Alternatively, have you confirmed that this behavior is not occurring?
    - 2b. In Figure 2, you suggest keeping the previously learned parameter tokens unfrozen for model growth; is this necessary? It may be more beneficial to keep the previously learned parameters frozen to essentially learn “residuals”. 
3. While you present the mechanism by which Pattention works, there is no explicit connection to linear layers, or is it closer to an FFN? Could you clarify this connection in the appendix? Intuitively, if the GELU scoring is one-hot, then you have a K-V lookup.
4. Your scaling conclusions in Section 4.3, Table 3, and Figure 5 are contingent upon which term dominates in the asymptotic analysis. You require that $l \cdot n \cdot d \cdot T < l \cdot d \cdot T^2$, otherwise Pattention will become more expensive than simply using linear layers. The scaling argument should be more precise, and both acknowledge this point and demonstrate which regime the different model scales are in. Notably, a significant number of parameters were added to the 1.4B model (6.6M per matrix vs. 2.4M), where it would also be appropriate to report the FLOPs required for a forward pass of sequence length $T$, verified with empirical measurement. Furthermore, relying on a fixed $l$ while scaling $n$ may significantly impact parallelism during inference. 
5. To confirm, did you retain the original softmax attention for the MHA layers, or did you apply the GELU scoring there as well? This should be made clear to avoid confusion.
6. Questions about the Future Work section. While it is understandable to suggest exploratory directions in the future work section, some of the proposed ideas would benefit from either preliminary results or a more concrete implementation plan.
    - 6a. You mention adjustable computational cost: can you elaborate on how this could be done? Your current methodology seems to be restricted to $n_\mathrm{new}$ once scaled up, without any way to “adjust” the computational cost. 
    - 6b. The idea of combining vision and language parameter tokens is interesting; however, I am highly skeptical whether this would work. Do you have preliminary results to back up this idea? If not, perhaps the statement should be weakened, as it's not improbable that some features may have significant overlap, allowing for $\mathrm{max}(n_\mathrm{lang}, n_\mathrm{vision}) < n < n_\mathrm{lang} + n_\mathrm{vision}$, where the vision and language tasks are not simply combined after independent training. 
    - 6c. The paragraph on “Increasing interpretability” is vague. How exactly will Pattention improve interpretability? Can you make a connection to Sparse Auto-encoders (SAEs)? This seems like something that should have been added to the appendix.




Additional Comments:

-	It appears you are missing a factor of 2 in the parameter counts in Table 3. According to your description, you learn $n$ Key and $n$ Value tokens, resulting in a total of $2n$. The parameter counts and cost would then be identical to a traditional transformer when $n=d_\mathrm{model}$ and $d_\mathrm{token}=d_\mathrm{model}$. 
-	Please include a discussion connecting the GELU scoring function to your scaling ability. Scaling by setting new keys to zero relies on the GELU function, where GELU(0)=0; in contrast, exp(0)=1, which causes a norm shift.


Yao, Yiqun, Zheng Zhang, Jing Li, and Yequan Wang. ‘Masked Structural Growth for 2x Faster Language Model Pre-Training’. In The Twelfth International Conference on Learning Representations, 2024.
soundness: 3
strengths: This paper proposes a unique approach to scaling parameters and demonstrates compelling results under those conditions, including vision and language tasks. The authors further support their claims with extensive experimentation, demonstrating a substantial reduction in training costs. Additionally, the authors propose an interesting attention mechanism for their parameter tokens, which solves the issue of vanishing gradients caused by the softmax function in traditional attention. Furthermore, this paper is well-written and easy to follow, with clear figures and tables.
summary: This paper proposes an innovative modification to the standard transformer architecture by replacing all linear layers with learned tokens via attention operations. By making this alteration, the proposed architecture can effectively remove all learned parameters (except embedding layers), instead replacing them with learned tokens. This token-parameter-attention mechanism (Pattention) therefore provides an efficient mechanism for scaling parameter counts. The authors validate their architecture with several scaling experiments, on both causal language modeling and vision tasks, demonstrating comparable performance with a significantly reduced training budget. 

Overall, this paper makes a significant architectural contribution to the lineage of transformer architectures. However, the authors primarily focus on the empirical validation of their methodology without sufficient theoretical backing or exploring limitations of this approach. I am therefore inclined to reject, but would be willing to reconsider pending further discussion.
weaknesses: -	Throughout the paper, the authors seem to imply that the scaling issue is unique to transformers. However, the increased training cost with scale applies to most DNN architectures. Furthermore, they contradict themselves when attributing training cost to increasing parameters while later showing that cost is dominated by self-attention (which is well known). The proposed training improvements are therefore solely due to fixing the hidden dimension size and resuming, rather than directly due to the Pattention mechanism. The authors should make their motivation more precise, and avoid overly vague and potentially misleading arguments to support their approach.
-	The paper lacks a theoretical foundation as to why their approach works, only briefly mentioning a potential explanation in Section 5, Future Work. Although scaling parameters generally improves performance, this is typically done while also scaling the hidden dimension, which has other benefits beyond parameter count (e.g. larger representation space). While the empirical results are promising, discussing deeper theoretical insights would allow application to other (hybrid) architectures beyond TokenFormer. 
-	The authors compare against an old baseline (Net2Net) to demonstrate their model's effectiveness in growth when resuming from a smaller model. It would be more appropriate to compare with newer methods, which have been shown to perform much better (Yao et al., 2024). Furthermore, the practice of initializing extended queries to zero may negatively impact potential performance due to symmetry, where the proposed method instead relies entirely on the value token initialization. The authors should clarify why Net2Net was chosen over newer methods, or why the newer methods would not be a more suitable comparison.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
questions: same as weakness.
soundness: 4
strengths: This paper proposes the novel TokenFormer that replaces the original linear layers in transformers with Pattention layers to enable more efficient and progressive model scaling without retraining from scratch. Pattention tokenizes the dense weights of query, key, value, and output layers with flexible and scalable token parameters. It then introduces cross-attention between input tokens with those token parameters. Instead of increasing the hidden dimension in transformers, TokenFormer achieves scaling by increasing the number of token parameters, incrementally adds new token parameters, and train all old and new token parameters with less training data. Therefore, it enables effectively reusing smaller model knowledge or weights and faster training of larger models with similar architecture. Experimental results show that the training cost of a 1.4B model is significantly reduced by training a 124M model initially and progressively scaling to 1.4B parameters, compared with training from scratch of a 1.4B model with 300B tokens.
summary: The propsoed TokenFormer replaces dense linear layers of transformers with learnable token parameters to achieve progressive model scaling and reduce training cost compared with retraining from scatch of larger models with full data. The original matrix vector production is also replaced with the novel cross-attention between input token with tokenized parameters. Therefore, more token parameters can be easily added to well-trained smaller models incrementally due to flexibility of attention. The training cost is significantly reduced by 3x with similar or even better model performance. This paper proposes a interesting direction to scale LLMs and uses attention for almost all components in LLMs.
weaknesses: 1. In this paper, only results of scaling only up to 1.5B parameters are given. What are the complexity and effectiveness of scaling to larger scales such as 7B, 13B, 70B, or even more? In which case, the number of token parameters will increase multiple times and may be larger than token or model dimension.
2. The hidden dimension is relatively constrained to enable progressively scaling and reusing token parameters. Will it limit the capability of LLMs as FFN layers contain knowledge?

---

confidence: 4
contribution: 4
initial_rating: 6
presentation: 2
questions: Have the authors considered a linear attention version of Pattention, i.e. a version without the gelu activation in the cross attention? Also, it seems that the tokenformer replacement of linear layers can be made in any architecture, i.e. not just transformers. Have the authors experimented with alternative backbone architectures?
soundness: 3
strengths: The main strength of the proposed method in the paper is the flexibility of the architecture to accommodate additional parameters without jeopardizing the capabilities of the original model. Consequently there is no additional cost to "healing" the performance of the original model before continually training and allowing the model to exploit the added parameters. The experimental results reported in the paper are compelling and the key ideas are explained adequately clearly.
summary: A novel growing architecture called "tokenformer" is introduced. The idea is to replace linear layers in the transformer architecture with a cross attention mechanism where input dependent queries cross attend to key and value parameters. This framing allows for the architecture to accommodate new parameters concatenated to the original key and value parameters. The architecture and training method are validated empirically by pretraining a 124M parameter tokenformer and incrementally growing it to 1.4B parameters. The resulting model performs competitively on evaluations compared to other models of similar size pretrained from scratch.
weaknesses: The presentation of the mathematical form of the "Pattention" operation can be improved - specifically, the notation in equations 4 and 5 are confusing - the function \Theta is described as a "a modified softmax operation for stable optimization of Pattention layer" when in fact the function used in subsequent sections and experiments isn't really a softmax and what's more the function name \Theta doesn't appear. Also it is unclear how equation 5 connects to equation 4, i.e. what exactly the function f has to do with \Theta. Perhaps it would be better to remove any reference to \Theta, and to solely describe the operation in terms of the function f, and to make explicit the choice of the specific form of f clear. Ideally a glossary defining all mathematical symbols in the paper would really help. 
 In fact, the authors should make it clear that the structure of the Pattention operation can be seen as an MLP with standard activations: GeLU(RMSNorm(XK)V). The analogy with attention is interesting, but it somewhat obscures the key mathematical form of the operation - especially since there are no tokens (i.e. data) involved in the K and V parameters in Pattention. Also it isn't made clear how the V parameters are expanded - are they also padded with 0? 

Also, in addition to the comparison with the Net2Net method, have the authors considered comparing to say the method in  arXiv:2409.12903v2 [cs.CL]  which also aims directly at growing transformer models?

Meta-review: Summary
This paper introduces TokenFormer, an innovative framework that replaces linear layers in transformers with a novel "Pattention" mechanism. The proposed architecture tokenizes model parameters and employs attention to facilitate interactions between input tokens and model parameters, enabling incremental scaling of models without retraining from scratch. This approach significantly reduces training costs while maintaining or surpassing baseline performance. The paper is well-executed and presents strong empirical results, making a substantial architectural contribution to model scaling and efficient resource utilization.

Decision
Based on the reviews and discussions, I recommend accepting this paper for a spotlight presentation. The paper is a significant step forward in scalable and efficient transformer architectures, and its findings are likely to inspire further research.

Decision: Accept (Spotlight)