Questions about model scaling

[Went-Liang]

This paper is excellent! I have some questions to consult:

1. Is the d_token mentioned in the paper the same as hidden_size?
2. Section 3.3 of the paper mentions: The dimension d of the new parameters is the same as that of the old parameters, and the number of new and old learnable parameter tokens are m and n, respectively. When expanding the model, it is only necessary to concatenate them into a size of (m+n)d. In the code , hidden_size is used to assign value to dimension d. However, in the 150M_eval.yml and the 450M_eval.yml, the size of hidden_size changes. This suggests that when expanding from a smaller model to a larger one, it would not be possible to concatenate parameters along the dimension of the number of learnable parameter tokens as mentioned in section 3.3. This seems inconsistent with what is written in the paper. Could it be that I am misunderstanding this? Could you please explain how exactly you implemented the model scaling?

Looking forward to your reply.

[Haiyang-W]

Thanks for your interest! Sorry for the late reply.
The answers are as follows,

1. d_token is fixed for long-context modeling, as mentioned in paper. Scaling the token dimension (d_token) and network depth are both unsustainable. For instance, scaling d_token leads to a significant increase in inference costs. The computational overhead of token-parameter interactions can be controlled and reduced by combining with MoE, but the computation costs of token-token interactions cannot be reduced once the token dimension increases. Network depth faces similar issues, as continually increasing depth is unsustainable; once it exceeds 100 layers, various optimization issues may arise. With TokenFormer, once combined with MoE, we can scale sustainably without increasing computational overhead. Of course, if necessary, we could increase token dimension and depth, though it may require higher costs to bring the model back to a satisfactory level. Our Figure 7 also shows some results of increasing the token dimension on Transformer. I believe there will be similar phenomena in TokenFormer.
2. 150M and 450M configs are standard settings, trained from scratch to match the performance of the Transformer, so the number of layers and channel dimensions are the same as Transformer. However, when scaling, our approach is to temporarily keep d_token fixed. If needed, we can later increase the token dimension and network depth. But as I said, increasing d_token and network depth is unsustainable, whereas increasing the token number and combining it with MoE is sustainable.

[Went-Liang]

Thank you very much for your response. I have another question: In Appendix "ZERO INITIALIZATION IN TOKENFORMER," equation (28) is missing the term \tau from equation (5). This coefficient \tau increases as the model scales up, which makes the statement "Since the newly added key parameters are initialized to zero, the attention mechanism does not modify the original output" seem invalid. This can also lead to very high initial loss when scaling up the model, although the loss will drop quickly afterward. I'm not sure if my understanding is correct, and I'm unsure if this is an important issue. Have you tried not using this coefficient? Looking forward to your reply.

[Haiyang-W]

Sorry for the late reply. I just underwent a surgery. Thanks for your interest. :) In our paper, we fixed \tau in the enwik8 experiment. I think it might be possible to design a better normalization method to ensure the loss doesn't fluctuate abruptly.