Summary
This paper proposes a new technique for training Transformer models while ensuring that the learned weights obey certain constraints, designed to facilitate mechanistic interpretability. In particular, the technique works by (1) freezing certain weight matrices throughout training (e.g., to always write the output of an attention layer to a previously unused linear subspace of the residual stream) and (2) introducing Gumbel-softmax sampling at various steps of the computation, and annealing the temperature down during training to ensure that at convergence, those computations perform discrete operations. The resulting Transformers can be automatically converted into symbolic programs that read, write, and perform lookup-table-style computation with a fixed number of categorical and integer variables. The authors train (quite small) Transformers in this way, to solve several tasks, including named entity recognition and sentiment classification on short sequences. The authors also provide some analysis of the programs extracted from these Transformers.
Strengths
This is a neat paper that advances two exciting lines of research:
- RASP and Tracr been helpful as intuition pumps for how trained Transformers *could* work, but it has not been clear whether the (correct) RASP algorithms researchers manually developed could be discovered automatically by gradient descent, and if not, what different properties these learned Transformer algorithms would have. This paper proposes a method for learning Transformer programs that could be used to start to answer these questions.
- A key challenge in mechanistic interpretability research (e.g. on Transformer circuits) is that Transformers trained by (unconstrained) gradient descent appear to collapse multiple conceptually distinct concepts via superposition. This paper proposes training methods that may help to mitigate superposition by forcing the network to read, write, and manipulate individual ‘variables’ that live in orthogonal subspaces of the residual stream.
Although terse in some places (perhaps due to length constraints), the methodology is for the most part clearly described, and the limitations of the approach seem sensible for a first attempt at this kind of constrained training. Another strength is that the framework seems to be extensible—the examples in the paper paint a clear enough picture of how to modularly extend the target symbolic programming language with new features (and then how to extend the neural net architecture and training procedure accordingly).
Weaknesses
- Although the motivation is interpretability, I was not convinced that the learned programs in the experiments were particularly interpretable. It would have been nice to see more detailed analysis of the programs. Outside the toy in-context learning example, are there any interpretable algorithms that are learned? Can examples on which the learned programs fail be traced back to particular ‘bugs’ in the symbolic programs? Can the symbolic programs be used to engineer adversarial examples on which the Transformer will fail? Do ordinary transformers trained without constraints also fail on these adversarial examples (indicating that they might be using similar flawed algorithms)? How stable are the learned programs across different training runs?
- It is unclear how scalable the presented technique is, to more complex problems, longer sequences, or bigger networks. I also would have appreciated more discussion / investigation of how easy or tricky the optimization was, and why it appears to work more poorly for larger Transformers (e.g., that handle larger sequence lengths). What are the key challenges in scaling up? Do the Gumbel gradient estimates become too noisy, e.g.? Or are there other main challenges?
Questions
- At the end of training, your Gumbel-Softmax's have been annealed to Categorical distributions, but these distributions are not necessarily *peaked* or low-entropy categoricals. Therefore, your extracted program (which takes the argmax of the categorical) may behave quite differently from your trained network (which randomly samples the categorical). Have you investigated (a) the extent to which the learned categoricals are peaked at a single mode, and (b) the difference in task performance / accuracy between the final learned Transformer and the extracted program?
- Why limit the predicate matrices to contain only one ‘1’ per row? Is it just because this constraint is more amenable to the Gumbel-softmax training procedure?
- In the description (Appendix A.1) of the categorical attention mechanism, you appear to use an argmax operation that is not differentiable. During training do you actually use Gumbel-softmax? If so, how are the resulting samples used to compute the output of the attention layer? (More generally, it would help if the paper more clearly laid out exactly where in the architecture all the Gumbel-softmax samples are used. Perhaps algorithm boxes in the style of https://arxiv.org/pdf/2207.09238.pdf could be provided for each component of the neural architectures you build?)
- For the in-context learning experiment, did you try (1) training an ordinary Transformer of the same size on the task? and/or (2) training a larger (i.e., overparameterized) Transformer program on the task? I’d be curious about the results of either experiment, if you did run them. Under the ‘lottery ticket hypothesis,’ overparameterization is an important part of how algorithms are learned by gradient descent — it would be interesting to see what this looks like in a Transformer program. E.g., in a larger model, do we end up with a small interpretable subprogram that correctly completes the task?
- Can you clarify how the word embeddings are trained? What does it mean that the 100-d GloVe embeddings are “randomly projected” to e.g. 32*4=128-d embeddings? Do you use a randomly generated dense 100x128 matrix? Then how are the embeddings trained? Is there a Gumbel sample for each of the categorical variables?
- Why not include MLP layers in the NLP tasks?
- For the RASP tasks that were not successfully solved by learned Transformer programs, do you believe this is because (1) good solutions exist, but are not found by gradient descent, or (2) your subset of RASP is too limited and the tasks cannot be solved with this subset? If (2), what are the prospects for handling a broader subset of RASP?
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Limitations
Limitations are adequately discussed, but I do think it could be useful to say a bit more about the challenge of scaling this approach -- what are the key obstacles & prospects for addressing them?