Comment by amluto

2 days ago

This design implicitly does something similar to something that I sometimes think conventional transformers should try: allowing later layers to query the KV data from earlier layers. As far as I can tell, with a conventional transformer, if a layer (and presumably higher-level-thinking) layer wants wants to take input from earlier tokens from something lower down, it needs to get it from the output and “remember” it by itself instead of just reading it directly.

But suppose an extra attention head were added that queried the KV data from lower layers. At the very least, I imagine this might cleanly solve the STRAWBERRY problem: whatever layer has figured out that the prompt wants to count instances of R could attend to lower layers that actually perceive those Rs.

This architecture does not allow later layers to directly query KV data from earlier layers. Each iteration of the loop uses the same layer parameters, so the KV data in later layers may well end up being the same, but only if the model stops changing it in response to other tokens in the context. Which is also something a traditional multi-layer transformer could do. (But might not end up doing due to lack of corresponding inductive bias.)

None of this helps with the strawberry problem, where the very first layer already gets a tokenized representation, so there is no layer that "actually perceives those Rs."

  • Is it fair to say that the “Rs in strawberry problem” will not be “cleanly” solved unless we advance beyond tokenization?

    • I don't see why that follows.

      The “Rs in strawberry problem” is literally "count the token R" in the word "strawberry".

      One could argue that the learnt tokenization model where it is tokenized into 3 tokens (see https://platform.openai.com/tokenizer) is problematic, but one solution to that is to double-down on it and learn tokenization as part of the end-to-end training instead of separately.

      If you mean that the idea of the current idea of the tokenization model being entirely fixed then I agree.

      (I'm not entirely sure how multi-modal models function in this regard - they must have a idea of the bytestream, but not familiar enough with that to comment intelligently.)

    • I think tokenization is probably not going anywhere, but higher layers need the ability to inspect 'raw' data on demand. You don't spell out most words as you read them, but you can bring the focus of your entire mind to the spelling of the word strawberry if you so choose. Models need that ability as well.

      1 reply →

> extra attention head were added that queried the KV data from lower layers

Isn't this sort of similar to latent looping? E.g. [1]. But actually as [2] argues, even that wasn't a good experiment because it used the very last hidden state, which is too close to the logits and loses most of the rich embedding structure. Perhaps you don't even need access to the state of anything except the penultimate hidden layer, since based on my vague reading of [3] the residual stream doesn't "lose information" as it passes deeper down the attention layers, so each block maybe manipulates a different subspace of the residual stream.

[1] https://news.ycombinator.com/item?id=45758093

  • > Perhaps you don't even need access to the state of anything except the penultimate hidden layer, since based on my vague reading of [3] the residual stream doesn't "lose information" as it passes deeper down the attention layers, so each block maybe manipulates a different subspace of the residual stream.

    I imagine that conventional transformers kind of force this. If you train a transformer such that it needs to learn the ability to do tasks like “Repeat the following words: apple banana cat” then the model is sort of forced to internally propagate the input far enough along to be able to perform the task. But maybe if you pre-trained from scratch with an architecture where later layers get direct access to earlier layers and/or the raw input, then the model wouldn’t need to propagate information.

    Or maybe it would all fall apart and something would go wrong with the gradients.

I remember doing this kind of test in a vanilla transformer trained on my laptop on a small text dataset. I basically added N^3 attention where each layer could pay attention to previous layers. It didn't improve anything and was much slower.

Hard to say whether something scales or not from a couple dozen million parameters to an actual billion-sized model, but I have the impression that the nature of the residual stream and its high dimensionality allows any layer to access information of previous layers if the transformers needs it.

I can't instinctively process how many R's are in STRAWBERRY. I use my vision to get it though almost immediately.

I feel simple transformers simply don't get access to those modalities that a human would use. I can't use my "talking" centers to count letters in words either.

You just need to pay attention to understand you don't use your language skills to count words.