TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters
1–10 of 34 posts
Re: TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters
#2The authors factorize every weight matrix with an attention mechanism:
weight = attention(token_query, weight_keys, weight_values).
Increasing model size becomes a matter of adding more weight_keys and weight_values, and incrementally training them.
Simple, clever, and it seems to work well. Beautiful.
Re: TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters
#3The authors factorize every weight matrix with an attention mechanism: weight = attention(token_query, weight_keys, weight_values). In other words, they query weight_keys to fetch the weight_values, and mix them to compute each weight on the spot. Increasing model size becomes a matter of adding more weight_keys and weight_values, and incrementally training them. Simple, clever, and it seems to work well. Beautiful.
There is a particularly nice geometric interpretation of attention I just realised recently in a flash of enlightenment, best explained with an interactive Desmos plot (black dot is draggable):
https://www.desmos.com/calculator/3rtqsyapxo
The above assumes the columns of K are normalised but bear with me. K and V together form a vector database. V are the payloads, each row containing a vector of data. K describes the position of these points in space, on the surface of a hypershpere. The query vector describes the query into the database: the vector direction describes the point in space that's being queried, the vector magnitude describes the radius of the query. The result is the weighted average of vectors from V, weighted by their distance from the query vector scaled by the query radius (which has a smooth Gaussian falloff). A recent paper from Nvidia I recommend, which derives a significant speedup by normalising vectors to a hypershpere: https://arxiv.org/abs/2410.01131v1
Re: TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters
#4The authors factorize every weight matrix with an attention mechanism: weight = attention(token_query, weight_keys, weight_values). In other words, they query weight_keys to fetch the weight_values, and mix them to compute each weight on the spot. Increasing model size becomes a matter of adding more weight_keys and weight_values, and incrementally training them. Simple, clever, and it seems to work well. Beautiful.
I believe there have been studies showing that the attention mechanism allows estimation of gradients for one-shot learning (i.e, based on what you tell the model you want in the input, it will use attention to 'update' the weights of the linear layers to 'learn' new information). This seems to be taking that one step further and just using attention for the weight estimations itself. The key insight here is that by adding more tokens to the weight estimation calculation, you can get more degrees of freedom.
Total aside, but imagining how many levels of functions are present in the calculation of each activation here, and thinking about how regular old differentiation and gradient descent actually work to train these nested parameters, is truly amazing, in my opinion.
Re: TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters
#5The authors factorize every weight matrix with an attention mechanism: weight = attention(token_query, weight_keys, weight_values). In other words, they query weight_keys to fetch the weight_values, and mix them to compute each weight on the spot. Increasing model size becomes a matter of adding more weight_keys and weight_values, and incrementally training them. Simple, clever, and it seems to work well. Beautiful.
I believe there have been studies showing that the attention mechanism allows estimation of gradients for one-shot learning (i.e, based on what you tell the model you want in the input, it will use attention to 'update' the weights of the linear layers to 'learn' new information). This seems to be taking that one step further and just using attention for the weight estimations itself. The key insight here is that by…
Yeah. This thing is "assembling a different transformer" on the spot for each token.
If one thinks about it for more than a moment, it's kind of incredible that it works.
Re: TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters
#6The authors factorize every weight matrix with an attention mechanism: weight = attention(token_query, weight_keys, weight_values). In other words, they query weight_keys to fetch the weight_values, and mix them to compute each weight on the spot. Increasing model size becomes a matter of adding more weight_keys and weight_values, and incrementally training them. Simple, clever, and it seems to work well. Beautiful.
There is a particularly nice geometric interpretation of attention I just realised recently in a flash of enlightenment, best explained with an interactive Desmos plot (black dot is draggable): https://www.desmos.com/calculator/3rtqsyapxo The above assumes the columns of K are normalised but bear with me. K and V together form a vector database. V are the payloads, each row containing a vector of data. K describes th…
Yeah, I believe this intuition first introduced by the Neural Turing Machine line-of-work and later simplified into AIAYN paper (NTM maintains "external memory" a.k.a. weight_keys, weight_values here).
Disclaimer: these are from my memory, which can be wrong entirely.
Re: TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters
#7Seems like a big deal. I feel like this could enable a new level of modularity and compatibility between publicly available weight sets, assuming they use similar channel dimensions. Maybe it also provides a nice formalism for thinking about fine tuning, where you could adopt certain heuristics for adding/removing key-value pairs from the Pattention layers.
One interesting thing to note: sounds like model scaling happens on the fly by adding key-value pairs as rows in the K and V matrices on the Pattention layer. That suggests that weights represented by tokens in the first rows may be more important than weights in later rows. There may be a lot you could do with that ordering of weights in terms of pruning and such.
Re: TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters
#8Surprised not to see a comparison to https://paperswithcode.com/paper/augmenting-self-attention-w...
Re: TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters
#9Earlier quoted context omitted.
I believe there have been studies showing that the attention mechanism allows estimation of gradients for one-shot learning (i.e, based on what you tell the model you want in the input, it will use attention to 'update' the weights of the linear layers to 'learn' new information). This seems to be taking that one step further and just using attention for the weight estimations itself. The key insight here is that by…
Yeah. This thing is "assembling a different transformer" on the spot for each token. If one thinks about it for more than a moment, it's kind of incredible that it works.
I think the same about regular neutral networks
Re: TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters
#10This could be revolutionary. The PPL/compute graphs are damning. If the Transformer is a function, then the TokenFormer feels like a higher-order function. Perhaps this approach is a natural direction for producing System Two reasoning? There's so much to digest here...