Grokked Transformers Are Implicit Reasoners

Since I first learned about grokking, I&#x27;ve had a strong suspicion that getting a handle on it and figuring out how to aid it should be the central question in AI. We are currently stuck in a local minimum, where memorizing circuits perform well-enough to handle a whole lot of economically viable use cases. But the profit function has guided us into a valley dominated by a data and compute hungry architecture that isn&#x27;t ideal for learning generalizing circuits (partially because the memorizing circuits are so effective! We relatively quickly get to a flat loss landscape, after which we blindly jump around for countless epochs in a kind of Brownian motion until we get into an area where regularizers can drive generalization). Research like this paper is incredibly important.<p>I thought this was the most interesting bit from the paper:<p>&gt; Training data distribution, instead of training data size, qualitatively influences generalization behavior.

Reminds me of that old quote about “the difference between average and state of the art is forgetting to turn it off over summer break” or similar.<p>I wonder if this is why smaller LLMs seem to punch above their weight, are they further along in the process of distilling the data down into understanding?

Conceptually, grokking reminds me of the concept presented in the book The Dip<p>Most people, for most tasks, will only learn&#x2F;train&#x2F;try to improve, up to where they get to a flat or negative return curve per unit of effort put in<p>But, the people that are the best at a certain task, usually implies they got through The Dip in the curve of return per effort

I just learned about grokking; reminds me of double descent, and I looked up a 2022 paper called &quot;Unifying grokking and double descent&quot;. I&#x27;m still unclear on what the difference is. My basic understanding of double descent was that the regularization loss made the model focus on regularization after fitting the train data.

This paper feels way too abstract, to the point it makes it hard to understand what the team actually did.<p>For instance, the paper claims it beat GPT-4-Turbo and Gemini-Pro-1.5 on certain tasks... but it doesn&#x27;t include any of the questions they asked GPT4 or Gemini, so it&#x27;s hard to guess whether these results have any value at all.<p>It&#x27;s also unclear what they even trained their custom transformer to do. It has a custom tokenizer, but they don&#x27;t give a list of tokens (aside from a few examples in the diagrams like &quot;Barrack&quot;, &quot;Michelle&quot;, &quot;Trump&quot;). They talk about in-distribution and out-of-distribution tasks, but they don&#x27;t give any examples of these tasks and what they look like.<p>This feels like accidental complexity. It wouldn&#x27;t have been hard to add a few more appendices with eg a list of 20 or so in-distribution sentences they asked the model to complete and 10 out-of-distribution sentences. Instead all they include is diagrams comparing performance for different hyperparameters and stuff, but we don&#x27;t even know <i>what the models are being tested on</i>.

I sort of wish that we would move on from the &quot;grokking&quot; terminology in the way that the field generally uses it (a magical kind of generalization that may-or-may-not-suddenly-happen if you train for a really long time).<p>I generally regard grokking as a failure mode in a lot of cases -- it&#x27;s oftentimes not really a good thing. It tends to indicate that the combination of your network, task, and data are poorly suited for learning {XYZ} thing. There are emergent traits which I think the network can learn in a healthy manner over training, and I think that tends to fall under the &#x27;generalization&#x27; umbrella.<p>Though I&#x27;d strongly prefer to call it &#x27;transitive&#x27; rather than &#x27;compositional&#x27; in terms of generalization, as transitive is the formal term most disciplines use for such things, compositional is a different, more general meaning entirely. Similarly, I&#x27;d replace &#x27;parametric&#x27; and &#x27;non-parametric&#x27; with &#x27;internal&#x27; and &#x27;external&#x27;, etc. Sloughing through the definition salad of words (this paper alone takes up roughly half of the top Kagi hits for &#x27;parametric memory&#x27;) makes actually interpreting an argument more difficult.<p>One reinterpretation of the problem is -- of course external memory models will have trouble generalizing to certain things like models relying on internal memory do! This is because, in part, models with internal memory will have much more &#x27;experience&#x27; integrating the examples that they&#x27;ve seen, whereas, for an external-memory model like a typical RAG setup, anything is possible.<p>But, that being said, I don&#x27;t think you can necessarily isolate that to the type of memory that the model has alone, i.e., I don&#x27;t think you can clearly say even in a direct comparison between the two motifs that it&#x27;s the kind of memory itself (internal vs. external) that is to blame for this. I think that might end up leading down some unfruitful research paths if so.<p>That said, one positive about this paper is the fact that they seem to have found a general circuit that forms for their task, and analyze that, I believe that has value, but (and I know I tend to be harsh on papers generally) the rest of the paper seems to be more of a distraction.<p>Definitional salad buffets and speculation about the &#x27;in&#x27; topics are going to be the things that make the headlines, but in order to make real progress, focusing on the fundamentals is really what&#x27;s necessary here, I think. They may seem &#x27;boring&#x27; a lot of the times, but they&#x27;ve certainly helped me quite a bit in my research. &lt;3 :&#x27;))))

One of the biggest bottlenecks of multi layer transformers is that reasoning can only happen in the hidden layers. Past the final layer, the model must generate a token that conforms to the training process. This token can then be fed back into the transformer from the beginning, but since it necessarily must be in natural language, it limits the type of reasoning the model can perform to the &quot;thoughts&quot; it has seen in the dataset and is therefore allowed to express. If you could figure out how to have the first layer take both the KV of the first layer and the KV of the final layers in the attention mechanism into account, the model would become capable of infinite length reasoning.

How do we actually implement this? Imma struggling to work out how I could use this rather than my stupid langraph, recursive rag checking crap that takes too much time and never really does justice.

Grokking is all you need?

This is the interesting result where their GPT-2 sized transformer blows away GPT4 and Gemini 1.5 in connecting together facts<p><pre><code> The difficulty of such a task is two-fold. First, the search space is large. For example, on average, each query entity connects with more than 50 facts, and each bridge entity in the ground truth proof connects with more than 900 facts. Second, there are no surface form clues to exploit and bias the search towards the ground truth proof, unlike most conventional QA benchmarks where the proof steps are transparent from the query. To test LLMs based on non-parametric memory, we translate the facts into natural language by simple templates (Appendix F). Facts&#x2F;queries for each attribute are grouped&#x2F;tested separately. We test both the vanilla setup where all facts (28.2K on average) are loaded into the LLM context, and the retrieval-augmented setup (5.4K facts retrieved on average) where the two-hop neighborhoods of the two query entities are retrieved, which includes enough facts to deduce the answer. We also try both standard prompting where the model answers directly, and chain-of-thought (CoT) prompting where the model is prompted to verbalize the reasoning. We test GPT-4-Turbo and Gemini-Pro-1.5, where for GPT-4-Turbo we only test the retrieval-augmented setup due to context length limit. Table 1:Results on the complex reasoning task. Direct&#x2F;CoT: predict the answer directly&#x2F;verbalize the reasoning steps. “+R”: retrieval augmentation. GPT-4-Turbo Gemini-Pro-1.5 Grokked Transformer Direct+R CoT+R Direct CoT Direct+R. CoT+R Accuracy (%) 33.3 31.3 28.7 11.3 37.3 12.0 99.3</code></pre>

&gt; delve