Added the paper pdf.

Author: golkar

_posts/2024-05-30-counting.md

@@ -1,23 +1,23 @@
 ---
 layout: post
-title: "Contextual Counting: A Mechanistic Study of Transformers on a Quantitative Task"
+title: "How Do Transformers Count in Context?"
 authors: Siavash Golkar, Alberto Bietti, Mariel Pettee, Michael Eickenberg, Miles Cranmer, Keiya Hirashima, Geraud Krawezik, Nicholas Lourie, Michael McCabe, Rudy Morel, Ruben Ohana, Liam Holden Parker, Bruno Régaldo-Saint Blancard, Kyunghyun Cho, Shirley Ho
-shorttitle: "Contextual Counting: A Mechanistic Study"
+shorttitle: "Counting in Context?"
 date: 2024-05-30 09:23
 image: counting-splash.jpg
 smallimage: counting-s.jpg
 blurb: We introduce the Contextual Counting task, a new toy problem aimed at exploring interpretability of Transformer models in quantitative domains. We compare the performance of causal and non-causal models with different position codes and find causal models with RoPE and NoPE significantly outperform other configurations. We provide detailed explanation of how the circuits function and what makes them succeed or fail in generalization to out-of-distribution samples.
 shortblurb: We introduce the Contextual Counting task, a new toy problem aimed at exploring interpretability of Transformer models in quantitative domains.
 splashimage: /images/blog/counting-splash.jpg
-link: 
-github_link: 
+link: https://arxiv.org/pdf/2406.02585
+github_link: https://github.com/golkar/counting/blob/main/
 permalink: /blog/contextual_counting/
 ---
 
-One of our goals at Polymathic-AI is to utilize the recent advances in machine learning to implement and deploy state of the art models that can aid in scientific exploration and discovery. As such, we believe that it is very important for us to be able to understand how these models behave and interpret the algorithms that are learned in these networks. The benefits of this direction of research are two-fold. First, by understanding the inner workings of the model, we get better insight into the scientific domain of interest. Second, by understanding the strengths and weaknesses of these architectures, we can design networks that are better suited for their task.
 
-In this blog post, we summarize a recent paper which is part of an ongoing effort in our team in this direction. In this work, we introduced a new toy problem specifically designed to advance the interpretability of Transformer models in quantitative and scientific contexts. This task, called **contextual counting**, requires the model to identify a specific region of interest within a dataset and perform accurate counting. As such, it simulates scenarios where precise localization and subsequent computation are critical, such as in object detection or region-based analysis in scientific data. 
+At Polymathic-AI, part of our mission is to develop foundation models that help with scientific exploration and discovery. But it's not enough to build these models, we also want to understand them! What algorithms do these networks actually learn under the hood? By uncovering them, we might discover improvements to our foundation models or even new insights about the scientific domains they represent.
 
+To understand how Transformers solve complex problems, it helps to start with simpler tasks. In [a recent paper](https://arxiv.org/pdf/2406.02585), we do exactly this. We introduce a new toy problem designed to help us understand how Transformers can count in a context-dependent way---a core capability for scientific and quantitative reasoning. We call this task *contextual counting*. Contextual counting asks the model to count tokens in different regions of a sequence. As such, it idealizes scenarios where precise localization and subsequent computation are critical, such as counting specific neuro-receptors within a neuron in biological research. While seemingly simple, this task is surprisingly hard for state-of-the-art LLMs.
 <br>
 ## Introducing the Contextual Counting Task
 <br>
@@ -34,17 +34,6 @@ To tackle this task using a Transformer architecture, we use an encoder-decoder
 </p>
 For our experiments, we fix the number of regions to 4 and the sequence length to 512. This allows us to explore how solutions generalize to different numbers of regions and sequence lengths.
 
-#### Relevance
-
-The contextual counting task is not just an instructive example for understanding Transformers; it emulates real-world quantitative problems requiring precise sensitivity to regional boundaries. Examples include counting specific neuro-receptors within a neuron in biological research.
-
-Current state-of-the-art LLMs struggle with this task, indicating the need for specialized models and interpretability techniques tailored to quantitative and scientific applications.
-
-#### Why Study Toy Problems?
-
-Toy problems serve as simplified models that help us understand complex systems. By stripping down the intricacies of real-world scenarios, toy problems allow researchers to isolate and examine specific mechanisms within machine learning models. This focused approach is particularly valuable in the study of interpretability, where the goal is to unravel how models make decisions. Simplified tasks such as contextual counting provide a clear, controlled environment where researchers can systematically manipulate variables and observe the effects. This clarity is often lost in more complex, real-world problems, where numerous interacting factors can obscure the underlying processes. By starting with toy problems, we gain foundational insights that can later be applied to more complicated tasks.
-
-Moreover, toy problems are instrumental in benchmarking and testing new theories and methods. They act as proving grounds for hypotheses about model behavior and performance. For instance, by using toy problems, researchers can quickly iterate on models and interpretability techniques, refining their approaches before deploying them on more sophisticated and critical tasks. This iterative process accelerates the development of robust methods that can be confidently applied in high-stakes domains like healthcare, finance, and scientific research. In the context of Transformers, toy problems help uncover how different architectures and encoding methods influence model performance and interpretability, providing essential knowledge for advancing machine learning technologies.
 
 <br>
 ## Theoretical Insights
@@ -88,13 +77,13 @@ The theoretical results above imply that exact solutions exist but do not clarif
 
 We summarize the results of this empirical exploration below.
 
-#### 1. Causal transformers significantly outperform non-causal ones.
+#### 1. Causal Transformers significantly outperform non-causal ones.
 
 <p align="center">
-  <img src="/images/blog/counting/accuracy.png" alt="Performance of the different configuraiton" width="65%" style="mix-blend-mode: darken;">
+  <img src="/images/blog/counting/accuracy.png" alt="Performance of the different configuration" width="65%" style="mix-blend-mode: darken;">
 </p>
 
-The above figure shows the performance of different Transformer configurations. The most prominant feature of this figure is that non-causal transformers with any positional encoding fail to get good performance. In contrast, causal Transformers can achieve close to 100\% accuracy.
+The above figure shows the performance of different Transformer configurations. The most prominent feature of this figure is that non-causal Transformers with any positional encoding fail to get good performance. In contrast, causal Transformers can achieve close to 100\% accuracy.
 
 #### 2. NoPE is best but harder to train than RoPE.
 
@@ -131,21 +120,19 @@ The figure below shows the behavior of three different type of solutions when ge
 
 We can get a hint at what might be the culprit by looking at the attention pattern of the decoder. The attention pattern given in the previous point pertains to the blue dots on this figure, i.e. the model that generalizes best. 
 
-The figure below, shows the attention pattern of the orange dots, i.e. the model that generalizes do different seuqence lengths but not to different region numbers. We see that as before, the decoder pays attention to the 1-tokens of the relevant region (in this case the first region), however this time the role of the bias term is played by the ]-tokens. During training, the number of regions is fixed at 4, and therefore the number of ]-tokens can be used as a constant bias. However, this is not the case when the number of regions changes. This explains why this model does not generalize to other number of regions.
+The figure below, shows the attention pattern of the orange dots, i.e. the model that generalizes do different sequence lengths but not to different region numbers. We see that as before, the decoder pays attention to the 1-tokens of the relevant region (in this case the first region), however this time the role of the bias term is played by the ]-tokens. During training, the number of regions is fixed at 4, and therefore the number of ]-tokens can be used as a constant bias. However, this is not the case when the number of regions changes. This explains why this model does not generalize to other number of regions.
 
 <p align="center">
   <img src="/images/blog/counting/decoder_nongen.png" alt="The attention profile of the decoder of a non-generalizing model." width="55%" style="mix-blend-mode: darken;">
 </p>
 
 In our exploration, we found that the model can use any combination of quantities that are constant during training as biases. 
 
-#### 7. (Technical) The network generates its output by balancing two learned shapes.
-
-This point is a little technical and it pertains to the detail of how the network explicitly generates its output. I think it is cute enough to be worth mentioning. 
+#### 7. The network generates its output by balancing two learned shapes.
 
 In some of our experiments, we chose to remove the MLP and self-attention layers from the decoder block. That is, the decoder is just a cross-attention layer. This configuration is less expressive but has the advantage that the output of the model is a linear combination of the value vectors derived from the embeddings of the encoder.
 
-In a preious case we saw that the decoder only attended to the 1-tokens of the relevant region and the beginning-of-sequence token. The figure below shows the value vectors of these two tokens.
+In a previous case we saw that the decoder only attended to the 1-tokens of the relevant region and the beginning-of-sequence token. The figure below shows the value vectors of these two tokens.
 
 <p align="center">
   <img src="/images/blog/counting/values.png" alt="The value vectors." width="55%" style="mix-blend-mode: darken;">
@@ -172,9 +159,9 @@ We see that as n becomes large, the difference between n and n+1 becomes smaller
 <br>
 The contextual counting task provides a valuable framework for exploring the interpretability of Transformers in scientific and quantitative contexts. Our experiments show that causal Transformers with NoPE can effectively solve this task, while non-causal models struggle. These findings highlight the importance of task-specific interpretability challenges and the potential for developing more robust and generalizable models for scientific applications.
 
-For more details, check out our preprint on the [arXiv](link).
+For more details, check out the [paper](https://arxiv.org/pdf/2406.02585).
 
 *-- Siavash Golkar*
 
 
-Image by [Tim Mossholder](https://unsplash.com/photos/blue-and-black-electric-wires-FwzhysPCQZc) via Unsplash.
+Image by [Tim Mossholder](https://unsplash.com/photos/blue-and-black-electric-wires-FwzhysPCQZc) via Unsplash.