Open PicoCreator opened 1 year ago
I already tried similar methods (RWKV-4b) :-)
Please check https://discord.com/channels/992359628979568762/1083107245971226685/1085912612035903549 (in RWKV Discord)
Ahh you mean the "However some tiny amt of QKV attention (as in RWKV-4b)" part of the message. Glad to see my understanding / theory / some validation in this direction all in one post.
If i understood right RWKV-4b is v4neo, with RWKV_MY_TESTING
enabled wtih def jit_funcQKV(self, x):
at https://github.com/BlinkDL/RWKV-LM/blob/6fea06e671ecf07d037caf3b8bdf415ddd7f3984/RWKV-v4neo/src/model.py#L221
And is not part of the current raven model?
Haha what would happen if there is even more attentions at the upper layer, as i see the jit_funcQKV is quite linear across the layers (unless i read the code wrongly)
Motivation/Theory
I suspect that because the upper block layers, can only "see" the block layer before it. There is lots of redundant information being forwarded between the layers.
Current flow
The following is our current data flow (somewhat)
So for example, if we have input like
The following would probably be approximate "level of thought on the layers"
Layer 0: story Layer 1-5: What is the species of the pet? Layer 6-10: (parts of the story), what the pet species?, maybe its a goat? Layer 20-30: (parts of the story), what the pet species?, maybe its a goat? probably not, probably a cat
While we do not accurately understand the exact information stored in latent space embedding, for an information theory standpoint, because the "upper layers" will require information from the "lower layers", there will be multiple redundant information flowing through our limited embedding through each layer. And while we probably have not hit the limit yet (at 8k tokens) of compressing information to our current embedding size, we would at some point.
Suggested change
The idea here, is to add mini attention layer from block 2 onwards, where it would read the time mixing / channel mixing embedding for all the previous layers + the input embedding.
The size of this layer, would grow X^2 based on the number layers. But the computation cost will remain as N per token
This allow the later layer, to simply depend on the lower layer for certain key latent information, without needing to "retransmit them to the upper layer"
So something like the following
So that instead it can be more like the following
Layer 0: story Layer 1-5: What is the species of the pet? Layer 6-10: (parts of the story) Layer 20-30: maybe its a goat? probably not, probably a cat
Allowing the various layers to be more focus onto a certain set of information which can be read by the upper layers
I also split up the time mixing, and channel mixing path, to allow each of it to be computed in parallel, to hopefully somewhat improve GPU utalization
Not made clear in the digram - is that the output for the attention layer should be (2, embedding size). Which means that that it output separate embedding for time and channel mixing respectively
So my question to @BlinkDL - does the above make sense? is it something you have tried and found that it made things worse?
Or is my basis and theory behind it flawed, and this seems like needless complications