Adaptive Transformer Programs: Bridging the Gap Between Performance and Interpretability in Transformers

Authors: Quoc-Vinh Lai-Dang, Taemin Kang, Seungah Son
Primary area: interpretability and explainable AI
Keywords: Mechanistic Interpretability, Transformers, Interpretable AI

Abstract:
Balancing high performance with interpretability in increasingly powerful Transformer-based models remains a challenge. While mechanistic interpretability aims to specify neural network computations in explicit, pseudocode-like formats, existing methods often involve laborious manual analysis or struggle to fully elucidate learned internal algorithms. Recent efforts to build intrinsically interpretable models have introduced considerable expressivity and optimization challenges. This work introduces Adaptive Transformer Programs, an enhanced framework building upon RASP language and Transformer Programs to create more robust and interpretable models. The proposed method increases expressivity by redesigning two primary attention modules to improve categorical and numerical reasoning capabilities. To overcome optimization hurdles, we introduce a novel reparameterization scheme that enhances the exploration-exploitation trade-off during training. We validate our approach through extensive experiments on diverse tasks, including in-context learning, algorithmic problems (e.g., sorting and Dyck languages), and NLP benchmarks such as named entity recognition and text classification. Results demonstrate that Adaptive Transformer Programs substantially narrow the performance gap between black-box Transformers and interpretable models, enhancing transparency. This work advances the development of high-performing, transparent AI systems for critical applications, addressing crucial ethical concerns in AI development.

Decision: accept (poster)
Reviewer scores: 8, 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1730642145750

thread_type:
official_review

<thread_object>
creator: reviewer_j1b6

<content>
summary:
The paper proposes Adaptive Transformer Programs, which are a method to increase the expressive capability of transformer programs (interpretable transformers which can be converted into succinct code). 

The main architectural modification proposed by the paper are:
1. Providing a method to smoothly switch between Gumbel-softmax and sparsemax as a method to control switching between exploration and exploitation.
2. A method to weight between categorical attention and score-based attention
3. A position-aware attention which allows numerical Value, while keeping Q, K categorical.

The paper includes experimental results on synthetic tasks on formal languages and natural language tasks. An ablation study is provided to determine the contribution of each architectural modification to the performance of the model.

strengths:
- The paper considers an important problem: that of making intrinsically interpretable models more expressive.
- The architectural modifications are well-motivated and largely well described in the paper.
- Experimental results across a wide range of tasks provide evidence for the performance of the model.
- A thorough ablation study is provided.
- The use of lines of code generated for the transformer program as a measure of interpretability is interesting.

weaknesses:
- There are sections of the paper which could do with more details. Section 3.3 in particular could provide more details about the reference attention e.g.

- Given that a stated goal of the proposed architecture is to make intrinsically interpretable architectures more performant, I would have liked to see experimental results which supported the claim of adaptive transformer programs narrowing the gap between standard transformers and transformer programs. Results are shown where the proposed architecture outperforms adaptive transformers (Table 1), but often adaptive transformers themselves already outperformed standard transformers (Tables 3, 4). I believe the paper would be greatly improved by the inclusion of results on a task in which a performance gap between standard transformers and transformer programs was narrowed by the proposed arch.

- In line 377 and the following line, it is said that removing all enhancements results in the lowest accuracy (76.44). This does not appear true however, as keeping only the smooth transition mechanism yielded lower accuracy (73.61), as shown in table 2.

questions:
- Could more details be provided about the hyperparameters and optimizers used for the various experiments?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730675338321

thread_type:
official_review

<thread_object>
creator: reviewer_ewyt

<content>
summary:
This paper introduces the Adaptive Transformer framework, designed to develop robust, expressive, and interpretable Transformer models that can be translated into human-readable programs. The authors address challenges in finding optimal solutions and promoting sparsity, limitations seen in the original framework's use of Gumbel-Softmax reparameterization. To overcome these, they propose a Smooth Transition Mechanism for discrete optimization, alongside Uncertainty-Aware Categorical Attention and Position-Aware Numerical Attention. Experiments on both synthetic and real-world NLP tasks demonstrate that this approach achieves superior performance over benchmarks and provides concise, insightful interpretability analysis.

strengths:
This paper is well-organized, presenting the concept of Adaptive Transformer clearly, from problem setting through to challenges and proposed solutions. Building upon the Transformer Programs framework, the authors introduce novel designs to address its limitations: the Smooth Transition Mechanism addresses the local optima issue in Gumbel-Softmax, Uncertainty-Aware Attention improves handling of categorical and contextual information, and Position-Aware Attention captures nuanced positional information and enables processing of non-integer values. Together, these enhancements expand the model’s numerical capabilities, supporting tasks that require fine-grained representations.

These methods are well-motivated and shown to be effective through ablation studies. Experiments on both synthetic and real-world NLP tasks demonstrate impressive performance compared to current baselines. The proposed method shows great potential for interpretable and reliable AI systems and is broadly applicable across transformer-based models, suggesting substantial future impact.

weaknesses:
While the overall performance is promising, the paper lacks case analysis for failure instances in the benchmarks, and the evaluated NLP tasks are limited to classification and NER. Including generative tasks would provide a more comprehensive understanding of the method’s impact.

**Minor issues:**

1.  In line 96, the MHA computation omits the normalization factor $\frac{1}{\sqrt{n}}$, where $n$ is the hidden dimension.
2.  In line 99, the double quotes before “Transformer” need correction.
3.  In line 109, please insert a space between "RASP" and "Weiss et al."

questions:
1.  In Figure 1's RASP accuracy comparison, why does the Adaptive Transformer Program outperform Standard Transformers on the Sort task but not on the Most Frequent task?
2.  In Table 2, accuracy appears to decrease after integrating the Smooth Transition Mechanism. How should this be interpreted?    
3.  How does the proposed method perform on generative tasks beyond NER and text classification?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730696152436

thread_type:
official_review

<thread_object>
creator: reviewer_ezgs

<content>
summary:
This paper introduces Adaptive Transformer Programs, an enhanced framework that improves upon existing Transformer Programs to create more interpretable AI models.

Specifically, the authors introduce three innovations: (1) a Transition Mechanism that combines Gumbel-Softmax and Sparsemax. (2) an dynamic attention mechanism that combined categorical and score-based attention using Jensen-Shannon Divergence. (3) a integrated attention mechanism that leveraging both Learnable and Sinusoidal positional encoding.

Then the authors conduct the experiement on various benchmark against the original Transformer Programs and achieve clear improvement. Meanwhile, the ablation study shows the effectiveness of the proposed components.

strengths:
- The proposed method removes the limitations of previous Transformers Programs from different perspectives, which can potentially bring huge benefits for improving the community's understanding of transformer models.
- The experimental results are solid and the proposed method clearly surpasses the Transformers Programs baseline in both accuracy and "lines of code."
- The experimental setting is comprehensive, consisting of both RASP and NLP tasks.
- This paper is well-written and easy to follow.

weaknesses:
The main concern to me is that the novelty is marginal. This paper proposes no new method, and all three methodologies are the integration of existing methods in this area. It either proposes a transition or direct combination of two existing methods. To me, these direct "transitions" are not interesting enough and are not fully supported by a simple ablation study on one benchmark. And at least one interpretation of "SMOOTH TRANSITION" is necessary to show the readers when the proposed transition takes effect and why it can be helpful.

questions:
(1) Why can reducing the "lines of code" be directly considered as improving interpretability?  
To me, the connection between "reduced program length" and better interpretability is not obvious enough. It would be helpful to provide a representation comparison example of the proposed model and original Transformer Programs to clarify that it can learn better representations.  

(2) Regarding sparsity, how does the proposed method behave better than Gumbel-Softmax?  
In section 3.2, one of the motivations is that Gumbel-Softmax fails to encourage sparsity, but I see no discussion following this point.  For example, some direct comparisons (like sparsity rate) between these two methods would be expected.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730754484261

thread_type:
official_review

<thread_object>
creator: reviewer_u9xh

<content>
summary:
This paper builds off previous work, "Learning Transformer Programs", which presents a Transformer variant which uses discretization and other simplifications so that the learned network can be represented as programs. The present work adds three modifications to the original architecture to improve representational ability and learning optimization. 

1. they replace Gumbel-Softmax with a smooth transition between Gumbel-Softmax and Sparsemax, a sparse version of softmax. 
2. They expand the representational capacity of the program-compatible attention by giving the model the option to use score-based attention in lieu of the representationally limited but interpretable and program-compatible categorical attention. To give the model the option, there is an uncertainty-aware gating between categorial attention and score attention which is based on the Jensen-Shannon Divergence. 
3. They add position embeddings.

The authors show that this improved model performs better or equal to the previous model, while having increased interpretability (fewer lines of code in the learned model).

strengths:
Originality: I looked on Google Scholar for work that cites the original Learning Transformer Programs that this work builds upon, and did not see anything. So, it seems like an original contribution.

Quality and Significance: The authors back up their proposed modifications with a solid array of experiments: in context learning, RASP tasks, and two NLP tasks. They show a scalability result and conduct an ablation study on one RASP task. The main way in which the new model improves over the previous model is learning programs with fewer lines of code that are slightly more accurate. However, the new model also scales slightly better (performance degrades less as you increase vocab size and max sequence length). 

Clarity: the paper is well written and presented. In particular, the authors do a good job explaining their Gumbel-Sparsemax and uncertainty aware attention contributions. There are some typos that should be fixed. I have some suggestions to improve the presentation, but nothing that would detract from my overall score of the paper.

weaknesses:
I find the main weakness to be that the evaluation is focused on accuracy and number of lines of code. Without seeing the programs, it is hard to understand what "lines of code" means for the simplicity of the program. See the questions.

Another main weakness is that there is no appendix giving further details. 

I think the paper would be much stronger with a section that looks at an example program found, and compares it to the program found by the original method, to understand the improvements of the modifications presented.

questions:
1. What is the intuition for how Gumbel-Sparsemax improves the optimization process? Is the main benefit that the final programs learned are simpler (fewer lines?) How does transitioning between Gumbel-Softmax and Sparsemax differ from what Gumbel-Softmax already does when the temperature is annealed over time for Gumbel-Softmax? One big difference seems to be that Sparsemax is deterministic? If i understand correctly, temperature annealing changes how discrete the Gumbel-Softmax is, but does not affect the randomness. even when fully discrete, one's sample might change each time. With sparsemax, the choice is deterministic (basically choose the most likely option?) — some explanation here would be helpful. 

2. Is there any mathematical relation between pure Gumbel-Softmax and pure sparsemax? It seems like they might not "mix" perfectly well, and there could be a more principled version of one or the other such that you can smoothly interpolate between the two and get both, but not have to just mash up the formulas with equation (8) (like imagine a second temperature parameter that controls randomness in the gumbel-softmax calculation)

3. What is A_ref? it is "introduced" but I don't see how it is calculated during training.
4.  "In low uncertainty (low JSD), CatAttention is preferred, as it is more confident in its categorical decisions" — I'm confused why this would be true. can't ScoreAttention regular attention? shouldn't it's performance be >= CategorialAttention's performance? so if the only objective is accuacy, I don't see why CategoricalAttention would ever be used.

5. "Our preliminary experiments with separate binary and continuous functions revealed scenario-dependent performance variations, motivating the development of a hybrid mechanism to create more Adaptive Transformer Programs" What are "performance variations", and how did they motivate this development? — I'm assuming this means that for some tasks, categorial worked better than score-based, and for others, the opposite was true. I think this should be stated more directly. this line would also be clearer if you clearly say that the hybrid mechanism can be both binary or continuous depending on how it is trained. 

6. The original Learning Transformer Programs paper mentions the possibility of using score-based attention, but says "It would be relatively straightforward to extend Transformer Programs to include score-based attention modules, but this would also result in longer and more complex programs, because each query must now be defined in terms of an ordering of all |K| possible keys, instead of a single key" — but in the ablation in your paper, adding uncertainty-aware attention (which has the possibility of score-based attention) decreases lines of code. Maybe having an ablation for just score-based would be a helpful comparison too. How can I square what the previous paper says with your results?

7. Is there no appendix to the paper? No examples of learned programs, or attempts to interpret a learned program? Or compare a program with the new method vs the old method? 

8. how are the programs extracted? I assume this is the same as the previous paper, but maybe something different is needed for the new modifications? what does "lines of code" even mean? this is not really made clear in the paper. 

9. how can lines of code differ between trained models?

# Suggestions
1. Fom the ablation study, it looks like just switching on Gumbel-Sparsemax indeed dramatically decreases lines of code, but also hurts performance a bit. Maybe this should be mentioned when motivating/justifying how it helps, or when analyzing the ablation.
2. some typos in citations — line 109
3. Line 178: citation for the claim that " Gumbel-Softmax often yields sub-optimal programs due to local optima and fails to encourage sparsity, hindering interpretability and efficiency" ?
4. score based attention is never clearly defined. it is confusing because categorical attention also uses a score function. so you should define it clearly as being called score-based attention (line 214 I believe)
5. typos in line 235-236.
6. lines 235-236: is the "simple neura network layer" a MLP? aren't MLP's multiple layers? I would maybe call it a network module, not a layer.
7.line 310 "we tested our models on small-scale datasets" — what does small-scale mean? reference appendix for more details?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732387235349

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to Reviewers and Revision Submitted

comment:
We thank all the reviewers for their insightful comments and constructive suggestions. We have carefully revised the paper to address the reviewers’ concerns, and we believe these revisions significantly strengthen our work. The minor revisions are summarized below, and detailed responses to each reviewer's comments follow separately. These revisions, along with the more extensive additions detailed in our individual responses, are marked in blue text in the revised PDF. We sincerely apologize for any omissions in the initial submission and commit to including all explanations and additions discussed below in the supplementary material of the camera-ready version.

The minor revisions are:

1. Added Gumbel-Sparsemax performance analysis in the Ablation Study (Section 4.2) [Reviewer U9Xh]
2. Corrected citation typos in lines 109-110 [Reviewers U9Xh, ewyT]
3. Added Gumbel-Softmax citation in line 178 [Reviewer U9Xh]
4. Refined score-based attention definition and clarified predicate/score-based terminology (line 214-216) [Reviewer U9Xh]
5. Fixed grammatical issues in lines 235-236 [Reviewer U9Xh]
6. Updated terminology from "simple neural network layer" to "network module" [Reviewer U9Xh]
7. Included normalization factor in line 96 [Reviewer ewyT]
8. Fixed quotation marks around "Transformer" in line 99 [Reviewer ewyT]
9. Clarified Transformer Program accuracy statement (lines 393-394) [Reviewer J1B6]

We greatly appreciate the reviewers' valuable feedback and are committed to incorporating all detailed explanations from our individual responses, including examples of learned programs, comparisons to the baseline, and implementation specifics, into the supplementary material of the camera-ready version. We believe that these revisions and additions will significantly strengthen the paper and better showcase the contributions of our work.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper studies an improved version of the newly proposed Transformer Programs, which are learned Transformer networks that can be converted to code. The paper proposes three architectural modifications: a smooth transition mechanism that combines Gumbel-softmax and sparsemax, uncertainty-aware gating between categorial attention and score attention, and position-aware attention.

Main strengths:
- Improvements on the novel Transformer Program framework 
- Improved performance on diverse tasks, including RASP and NLP tasks

Main weaknesses:
- It would be great to give some examples of reduced-size codes/programs.
- Some architectural improvements are not principled or novel enough.
- Reducing program size does not directly mean improved interpretability.
- The evaluation could have more NLP tasks, such as generative tasks.