TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters

We are glad the reviewer found the ideas in the paper compelling and clear. Below are our responses to the questions and suggestions raised by the reviewer.

Q1: The presentation of the mathematical form of the "Pattention" operation can be improved.

A1: Thank you for your feedback. We have removed the reference to $\Theta$ and added the notion table in the paper as requested. Please refer to lines 184-187 and lines 199-203 of the main paper for more details.

Q2: 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.

A2: Thanks. More discussions have been added to the main paper. Please refer to section G.1 in the appendix and to the general answer above. Our goal is to provide a unified framework for handling computations in models by treating parameters as tokens and using the attention mechanism to manage interactions between different types of tokens (e.g., data tokens and parameter tokens). We started by drawing an analogy between a two-layer MLP and attention, and used this module as the fundamental computational unit to design the model, rather than the previous linear projection, thereby unifying all the computations within a single framework. Then it was applied to incremental model scaling, enabling the expansion of a model based on a pre-trained one, significantly reducing training costs.

Furthermore, the concept of tokenizing all components and leveraging attention to establish relationships has the potential to impact topics beyond pretraining models. For instance, it can be applied to the hidden states in linear models, allowing for incremental memory expansion while addressing the issue of excessive information loss often observed in such models.

Q3: It isn't made clear how the V parameters are expanded - are they also padded with 0?

A3: The value parameters are initialized randomly. More clarifications have been added to the main paper. Please see line 269 in the main paper.

Q4: 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?

A4: Thank you for pointing out this baseline. We implemented the method from this paper and conducted a preliminary experiment. The results are as follows:

The unit of Flops is $10^4 \times T^2$. Our method outperforms HyperCloning while requiring less computational effort for token-token interactions. Due to time limitations, we scaled our model only up to 757M parameters. The complete comparison with this baseline will be included in the final camera-ready version.

Q5: Have the authors considered a linear attention version of Pattention, i.e. a version without the gelu activation in the cross attention?

A5: Thank you for this valuable suggestion. It highlights an excellent opportunity to further minimize the computational overhead associated with Token-Parameter interactions. As part of our future research, we will explore methods to reduce this cost even further, including techniques like Mixture-of-Experts (MoE) or linear attention.

Q6: 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?

A6: Thanks. Due to time constraints, we will explore additional backbone architectures in the future to broaden our approach. The concept of tokenization is not limited to parameters—it can also be applied to memory and hidden states. For instance, if the hidden states in models like Mamba or RNN are treated as tokens and processed using the attention mechanism, they can be expanded arbitrarily, thereby addressing the issue of excessive information compression in linear models.

Thank you very much for your valuable suggestions and kind acknowledgment of our work! The complete experimental results of hypercloning (from 124M to 1.4B) will be included in the final version.

We sincerely thank the reviewer for providing thoughtful review and positive feedback. Below are our responses to the questions and suggestions raised by the reviewer.

Q1: 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.

A1: Pretraining LLMs requires huge computational resources, and given our own resource limitations, the model in the paper is scaled up to 1.4B. We are also exploring scaling to larger models and integrating the MoE framework to design even bigger and more powerful models. The results of larger models will be included in camera ready version.

Q2: 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?

A2: Our experiments have not observed a bottleneck in the hidden dimension. OpenAI's scaling laws [1] show that model capacity scales with size following a power law, with minimal impact from architectural details like width or depth. Please see lines 3-4 of the abstract in OpenAI's scaling laws paper.

Moreover, expanding parameters using the mixture-of-experts (MoE) approach while keeping the hidden dimension fixed is also a common industry practice [2,4]. TokenFormer is similar to this approach. This approach is validated by Krajewski et al. [3] and Fedus, William, et al [4], who found that increasing the number of experts while maintaining a fixed hidden dimension aligns with scaling laws for dense models. Please check Figure 1 of the popular switch transformer for more details..

[1] Kaplan, Jared, et al. "Scaling laws for neural language models." arXiv preprint arXiv:2001.08361 (2020).

[2] Dai, Damai, et al. "Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models." arXiv preprint arXiv:2401.06066 (2024).

[3] Krajewski, Jakub, et al. "Scaling laws for fine-grained mixture of experts." ICLR 2024 Workshop Oral.

[4] Fedus, William, Barret Zoph, and Noam Shazeer. "Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity." Journal of Machine Learning Research 23.120 (2022): 1-39.

Q12: 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.

A12: We would like to clarify a potential misunderstanding. The terms $l \cdot n \cdot d \cdot T$ and $l \cdot d \cdot T^2$ determine whether Token-Token or Token-Parameter computation dominates, which is irrelevant to whether pattention or linear layers have a higher computational cost. For the token-parameter part, the computational overhead is proportional to the number of network parameters, so it is the same whether pattention or linear layers are used. However, since we can scale the parameter size without increasing the hidden dimension, pattention significantly reduces the computational overhead for token-token interaction compared to a linear layer. As a result, compared to linear layers, the overall computational overhead is greatly reduced in long context modeling.

Q13. Furthermore, relying on a fixed $l$ while scaling $n$ may significantly impact parallelism during inference.

A13: In our experiments, we did not observe this issue. The specific training TPU hours are shown in the table below.

All the models are trained on 30B tokens with a sequence of 1024 on TPU v4 hardware. Since the sequence length here is 1024, the computational cost is primarily in the Token-Parameter interaction, allowing for an approximate comparison of the issues related to matrix parallelization. The table shows that cost scales with parameter count and is unrelated to the matrix shape. Therefore, fixing the hidden dimension and increasing the number of parameter tokens won't impact parallelization.

Q14: 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.

A14: Thanks. As outlined in Equations 8–10, we replaced all linear projections with pattention while preserving the standard attention mechanism for token-token interactions. This means that the standard softmax operation remains unchanged.

Q15. You mention adjustable computational cost: can you elaborate on how this could be done? Your current methodology seems to be restricted to $n_{\text{new}}$ once scaled up, without any way to “adjust” the computational cost.

A15: Unlike Transformers, which can only scale by increasing the hidden dimension at a given network depth, our approach decouples the hidden dimension from the model size and offers a new scaling dimension: the number of parameter tokens. Thanks to this new axis, the hidden dimension can be freely configured regardless of model size, thereby enabling adjustable control over the cost of token-token interactions.

Q6: 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?

A6: We did not observe representation collapse because our method of adding parameters is similar to LoRA and does not present symmetry issues. Similar to LoRA, our approach builds on the original model, with value parameters randomly initialized. This provides diverse gradients for different parameter tokens, effectively avoiding symmetry issues. Even if the key parameters are initially set to all zeros, by the end of training, the norm of the new key parameters is on the same order of magnitude as the old parameters.

Q7: 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.

A7: Currently, the Transformer is the dominant general-purpose architecture due to its flexibility and versatility. It has achieved remarkable success across various domains, particularly in NLP, where model scaling is highly emphasized. Therefore, we focus primarily on the Transformer and push its flexibility to the limit. However, the concept of tokenizing everything is not exclusive to Transformers and can also be applied to linear models, like Mamba and RWKV. Specifically, if the hidden states in linear models are treated as tokens and processed by attention, they can be expanded arbitrarily, thereby having the potential to address the problems of excessive information compression.

Q8: 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).

A8: Whether Token-Token or Token-Parameter interactions dominate the computational overhead depends on the sequence length. TokenFormer effectively reduces the computational cost of both. Research on model reuse and MoE has targeted reducing Token-Parameter interaction costs during training and inference, while methods like Mamba and linear Transformers focus on optimizing Token-Token interaction costs. Our TokenFormer combines the strengths of these two research directions, enabling incremental model reusing with potential applications in sparse inference. Additionally, it offers the advantage of controllable Token-Token Interaction costs, making it well-suited for long-text modeling tasks.

Q9: 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.

A9: We use an attention-like mechanism to model Token-Parameter interaction, replacing linear projection. This enables the network’s Token-Parameter interaction to handle a variable number of parameters, similar to how attention processes a variable number of input tokens. Based on this design, we incrementally scale the model by adding new parameter tokens to the original one and initializing key parameters to 0, allowing us to fully resume training from pre-trained checkpoints — a capability not possible with transformers.

Q10: 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”.

A10: In our experiments, we found that freezing old parameters leads to poorer performance, as shown in the table below:

Freezing old parameters will reduce the effective training budget, leading to poorer performance under the same number of training iterations.

Q11: 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.

A11: Thanks. The discussion has been added to the main paper. Please see the sections G.1 and G.2 in the appendix.

Q4: 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). The authors should clarify why Net2Net was chosen over newer methods, or why the newer methods would not be a more suitable comparison.

A4: Net2Net is a widely used method in the field of model reuse, valued for its simplicity and efficiency. It remains the most common strategy in the industry [1,2], requiring no extra loss functions or operations—only using parts of a smaller model to initialize a larger one. This simplicity makes it our primary baseline.

To better validate the effectiveness of our method, we included HyperCloning [1] as an additional baseline. This is the most recent paper on model scaling (published two months ago), and it was also suggested by Reviewer UJD1 as a strong baseline. All experiments were conducted on OpenWebText2, and the specific perplexities are as follows:

The unit of Flops is $10^4 \times T^2$. Our method still outperforms the latest HyperCloning and has lower computational complexity in the token-token interaction part. More baselines will be included in the camera-ready version.

[1] Samragh, Mohammad, et al. "Scaling Smart: Accelerating Large Language Model Pre-training with Small Model Initialization." arXiv preprint arXiv:2409.12903 (2024).

[2] Wang, Peihao, Rameswar Panda, and Zhangyang Wang. "Data efficient neural scaling law via model reusing." ICML, 2023.

Q5: 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?

A5: Simply scaling the model by increasing its depth is unsustainable. Very deep networks often suffer from gradient vanishing or exploding problems, which make training exceedingly difficult. Consequently, even X-AI's Grok-1 model with 314 billion parameters is limited to just 64 layers. We offer an alternative dimension to scaling the model beyond simply increasing its depth.

Notably, our approach does not conflict with methods that stack the depth. TokenFormer can also be scaled in depth by initializing new layers on top of the base model to further expand the model depth.

The table shows that TokenFormer can not only scale the model by increasing the number of parameter tokens but also be compatible with depth expansion.

Q2: 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.

A2: Given a fixed number of model parameters, whether a larger representation space leads to greater expressive power remains an open problem. In our experiments, we did not encounter the bottlenecks in the hidden dimension. With the same data and training budget, OpenAI's scaling laws [1] reveal that model capacity follows a power-law relationship with model size. The architectural details, such as network width or depth, have minimal impact on this relationship. Please see lines 3-4 of the abstract in OpenAI's scaling laws paper.

Additionally, expanding model parameters through the mixture-of-experts (MoE) approach—where the hidden dimension remains fixed while more experts are added—has become a popular and successful practice in the industry [2,3]. TokenFormer represents a variant of MoE. Fedus, William, et al [2] and Krajewski, Jakub, et al. [4] have demonstrated that by gradually increasing the number of experts without changing the hidden dimension, models can adhere to the scaling laws corresponding to dense models. This experimental conclusion validates the feasibility of our method. Please check Figure 1 of the popular switch transformer for more details.

Notably, our approach does not conflict with methods that expand the hidden dimension; instead, we provide an alternative scaling dimension beyond the hidden dimension. In Transformers, the qkv (query, key, value) and output projection matrices are square, meaning the channel and hidden dimensions are linked. Given a fixed network depth, this linkage forces Transformers to scale the model only by increasing the hidden dimension, which significantly raises inference costs. Such scaling is unsustainable in practical applications with limited computational resources. Our approach offers a new scaling dimension that is separate from the hidden dimension, helping to keep token-token interaction costs manageable when modeling long sequences.

[1] Kaplan, Jared, et al. "Scaling laws for neural language models." arXiv preprint arXiv:2001.08361 (2020).

[2] Fedus, William, Barret Zoph, and Noam Shazeer. "Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity." Journal of Machine Learning Research 23.120 (2022): 1-39.

[3] Dai, Damai, et al. "Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models." arXiv preprint arXiv:2401.06066 (2024).

[4] Krajewski, Jakub, et al. "Scaling laws for fine-grained mixture of experts." ICLR 2024 Workshop Oral.

Q3: The paper lacks a theoretical foundation as to why their approach works, only briefly mentioning a potential explanation in Section 5, Future Work. While the empirical results are promising, discussing deeper theoretical insights would allow application to other (hybrid) architectures beyond TokenFormer.

A3: The theoretical analysis of LoRA's representational capacity [2] can provide theoretical insights into our method to some extent. As described in Equations 4-13 of main paper, given input $X$ and learnable tokens $(K_P, V_P)$, we formulate token-parameter interaction as attention, $\Theta\left(X \cdot K_P^{\top}\right) \cdot V_P$. With this naturally expandable design, the model can be incrementally scaled by appending new key-value parameter tokens, $\Theta\left(X \cdot [K_P^{\text{old}} , K_P^{\text{new}}]^{\top}\right) \cdot [V_P^{\text{old}}, V_P^{\text{new}}]$. The effectiveness of this approach can be partly understood through the analysis of LoRA [1], as TokenFormer can be viewed as a high-rank extension of LoRA, where $K_P^{\text{new}}$ and $V_P^{\text{new}}$ are equal to $A$ and $B$ matrix. LoRA has been widely applied in foundation models, and its representational capacity has been theoretically proven by Zeng Y et al [2], which to some extent explains why TokenFormer’s incremental scaling can be effective. This preliminary analysis has been added to the main paper. Please see G.4 in the appendix for more details.

[1] Hu, Edward J., et al. "Lora: Low-rank adaptation of large language models." ICLR2022

[2] Zeng, Yuchen, and Kangwook Lee. "The expressive power of low-rank adaptation." ICLR2024

We sincerely thank the reviewer for providing thoughtful review and constructive suggestions. We answer the questions as follows.

Q1: The practice of initializing extended keys to zero may negatively impact potential performance due to symmetry, where the proposed method instead relies entirely on the value token initialization.

A1: We thank the reviewer for raising an interesting point. Initializing the key parameter tokens to zero does not lead to symmetry issues. This approach is validated by the successful LoRA [1], where the first layer is also initialized to all zeros, while the second layer is initialized with Gaussian random values. This strategy minimizes noise during the initial training phase and preserves the original network's distribution. The symmetry problem can be avoided by ensuring the update process builds upon a pre-trained model and that the value parameter tokens are randomly initialized. Firstly, the random initialization of the value parameters adds diversity to the gradients. Secondly, the newly added parameters are incrementally updated on top of the well-trained base model. Consequently, the key parameters will quickly evolve into non-zero values during training.

We further conducted validation experiments on OpenWebText2, scaling 124M model to 354M. The impact of different initializations and the final statistics of the key parameters in last Pattention layer are as follows:

As shown in the table above, even when the key parameter is initially set to zero, it quickly deviates from zero during training (increasing from 0 to a standard deviation of 0.0311 by the end of the process).

[1] Hu, Edward J., et al. "Lora: Low-rank adaptation of large language models." ICLR2022

Q16. 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 $\text{max}(n_{\text{lang}}, n_{\text{vision}}) < n < n_{\text{lang}} + n_{\text{vision}}$, where the vision and language tasks are not simply combined after independent training.

A16: We have conducted preliminary experiments on vision-language tasks, and the results are promising. Our approach introduces bridging parameter tokens to connect the vision and language parameters. Specifically, image tokens are allowed to attend only to image parameter tokens and bridging parameter tokens, while language tokens interact exclusively with language parameter tokens and bridging parameter tokens. However, tokens from all modalities can interact freely in the token-token attention mechanism.

The models used include a TokenFormer pre-trained on ImageNet for vision modality and another TokenFormer pre-trained on OpenWebText2 for language modality, both at the base model size. For validation, we adopted the widely used BLIP's training data settings. The COCO captioning results are summarized below:

The table above shows that our TokenFormer can perform vision-language tasks and achieve performance comparable to classic vision-language methods.

Q17. 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.

A17: The interpretability discussed here is not related to sparse autoencoders but rather to research analyzing the interpretability of Feed-Forward Networks (FFNs). Geva, Mor, et al. [1] modeled FFNs as key-value lookup mechanisms, interpreting FFNs in Transformers as key-value memory structures where each key-value pair encodes a learned pattern. Different tasks activate and utilize these key-value memories in unique ways. This discussion has already been incorporated into the paper. Please see section G.2 in the appendix.

[1] Geva, Mor, et al. "Transformer feed-forward layers are key-value memories." EMNLP (2021).

Q18. 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_{\text{model}}$ and $d_{\text{token}} = d_{\text{model}}$.

A18: Sorry for this typo. We have already corrected it in the main paper. Please refer to the updated Table 3 in the main paper. The computational overhead of Token-Parameter interactions is strictly proportional to the number of network parameters. Therefore, when the parameter count is kept the same, the computational cost for this part is identical between TokenFormer and Transformer. However, our method decouples the token dimension from the model dimension, allowing the computational cost of Token-Token interactions to be independent of the model's parameter count. This ensures that scaling the model does not increase the computational burden of Token-Token interactions.

Q19. 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.

A19: Thanks to the reviewer for raising that interesting observation. By replacing the exponential function in softmax with the GeLU function, we can add new parameters while fully resuming the checkpoint of the previous model. For the exponential function, if $\text{exp}(x) = 0$, $x$ would need to be negative infinity. It is not feasible to achieve this property by initializing the key parameters, making it impossible to ensure that newly added key parameters have no impact on the original key parameters. This is one advantage of using GeLU. More clarification has been added to section B (lines 769-773) in the appendix.

Below, we address the specific responses one by one.

A2: We appreciate the reviewer’s agreement with our perspective. It is possible that our models have not yet reached their saturation point, and further investigation is needed to explore the representation span within a fixed vector space in other use cases. As the reviewer noted, extending the saturation point further could be a significant advantage of pattention compared to standard linear projections.

As the reviewer mentioned, the analogy to MoE is exactly because "MoE models scale the hidden dimension of the base model and then hold this hyperparameter fixed while increasing the number of experts". In this context, each key-value parameter pair (or group) in tokenformer could be interpreted as a generalized expert, and TokenFormer scales these experts. This conceptual similarity is why we suggested in our future work that TokenFormer could be extended into an extreme form of MoE, where each key-value parameter pair operates as an individual expert.

Regarding inference costs, we will carefully rephrase the relevant sentences and include additional discussions on short-sequence scenarios of visual modeling (e.g., image classification and CLIP). However, it is worth noting that long-sequence scenarios also have wide applications in visual modeling. For example, high-resolution scenarios such as semantic segmentation. The introduction of Swin Transformer [1] and window-based ViT [2] was precisely to address the enormous computational cost of ViT when processing high-resolution images. There is also an increasing amount of work [3] focused on accelerating vision-language models. For example, FastV [3] addresses the issue of slow inference encountered by LLaVA when handling multiple image inputs or high-resolution vision-language tasks. Long sequence modeling is equally important in vision.

[1] Liu, Ze, et al. "Swin transformer: Hierarchical vision transformer using shifted windows." Proceedings of the IEEE/CVF international conference on computer vision. 2021.

[2] Kirillov, Alexander, et al. "Segment anything." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

[3] Chen, Liang, et al. "An image is worth 1/2 tokens after layer 2: Plug-and-play inference acceleration for large vision-language models." European Conference on Computer Vision. Springer, Cham, 2025.

A4: A4 was added to address the question regarding the newest baselines as requested by the reviewers, where we included the latest methods (from two months ago) to supplement our experiments. To further address the reviewer's concerns, we additionally conducted experiments replacing the linear layer with FFNs (using an equivalent parameter count to replace the linear layer). The results in the table further validate the effectiveness of our approach.

Since we have operations similar to softmax, the scale and variance of the pattention's output is effectively bounded, achieving better results.

We sincerely thank the reviewer for the prompt response, allowing us to have a better discussion.

Comparison with LoRA: We thank the reviewer for pointing this out. From this perspective, using LoRA as an analogy to model scaling is indeed inappropriate. To address this, we have removed the LoRA analogy in the revised version.

We also appreciate the reviewer for highlighting that our A16 results can be analogized to LoRA, as this insight strengthens the contribution of our paper. This analysis, integrated with vision-language discussions, will be included in the camera-ready version.

A6: We thank the reviewer for the valuable suggestions, which have helped us better formulate this metric.

The table above shows normalized values, where squared scores have been divided by their sum. Since the number of parameter tokens for the new parameters (2140) and the old parameters (576) is different, and the normalized values are around 1/n, we use different scales to depict the two distributions. From the table above, the distributions of the new parameters and the old parameters are roughly similar.

Additionally, we have also followed the reviewer's suggestion to consider the top-1 tokens and compute the normalized entropy over that distribution: diversity = $H / log_2(N)$.

From the table above, it can be seen that the new parameters exhibit even greater diversity, confirming that we did not experience a collapse in representation.

A15: We thank the reviewer for providing this valuable suggestion and perspective. We are currently exploring this direction, and your suggestion has been highly inspiring, greatly aiding the future extension and application of our method.

We sincerely appreciate the time and effort the reviewer has invested in evaluating our work. As we approach the end of the discussion phase, we kindly ask the reviewer to consider if our rebuttal might warrant a reevaluation of the current score. We truly value the reviewer's assessment and are grateful for the constructive feedback provided.

A9: The training improvement (reduced cost) arises from Pattention’s ability to reuse pre-trained parameter tokens, leveraging prior knowledge to achieve the same performance in fewer training steps.

We also provided a method of fixing the hidden dimension and then resuming, as mentioned in A4: replacing the linear layer with FFNs (using an equivalent parameter count to replace the linear layer).

Since we have normalization similar to softmax, the scale and variance of the pattention's output is effectively bounded, achieving better results.

Pattention provides a new dimension for scaling, which does not conflict with scaling hidden dimension. In the table below, we provide experiments that scale both the hidden dimension and parameter tokens simultaneously.

From the table above, we can see that our method is compatible with scaling the hidden dimension while also providing an additional scaling dimension beyond the hidden dimension. The advantages compared to previous methods are: (1) faster convergence when incorporating new data (Fig. 6); (2) suitable for long context modeling (Fig. 5); and (3) improved performance (Fig. 7 and above table).

A10: As mentioned above, to address the reviewers' concerns, we have removed the analogy to LoRA. However, we want to emphasize that the core idea of LoRA lies in using learnable low-rank matrices to perform incremental updates on the original model, which is similar to our core idea of parameter reusing. Freezing the old parameters is primarily for parameter efficient tuning and memory savings. In my practical use of the successful LLaVA model (for vision-language task with LoRA), not freezing the base model can achieve slightly better performance, but it requires significantly more memory and computational resources.

A12: As explained in A8, if TokenFormer and Transformer are required to have the same parameter count and the same number of layers, $n$ can only be smaller than $d$. This is why TokenFormer chooses $n=576$ and $d=768$ to align with the 124M parameter count of Transformer.

Our understanding is that under the condition of Transformers and TokenFormer having the same parameter count, identifying which term is dominant is not critical. Because the $N$ of both are equal and the comparison of FLOPs is only related to $n_{\text{layer}}^{\text{transformer}} \times d_{\text{transformer}}$ and $n_{\text{layer}}^{\text{tokenformer}} \times d_{\text{tokenformer}}$. Indeed, as the reviewer mentioned, there exists a point that we can adjust depth and width to achieve $n_{\text{layer}}^{\text{transformer}} \times d_{\text{transformer}}$ < $n_{\text{layer}}^{\text{tokenformer}} \times d_{\text{tokenformer}}$, but this point typically occurs when the model parameter count is small and requires $n_{\text{layer}}^{\text{transformer}}$ to be as small as possible and $d_{\text{transformer}}$ to be as large as possible. Such configurations are not meaningful, and as the network scale grows, this condition is difficult to satisfy. For more explanation, please refer to the response to A8.

Moreover, Figure 5 presents experiments on model scaling starting from a standard 124M base model (depth=12, hidden dim=768) while keeping the network depth unchanged. This setting does not have such a turning point, which makes it meaningful.

We will include this discussion in the camera-ready version.

A8: Thanks the reviewer for helping make our explanation clearer. We want to clarify that under the premise of the same model parameter count and network depth, $d_0 \cdot n$ will not be greater than $d^2$. When $d_0 \cdot n > d^2$, it is natural for TokenFormer to have more parameters and therefore a higher computational cost. This is why TokenFormer chooses $n=576$ and $d=768$ to align with the 124M parameter count of Transformer. All our discussions are based on the premise that TokenFormer and Transformer have the same parameter count.

If the network depths of Transformer and TokenFormer are allowed to differ, we acknowledge there is a point where the computational cost of pattention exceeds that of a linear layer. However, this only occurs when the model size is very small, resulting in an impractical network structure that is extremely shallow and overly wide. Such configurations are not meaningful, and as the network scale grows, this condition is difficult to satisfy. Additionally, it seems there may be a misunderstanding by the reviewer. The sequence length being processed does not determine whether TokenFormer or Transformer has fewer FLOPs; it only determines how many FLOPs are saved. Once TokenFormer's FLOPs are less than Transformer's, starting from a sequence length of 0, the longer the sequence, the greater the FLOPs savings. Let us clarify.

1. We analyze this using the two formulas on the right side of the last row in Table 3. If the parameter counts of the Transformer and TokenFormer are the same, meaning $N$ in both formulas is equal, then the additional computational cost of the Transformer compared to the TokenFormer comes only from the token-token interaction, which is $4 \times (n_{\text{layer}}^{\text{transformer}} \times d_{\text{transformer}} - n_{\text{layer}}^{\text{tokenformer}} \times d_{\text{tokenformer}}) \times T^2$. With the model scaling of Transformers, both the depth $n_{\text{layer}}^{\text{transformer}}$ and width $d_{\text{transformer}}$ of the network increase ($n_{\text{layer}}^{\text{transformer}} \times d_{\text{transformer}}$ is larger and larger), whereas TokenFormer expands the model using other dimensions ($n_{\text{layer}}^{\text{tokenformer}} \times d_{\text{tokenformer}}$ is constant), making this advantage increasingly significant. From the above formula, it can also be seen that the sequence length $T$ only determines how many FLOPs are saved; it does not determine whether the computational FLOPs of Pattention are greater or smaller than those of the linear layer.

2. Indeed, as the reviewer mentioned, there exists a point where $n_{\text{layer}}^{\text{transformer}} \times d_{\text{transformer}}$ < $n_{\text{layer}}^{\text{tokenformer}} \times d_{\text{tokenformer}}$, but this point typically occurs when the model parameter count is small. Let us clarify. Under the premise of having the same parameter count $N$ of transformer and tokenformer, if we want $n_{\text{layer}}^{\text{transformer}} \times d_{\text{transformer}} < n_{\text{layer}}^{\text{tokenformer}} \times d_{\text{tokenformer}}$ when the model continues to scale up, this requires $n_{\text{layer}}^{\text{transformer}}$ to be as small as possible and $d_{\text{transformer}}$ to be as large as possible. This is because the parameter count of transformer is $12 \times n_{\text{layer}}^{\text{transformer}} \times d_{\text{transformer}}^2$. This model configuration greatly limits the depth of the network, as the reviewer mentioned, which can affect the network's performance. Moreover, as the network's parameter count increases, it will become increasingly difficult to satisfy this condition.

We will include this content in the camera-ready version to make our descriptions more accurate.

A5: Thanks for the reviewer's recognition. Given the laboratory's resources, we are unable to validate TokenFormer at the scale of 300B (e.g., Grok-1). However, it is an academically recognized fact that simply increasing depth is not a sustainable approach for model scaling.

As the reviewer agreed, the additional experiment we provided suggests that this limitation may not apply to TokenFormer. For further clarity, we include another baseline of TokenFormer initialized to the expanded configuration in the below Table. It shows that the scaling results are indeed comparable.

The results in the table show that our extended configuration still performs slightly better than the Transformer under the condition of aligned parameter counts. Additionally, due to parameter reusing, it achieves much better performance than training from scratch when training with the same number of data tokens (test loss: 2.65 vs 2.73), validating the effectiveness of our reusing approach.

A6: Since we use a modified softmax (i.e., replacing exponential with gelu), the pattention scores can be negative, making entropy unsuitable as a metric. Instead, our strategy is as follows: for the pattention scores with shape of [batch size, number of parameter tokens], we square the scores and then take the mean along the batch dimension, resulting in a [1, number of parameter tokens] matrix. We then analyze the distribution of the mean corresponding to each parameter token.

The table above shows the distribution of the newly added parameters and the old parameters after finishing the model scaling. It can be observed that the distributions of the new and old parameters are similar, verifying that the newly added parameters did not encounter representation collapse.

We sincerely thank the reviewer for their detailed feedback and specific suggestions. Below, we address the key concerns:

Theoretical Motivation in the Main Paper: We appreciate the reviewer’s valuable suggestions. In the current version of the paper, this analysis has been included in the appendix due to its length and the space constraints of the main text. Incorporating it into the main paper within the page limit would require extensive rewriting. Due to the deadline for PDF revision, we assure the reviewer that we will thoroughly revise and integrate this theoretical analysis into the main text for the camera-ready submission.

Our paper’s key contribution is the proposal of a novel architecture for foundation models, validated through extensive experiments. This architecture is inherently scalable and reduces training costs by reusing parameters. Notably, the theory of deep network remains an open problem. While more solid theoretical backing would undoubtedly strengthen our paper, we believe this is a promising direction for long-term exploration in the future.

Inconsistencies in Theoretical Analogies: Our pattention is indeed different from LoRAs, MoEs, and FFNs, both in theory and in specific implementation. This is why we use the term "analogy". This approximate analogy is interesting (commented by reviewer UJd1) and helps readers understand at a high level why the pattention mechanism is effective. To address the reviewer's concern, the analogy of LoRA has been removed (G.4). Please check the revision.

Contradictory Evidence from Additional Experiments: The core idea of LoRA lies in using learnable low-rank matrices to perform incremental updates on the original model, which is similar to our core idea of parameter reusing — using learnable residual matrices (new parameter tokens) to incrementally update the original model. When this idea is specifically applied to parameter efficient tuning, it becomes LoRA. The corresponding specific designs are as follows:

1. Freezing the original parameters is intended to save memory costs and ensure parameter efficiency.
2. Applying it to linear matrices allows the trained low-rank matrices to be directly merged into the original parameters without adding new inference overhead.
3. Partial zero initialization is designed to ensure stable training.

Implications for Baseline Comparisons: We appreciate the reviewer highlighting this point. In response, we have conducted an additional experiment comparing TokenFormer and Transformer-replacing linear projections with FFNs of equivalent size. The results, shown in below Table, demonstrate that TokenFormer still outperforms this baseline, underscoring the superiority of the Pattention mechanism.

Q7: 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?

A7: The table below presents the norm of the new and old key parameter tokens in the last Pattention layer after incrementally scaling the TokenFormer model from 124M to 354M.

We observed that even when the key parameters are initially set to zero, they gradually evolve during training, and the new and old parameters eventually converge to the same order of magnitude. This behavior is consistent with experimental observations in LoRA. Moreover, this manner not only fully resumes the previously learned distribution of the smaller model, but also incorporates the additional learning capacity.

We are glad the reviewer found the ideas in the paper innovative, which will probably inspire further experimental work in this direction. Below are our responses to the questions and suggestions raised by the reviewer.

Q1: The computational complexity of inference and the model itself is only discussed late in the work.

A1: Thanks for the valuable comments. More discussion has been added to the revision. Please refer to lines 293-295 in the main paper for more details. TokenFormer has two major advantages: controllable token-token computation overhead for efficient long-context modeling and natural suitability for selective token-parameter interaction in sparse inference.

Q2: The work would benefit from an explicit discussion of the architectures practical limitations.

A2: The corresponding discussion has been added to the appendix. Please see section G.3 in the appendix. In summary, while Token-Token interaction complexity is controllable, Token-Parameter computation grows linearly with network scale, making it challenging to expand parameters without increasing inference overhead. However, TokenFormer's natural suitability for sparse inference, as each key-value parameter pair functions as an individual expert. This design inspires future integration with MoE to build more efficient and powerful networks.

Q3: 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.

A3: In Table 1, we present a comparison of the Transformer-based 1.5B and 3B models as a group. To avoid the misleading issue, we have made adjustments in the revised version. Please refer to the updated Table 1 in the main paper.

Q4: Figure 6: Would you say the (Token)Transformer converged on the test loss respectively (left)

A4: As explained in Section 4.3, both TokenFormer and Transformer were trained on Enwik8 for a single epoch. TokenFormer demonstrated rapid convergence within this single epoch, while Transformer, which employs the Net2Net method, was unable to effectively resume from the last checkpoint, resulting in large loss jumps and slower convergence. After one epoch, the Transformer had not yet converged, supporting our claim that TokenFormer's scaling method enables faster convergence.

Q5: 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?

A5: Thanks for this valuable suggestion; this is indeed a direction worth exploring in the future. TokenFormer's scaling method is highly versatile and adaptable to any model size. Currently, we simply follow the standard Transformer configurations: 124M, 354M, 757M, and 1.4B, which are also reported in concurrent works like Mamba.

Q6: Is there a different effect of scaling the channels and the dimensions?

A6: Scaling the hidden dimension causes larger loss jumps than the channel (i.e., the number of parameter tokens). The table below shows the loss increases observed at the moment new parameters are added through the hidden dimension and channel respectively. The experiment was conducted in the setting of scaling TokenFormer from 124M to 354M, with the newly added key parameters initialized in two ways: zeros and random initialization.

As shown in the table above, scaling the channel (increasing parameter tokens) causes smaller loss jumps compared to scaling the hidden dimension. This is because scaling the channel resembles LoRA, learning a residual that minimally impacts the original model's distribution. When key parameters are zero-initialized, like in LoRA, the previously learned distribution is fully preserved, preventing any loss jumps.

Thank you for your valuable suggestions, which greatly improve the quality of our work. When the new scaling experiments are completed (which we are currently working on, e.g., larger model size), we will include more discussion about the practical relation between parameters, dataset size, scaling stepsize, and performance in the camera-ready version.

In my view, the primary mathematical distinction of Pattention compared to FFN lies in its use of alternative activation functions and normalization techniques. However, its most notable innovation is the conceptual shift of treating model parameters as tokens, thereby unifying token-token and token-parameter interactions within a single framework.

As mentioned by Xingkui Zhu and the general response, our approach introduces a unified concept for computation by representing all components, including parameters, as tokens. Relationships between these tokens—such as data and parameter tokens—are governed by attention mechanisms.

We start with the formulation that a two-layer MLP can be expressed using attention, and use this module as the basic computational unit to design the neural network, rather than the linear projection, thereby unifying all the computations within a single framework.

With this all-attention architecture, TokenFormer enables scalable model expansion without retraining from scratch. Key-value parameter tokens act as learnable patterns, allowing incremental model growth by adding tokens to the parameter set. This method achieves performance comparable to training from scratch but at a lower computational cost.

This tokenization approach extends beyond pretraining and has implications for other areas, such as enabling memory expansion and mitigating information loss in linear models (e.g., Mamba and RWKV)

Thank you for raising this very valuable question. First, if both layers are zero-initialized, it would lead to a decrease in expressive capability. To ensure diverse gradients, at least one layer must be randomly initialized, or alternatively, all layers are randomly initialized.

The table below presents a performance (i.e., perplexity) comparison of FFN expansion. To maintain fairness with Net2Net (random initialization), all new parameters in these methods were randomly initialized. The FFN expansion method underperforms compared to our approach and even slightly falls short of Net2Net.

The table shows that expanding only the FFN slightly degrades performance, likely due to the lack of scaling in the qkv and output projection layers. This imbalance limits the expressive power of the self-attention block and weakens token-token attention. We argue that parameter scaling should be uniform—each token-parameter interaction module should scale proportionally. Otherwise, unscaled components create bottlenecks in the model's expressive capability.

The FFN expansion approach is a simplified version of our method, where only the FFN is replaced with Pattention, while the other linear projection components remain unchanged. Therefore, in terms of model flexibility, it is inferior to TokenFormer. From a computational efficiency perspective, the training time of a Transformer, the FFN expansion approach, and TokenFormer are similar, being roughly proportional to the number of parameters.

I sincerely appreciate the authors' thorough response and the considerable effort they put into addressing my questions with additional analysis.