Adaptive Transformer Programs: Bridging the Gap Between Performance and Interpretability in Transformers
Authors: Quoc-Vinh Lai-Dang
Taemin Kang
Seungah Son
Venue: ICLR 2025 Poster (2025)
Keywords: Mechanistic Interpretability
Transformers
Interpretable AI
Primary area: interpretability and explainable 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.

Published as a conference paper at ICLR 2025
ADAPTIVE TRANSFORMER PROGRAMS: BRIDGING THE GAP BETWEEN PERFORMANCE AND INTERPRETABILITY IN TRANSFORMERS
Quoc- Vinh Lai-Dang', Taemin Kang", Seungah Son'
Cho Chun Shik Graduate School of Mobility "The Robotics Program
Korea Advanced Institute of Science and Technology (KAIST)
Daejeon, South Korea
[idqvinh, tmkang98265 Beungahson]@kaisee ac. kr
ABSTRACT
Balancing high performance with interpretability in increasingly powerful Transformer-based models remains challenge. While mechanistic interpretabil- ity aims to specify neural network computations in explicit, pseudocode-like for- mats, existing methods often involve laborious manual analysis or struggle 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 create more robust and interpretable models. The proposed method increases expressiv- ity by redesigning two primary attention modules to improve categorical and nu- merical reasoning capabilities. To overcome optimization hurdles, we introduce a novel reparameterization scheme that enhances the exploration-explootation trade- off during training. We validate our approach through extensive experiments diverse tasks, including in-context learning, algorithmic problems (e.g., sorting and Dyck languages), and NL.P benchmarks such as named entity recognition and text classification. Results demonstrate that Adaptive Transformer Programs sub- stantially narrow the performance gap between black-box Transformers and inter- pretable models, enhancing transparency. This work advances the development of hegh-performing, transparent AI systems for critical applications, addressing crucial ethical concerns in AI development.
INTRODUCTION
Balancing high performance with model interpretability has emerged as a central challenge in arti- ficial intelligence. The introduction of Transformer architectures (Vaswani al., 2017) and the rise of large language models (LLMMs) (Brown et al. 2020) have significantly advanced natural language processing, However, these powerful models often operate as "black boxes," making it difficult to understand their decision-making processes. This issue is especially critical in fields like healthcare, finance, and law, where Al-driven decisions can have profound impacts. Addressing this challenge within the context of Transformers and LLMs is both timely and essential.
Various interpretability techniques have been proposed to illuminate how AI models make decisions, each offering unique insights yet presenting distinct challenges. Behavioral approaches, such as those by RRbbeiro et al.) (2020); Warstadt et al.) (2020). probe model responses to diverse inputs, providing an external view of model behavior but lacking access to internal reasoning mechanisms. Autribution methods like Integrated Gradients (Sundararajan et al., 2017) and SmoothGrad (Smilkov 2017) quantify the influence of input features predictions but often fail to capture underlying causal relationships. Concept- based interpretabilities (Kim et al. 2018 Belinkov 2022) adopt a top-down approach to unraveling model's decision-making processes but risk introducing biases through subjective concept selection. Mechanistic interpretability efforts (Elhage et al. 2021 Nanda
et al.
et al. 2023) delve into the internal computations of models but struggle with scalabiligg lity as model

---

Published as a conference paper at ICLR 2025
complexity grows. These limitations underscore the necessity for inherently interpretable models that offer transparent decision decision-making processes without sacrificing performance.
Advancements such RASP, Tracr, and Transformer Programs represent significant strides toward inherently interpretable Transformer models. RASP (Weiss et al., 202T) introduces a programming language that allows users to define Transformer operations in human-readabll format, effectively mapping neural computations to symbolic logic. Building on this, Tracr (Lindner et al. 2024) serves as compiler that translates RASP programs into actual Transformer weights, bridging the gap be- tween high-level specifications and low-level implementations. Transformer Programs (Friedman al. 2024) take this a step further by proposing a method train Transformers that can be directly translated into discrete, interpretable programs. While these innovations move us closer to trans- parent Al systems, challenges in expressivity and optimization persist. Our work addresses these challenges by introducing novel enhancements to Transformer Programs.
In this paper, we introduce three key innovations that enhance the expressivity and optimization of Transformer Programs (Friedman et al.) (024) while preserving interpretability. First, we pro- pose a seamless transition mechanism between Gumbel-Softmax and Sparsemax, improving the exploration-exploitation trade-off during training by allowing the model to dynamically adjust its at- tention distributions. Second, we develop an uncertainty- attention mechanism that integrates categorical and score-based attention through -hasennShannon Divergence, enabling the model to handle varying levels uncertainty in data processing. Third, we enhance the numerical mechanism by incorporating positional encodings. These contributions not only extend the functional capacity of Transformer Programs but also maintain their inherent interpretability, addressing limitations in previous approaches.
Our extensive validation on diverse tasks, including in-context learning, algorithmic problems Weiss al. 202I). and NLP benchmarks, demonstrates the effectiveness of our Adaptive Trans- former Programs. Experimental results show a substantial improvement in bridging the perfor- mance gap between black-box Transformers and interpretable models while offering enhanced trans- parency. This work not only advances the state-of-the-art in interpretable AI but also paves the way for the responsible and ethical integration of Al systems critical applications, potentially trans- forming how we develop and deploy AI in high-stakes environments.
2 BACKGROUND
Transformer Architecture and Circuits. The Transformer architecture (Vaswani et al. 2017) has revolutionized sequential data processing, achieving unprecedented performance across NLP tasks. It processes token sequences w = {w), ....... ) from a vocabulary V, converting each into high-dimensional embedding. The initial representation Xo E R combines learned token em- beddings with positional encodings, crucial for capturing sequential information. The architecture consists of L layers, each refining the input representation through two main components: Multi- Head Attention (MHA) and Multilayer Perceptron (MLP). The output of layer is computed as:
x;-1 (X- + (HA,(x;-)));
where MHA allows the model to attend to different positions within the sequence:
x x
MHA(x)
softmax
XW
Vdk
Recent research has focused on understanding Transformers through the lens of "Transformer cir- cuits" (Elhage et al. 202T), viewing them as a residual stream architecture where each component reads from and writes running representation:
x + f(x W W.
This approach has yielded insights into emergent behaviors, such as induction heads for in-context learning (Olsson al. 2022). While offering insights into attention heads and neurons, their ap-
2

---

Published as a conference paper at ICLR 2025
proach to interpretabilityy is limited by the complexity of modern models. This has led efforts to bridge the gap between symbolic reasoning and neural network models.
Bridging Transformers and Programs. Bridging the gap between Transformers and symbolic programs has emerged as a promising direction for enhancing interpretability. RASP (Weiss et al.) 202T) offers a programming language designed to express Transformer computations in a human- readable format. Its key select function, analogous to attention in Transformers, which takes quences of keys k E 'C and queries E QM along with a boolcan predicate p Q {0, 1]. to produce attention matrix A E (0, ) This is followed by agg aggregat e operation, akin to value aggregation in Transformers. Building on RASP, [indner et al. (024) introduced Tracr, a compiler that converts RASP programs into Transformer weights. This led to Learning Transformer Programs |Friedman et al. 2024), a method for training Transformers that can be automatically con- verted into discrete, human- readaale programs. This approach advances neural- saurrlllymolicc integration by combining neural network expressiveness with symbolic transparency. However, a key challenge remains implementing effective discrete optimization to ensure both accuracy and interpretability in the learned programs.
Discrete Optimization. Discrete optimization is crucial for training interpretable models with dis- crete representations like Transformer Programs. The Gumbel-Softmax estimator (Jag al., 2017) enables differentiable sampling from discrete distributions, generating one-hot encoded vectors ap- proximating discrete selections. This allows gradient- -aase optimization in Transformers. While effective, it has limitations in finding optimal solutions and promoting sparsity, crucial for inter- pretability, This paper explores alternative methods, including Sparsemax (Martins & Astudillo, 2016), and introduces a novel smooth transition mechanism addressing these limitations, leading to more effective, interpretable program learning.
3 ADAPTIVE TRANSFORMER PROGRAMS
3.1 OVERVIEW
Our approach builds upon the Transformer Programs framework (Friedman et al., 2024), which introduces two key constraints for interpretable Transformers: disentangled residual stream and rule-based modules.
The disentangled residual stream encodes each program variable in dedicated, orthogonal sub- space, preventing the entanglement often seen in standard Transformers (Vaswani et al., 2017) and facilitating clear reading and writing mechanisms. When reading, each module accesses specific variables using projection matrices parameterized by one-hot indicator vectors. Formally, if the residual stream encodes m categorical variables, each with cardinality k. resulting in input em- beddings x E 13 Nxmmk, then each projection matrix W e mm is defined by an indicator E (0, 1) W Ii.... I.] where I is the le x i identity matrix. Writing involves concatenating new information to maintain separation: x; hxi-1;((x;... where i denotes the layer and h is an attention head.
Transformer Programs enforce interpretable, rule-based mappings between input and output vari- ables. Categorical attention heads compute attention patterns using boolean predicate matrices and employ hard attention for aggregation. The attention pattern is determined using a boolean pred- icate matrix W predicate E {0. 1) kxh defining mappings between query and key values. This re- sults in an attention score matrix attention ensures each query attends to a single key, producing categorical output variable: A; = One-hot (arg max, S. Additional modules include factored categorical embeddings, lim- ited numerical attention, and feed-forward layers lookup tables.
{0, NxN where S xW Q W predicase (xWx)) TT Hard
reparameterization (Jang et al., 2017) faces challenges in finding optimal solutions and promoting sparsity. Our work addresses these challenges through three main contributions: (1) a Smooth Transition Mechanism for discrete optimization, (2) Uncertainty-Aware Categorical Attention, and (3) Position-Aware Numerical At- tention. These enhancements improve both interpretability and performance, leading to Adaptive Transformer Progeae Programs.
While effective, the original framework's use of Gumbel-Softmin
3

---

Published as a conference paper at ICLR 2025
3.2 SMOOTH TRANSITION MECHANISM FOR DISCRETE OPTIMIZATION
Our Smooth Transition Mechanism facilitates effective discrete optimization in Transformer Pro- grams (Friedman al. 2024 by gradually shifting from exploration exploitation during training. The inherently discrete nature of Transformer Programs, with modules containing discrete parame- ters like predicate matrices and gate vectors, poses challenges for traditional gradient-basee meth- ods. To address this, differentiable relaxation techniques have been employed, with the Gumbel- Softmax estimator (Jang et al. 2017) being widely adopted:
2j + 9i
2,
(4)
T
exp(2.)
Msoft, softmax;
(5)
Σ3 exp(Ez)
where is Gumbel noise, J temperature, 2; is raw logit, and ž is perturbed and scaled logit. However, Gumbel-Softmax (Jang et al., (017)) often yields sub-optimal programs due local optima and fails to encourage sparsity, hindering interpretability and efficiency.
To address these limitations, introduce Smooth Transition Mechanism combining Gumbel- Softmax and Sparsemax (Martins & Astudillo, 2016). This hybrid approach balances exploration and exploitation during training. Initially, it behaves like Gumbel-Softmax, encouraging diverse program structures. As training progresses, it shifts towards Sparsemax variant with Gumbel noise, which we term Gumbel-Sppassmax:
Ysparse = sparsemax (Z) == arg min ||p ž||2.
(6)
PEAK-1
where A KK is the (K - dimeenniil probability simplex, and the equation finds the closest point the perturbed logit ž. This promotes sparsity and more deterministic program choices, refining promising solutions and encouraging concise, interpretable programs. This balance enables the discovery of high-quality, interpretable, and efficient program structures.
The temperature parameter 7 controls the smooth transition between Gumbel-Softmax and Gumbel- Sparsemax. High т favors exploration (Gumbel-Softmaxy while low 7 promotes exploitation (Gumbel-Sparsemaxx) The transition is governed by a(T):
a(T)
T1 -
Ti
where TI, 72 are the transition points (Ti > T2). The hybrid distribution
= (1 a())) Yson + a(T) Ysparse
(8)
While improving discrete optimization, this mechanism falls short of fully addressing the need for adaptability and robustness in real-world scenarios, motivating our next contribution: Uncertainty- Aware Attention.
3.3 UNCERTAINTY-AWARE ATTENTION
Categorical attention heads in Transformer Programs, as in Traer [Lindner et al., 2024) and RASP (Weiss et al.) 2021), enforce one-to-one attention, excelling at discrete, rule-based relationships but struggling with nuanced or continuous relationships. Weiss et al.) (202I) proposed an extension to RASP combining binary function predicate Q K - {0, 1] with continuous function score Q x - R. capturing fine-grained relationships useful for tasks like semantic similar- ity in NLP. This continuous score function forms the basis of what term score-based attention. While Friedman et al.] (2024) suggested incorporating score- based attention in Transformer Pro- grams, this approach increases program complexity. Our preliminary experiments with separate
4

---

Published as a conference paper at ICLR 2025
binary (predicate- and continuous functions (score-based revealed scenario -eeerdiedeeendent per- formance variations, motivating the development of a hybrid mechanism to create more Adaptive Transformer Programs.
1991). a symmetric and smoothed version We employ Gensen-Shanner Divergence (JSD) Li
of Kullback-Leibler divergence (Kullback & Leibler 1951), to measure uncertainty in categori- cal attention. JSD's non-negativity, boundedness, and symmetry make it suitable for this task. In this context, JSD measures uncertainty in probability distributions, where higher values indicate greater uncertainty and a larger divergence between categorical attention and reference distri- bution. Given query q and keys ky..... ... .... we define categorical attention (CatAttention) as Acar, = predicate(q, k;) and score-based attention (ScoreAttention) as Ascore = score(g. kz). We introduce dynamic reference attention A that adapts during training, allowing flexible uncer- tainty estimation. The ISD is formulated as:
1 c Ardii) KL((ca,, Auvgs) KL
||
Awga)
where Amg. = (Acat, //2 This formulation enables uncertainty estimation in the attention mechanism, which used to adjust attention weights accordingly.
A learnable gating mechanism, driven by the JSD-based uncertainty estimate, dynamically weights the contributions of CatAttention and ScoreAttention. The gating mechanism: g = MLP(JSD(AAu.) Ae),,)), that implemented input and outputs a gating weight between and 1. This weight is used to combine the outputs of CatAttention and ScoreAttention:
network module that takes the JSD value as
A; = Acti +
Assor,
(10)
In high uncertainty scenarios (high JSD), the gate favors ScoreAttention, which is more reliable in uncertain contexts. In low uncertainty (low JSD), CatAttention is preferred, as it is more confident in its categorical decisions. This adaptive mechanism provides flexible and robust decision-making by dynamically adjusting the balance between attention types based on uncertainty.
While Uncertainty-Aware Attention improves the handling of categorical and contextual informa- tion, processing numerical data poses additional challenges. To address this, we introduce the Position-Aware Attention module, which extends the original Numerical Attention mechanism.
3.4 POSITION-AWARE ATTENTION
The numerical attention mechanism in Transformer Programs (Friedman al.) 2024 is restricted to outputting integer values within a bounded range. This limitation hinders the model's ''s ability to represent and process continuous fractional values, thereby reducing expressiveness and compli- cating tasks that require nuanced numerical representations or complex calculations. Additionally, numerical attention employs binary predicate matrix and computes a weighted sum instead of weighted average, simplifying standard Transformer attention and diminishing the model's 's capacity to capture intricate relationships between inputs. Furthermore, numerical variables are limited to being either constant (set to one at the input layer) or outputs of numerical attention heads, con- straining the model's ability to learn and represent arbitrary numerical values and restricting its problem-solving capabilities.
Our Position-Aware Attention mechanism extends the numerical attention in Transformer Programs (Friedman et al. 2024). It uses categorical variables keys and queries, and numerical variables as values. We incorporate both Learnable (Gehring et al. (017) and Sinusoidal (Vaswani et al.) 2017) Positional Encodings into the numerical value variable var, creating a position-aware value va pus- This allows the model to learn nuanced positional relationships. Given attention scores S E (0, 1) NxN, the output for the fth token is computed as:
A Apum,
Σ Sayaar
(11)
5

---

Published as a conference paper at ICLR 2025
The integration of learnable and sinusoidal positional encodings offers several advantages. Learn- able encodings capture nuanced positional information and enable processing of non-integer values, expanding the model's numerical capabilities for tasks requiring fine-grained representations, By incorporating positional information, each token can distinguish between identical numerical values different positions. Additionally, Sinusoidal Positional Encodings provide structured approach to embedding positional information, improving the model's ability to handle sequence-dependeee tasks and maintain interpretability. These enhancements, combined with the Smooth Transition Mechanism and AhccertaintyAAwar Attention, enable Transformer Programs to be effectively con- verted into interpretable programs, as discussed in the subsequent section on Experimental Results.
Table I: Accuracy (Acc.) and Program Length (Lines) for Transformer Programs (Baseline) and Adaptive Transformer Programs (Ours) on In-Context Learning (Friedman al. 2024) and RASP tasks Weiss et al. (202I)
Description
Example
Baseline
Ours
Acc.
Lines
Acc.
Lines
Induction
context learning.
induct (*a (*a1n2b2a"")
101
Reverse
Reverse order.
reverse (" ['abbc') "cbba"
99.74
859
99.99
779
Histogram
Count the number of tokens.
hist ("abbe") = "1221"
99.94
199
99.95
189
Double
number unique tokens sharing identical frequency count.
hist2 (*abbee *2112*
91.81
513
Sort
Arrange the input elements in alphabetically ascending onder.
sort ("cbba") "abbe"
99.98
945
99.86
895
Most-Preq
Order unique elements by queney, using carlier positions to break ties.
fre-
moat..treg (Mabbo") "bac"
76.44
1334
80.80
894
Dyck-1
Clussiy if each position is valid sTring(T). valid (Pe invalid(F).
dyckl (* (1)) *SUere"
99.69
1297
99.93
1086
Dyck-2
The same amalysis with abuve. but in Dyck-
cyck2 (1) (1)*)
"PPPPIE""
97.98
1316
98.14
1065
EXPERIMENTS
4.1 EXPERIMENTAL SETUP AND DATASETS
To evaluate the effectiveness of Adaptive Transformer Programs, we conducted experiments on di- verse set of tasks: a simple in-context learning task (Friedman et al. 2024). algorithmic RASP tasks (Weiss al., (021). and two standard NLP tasks -akked entity recogmition using CoNLL-2003 TTjon Kim Sang & De Meulder 2003) and text classification using TREC, MR, Subj. and AG News datasets (Voorhees & Tice| 2000; Pang Lee 2004 2005) Zhang al. 2015). In the in-context learning task, the model processed sequences of up to TO tokens of alternating letters and numbers from a vocabulary of four letters and four numbers, outputting the number following a repeated let- ter unk for new letter, using attention-only Transformer with two layers and one attention head per layer, fixed one-hot encoded token and position variables as input, and causal attention mask. For the RASP tasks, summarized in Table I] tested our models small-scale datasets with sequence lengths up to 16 (Dyck tasks) or 8 (others), using vocabularies matching the sequence lengths; the models employed fixed one-hot token and position embeddings, variable cardinality set to the maximum sequence length, and incorporated our enhanced egorical attention and Aesition-Awart numerical attention heads and MLPs with two input variables. In the NLP tasks, sentences were limited to 32 words for named entity recognition and 64 words for text classification, both using a 10,000-word vocabulary and initialized with 300-dimensional Glo Ve embeddings (Pennington et al.) 2014); for named entity recognition, only categorical attention heads and MLPs were employed, while text classification used averaged token embeddings for sentence representation.
cat-
4.2 IN- CN-CONEXTT LEARNING AND RASP TASKS
Performance.
Table compares Adaptive Transformer Programs (Ours) Transformer Programs
(Friemma et al. 2024) (Baseline) across eight tasks, showing improvements in accuracy and pro- gram complexity. Our approach consistently matches outperforms the baseline, with notable gains

---

Published as a conference paper at ICLR 2025
in challenging tasks: Double histogram (91.81% 66.78%) and Most-Freq (80.80% vs 76.44%). For tasks like Induction and Sort, where baseline accuracy was already high, our model maintains performance (100% and 99.86 99.86% respectively) while reducing program complexity. Slight accuracy improvements are observed in Reverse, Histogram, and Dyck-1 tasks. Interestingly. for the Dyck-2 task, which involves from 97. 77.8% to 98. 14%. These results demonstrate that Adaptive Transformer Programs maintain high accuracy across diverse RASP tasks while showing significant improvements in challenging scenarios, highlighting the effectiveness of our proposed enhancements in complex reasoning tasks.
complex string analysis, our model achieves small increase in accuracy
Interpretability. A key advantage of Adaptive Transformer Programs is their ability to achieve more concise and interpretable representations, evidenced by the reduced program length (lines of code) across all tasks in Table This improved sparsity not only suggests greater computational effi- ciency but also enhances interpretability, making it easier to understand the model's decision- making process. Notably, our approach achieves substantial reductions in program length for both complex tasks like Dyck-1 (16.3% reduction) and Dyck-2 (19.1 1% reduction) and simpler tasks like Induction (5.6% reduction) and Histogram (5% reduction). The most significant reduction is observed in the Most-Freq task (33% reduction, from 1334 to 894 lines). These improvements in conciseness, com- bined with maintained or improved accuracy across tasks, demonstrate the effectiveness of Adaptive Transformer Programs balancing performance with interpretability.
Reverse
Sort
Most Freq
100
Accuracy
20
Ours
Program Standard
Ours
Program Standard
Ours
Program Standard
| 8, =
V =
[ 16,
Figure 1: RASP accuracy comparison of Adaptive Transformer Programs (Ours), Program Trans- formers (Program), and Standard Transformers (Standard) across increasing input vocabulary sizes ((V) and sequence lengths (N).
Scalability. Adaptive Transformer Programs demonstrate robust scalability when faced with in- creasing input complexity, outperforming both Program Transformers and Standard Transformers in some scenarios. Figure illlustrates the performance across three representative RASP tasks (Re- verse, Sort, and Most Freq)
increase the input vocabulary size ((V)) from 8 to 16 and the maximum sequence length (N) from to 16. For the Reverse task, our model maintains high ac- curacy (99.99% with | 8 and N 8, and experiences less degradation 183.11% and 59.35%) compared to baselines as complexity increases. In the Sort task, Adaptive Transformer Programs consistently outperform Standard Transformers and show better resilience than Program Transform- ers, maintaining 87.6% accuracy even at |V| = 16 and N = 16. The Most Freq task presents challenge for all models, but our approach still demonstrates competitive performance, particularly at higher complexities. These results highlight the superior scalability of Adaptive Transformer Programs, showing their potential for handling more complex, wallworldd data distributions while maintaining interpretability.
Ablation Study. Table presents our comprehensive ablation study, evaluating the impact of three key enhancements: the Smooth Transition Mechanism, Uncertainty-ANn Attention, and Position- Aware Attention on the Most-Freq task. We manually enable or disable these components in various combinations, training separate models for each configuration. The results reveal that the full model, with all enhancements enabled, achieves the highest accuracy of 80.8% while maintaining rela- tively concise program length 894 lines. As observed the ablation study where only Smooth Transition is enabled and the others disabled, while the Smooth Transition Mechanism (which incor- porates Gumbel- Sparsee--pa contributes significantly to program conciseness, also shows a slight decrease in accuracy when used in isolation, highlighting the trade-off between interpretability and

---

Published as a conference paper at ICLR 2025
Table Ablation study on Most-Freq task: Impact of model enhancements on performance (Accu- racy) and program length (Lines).
Smooth Transition Mechanism
Uncertainty-Awe Attention
Position-Aware Attention
Accuracy
Lines
80.8
894
78.25
850
75.01
939
78,62
938
73.61
920
77.34
76.26
1028
76,44
1334
performance. Disabling individual components leads to performance drops, with Uncertainty-Aware Attention showing the significant impact (accuracy decrease 75.01 % when disabled). No- tably, removing all enhancements (equivalent to baseline Transformer Programs) results in lower accuracy (76.44%) and the longest program (1334 lines), highlighting the cumulative benefit of our proposed enhancements. When multiple components are disabled, the performance declines setther-reoovnng both Uncertainty-Aware Attention and Position-Aware Attention leads an ac- curacy of 73.61 1% and a less efficient program length of 920 lines. These results demonstrate the critical role each enhancement plays in maintaining high performance and concise, interpretable program structures.
4.3
NLP TASKS
Table NER Performance Metries and Program Length on CoNLL--2003 Dataset.
Model
Accuracy
Precision
Recall
F1
Lines
Standard Transformers
92.2
71.1
62.5
66.6
Transformer Programs
94.2
78.9
72.9
75.8
991
Ours
94.1
77.2
73.2
75.1
916
Named Entity Recognition (NER). On the CoNLL-2003 Named Entity Recognition (NER) task, standard benchmark for sequence labeling, Adaptive Transformer Programs demonstrate com- petitive performance while offering significant advantages in interpretability. Tabll 3 presents the results, comparing our approach Standard Transformers and Transformer Programs. Our model achieves 94.1% accuracy, closely matching the 94.2% of Transformer Programs and significantly outperforming the 92.2% of Standard Transformers. Although the scores are similar across Transformer Program and approach (75.8 and 75.1, respectively), Adaptive Transformer Pro- grams achieve this performance with a notably shorter program length (916 lines compared to 991). This conciseness, indicative of greater program sparsity, is key advantage, promoting easier anal- ysis and understanding of the learned programs, crucial aspect for interpretability. This result highlights the ability of Adaptive Transformer Programs maintain competitive performance while generating more interpretable program representations.
Table 4: Accuracy and Program Length for Various Text Classification Tasks.
TREC
MR
Subj
AG
Model
Acc.
Lines
Acc.
Lines
Acc.
Lines
Acc.
Lines
Standard Transformer
83,4
75.9
90,9
89.1
Transformer Program
84.2
5520
77.1
3972
92.3
3065
90.3
1881
Ours
83.6
827
77.9
773
1954
90.0
1790
Text Classification. Adaptive Transformer Programs exhibit robust performance across diverse text classification tasks, demonstrating their capacity for generalization to various real-world sce-
8

---

Published as a conference paper at ICLR 2025
narios. As shown in Table 4 our approach performs competitively on four distinct classification tasks: TREC (question type identification), MR (sentiment evaluation), Subj (subjectivity assess- ment), and AG news (topic categorization).. Notably, model achieves the highest accuracy the MR task at 77.9%, surpassing both Standard Transformers (75.9%) and Transformer Programs (77.1%). While performance in accuracy, the most striking difference lies in program length, Across all tasks, Adaptive Trans- former Programs consistently produce more concise programs. For instance, on the TREC task, our model requires only 827 lines compared to 5520 lines for Transformer Programs, representing an 85% reduction in complexity. This significant decrease in program length, coupled with competitive accuracy, underscores the efficiency and interpretability of our approach in handling diverse text classification challenges.
other tasks is comparable to the baselines, with slight variations
5 RELATED WORK
Learning Programs. Program synthesis has evolved from classical symbolic approaches to deep leamiig-based methods, driven by the need to scale to complex programs learned by modern neural architectures. Traditional paradigms like Inductive Logic Programming (Muggleton & de Raedi 1994] and Deductive Program Synthesis (Manna Waldinger 1980) relied on symbolic reasoning and expert knowledge. The field then shifted towards neural program induction, with works like (Reed de Freitas 2016) and Neuro-Symbolic Program Synthesis (Devlin et al. (017) learning programs directly from data. However, these methods struggle with scalability to large datasets, complex program structures, and incorporating domain-specinn knowl- edge (Gulwani et al.) 2017). Transformer Programs (Friedman et al., (2024) address these limitations by leveraging Transformer architectures' representation learning capabilities while imposing con- straints to learn interpretable programs.
Neural ceeeraamme--Ittrprettte
Transformers and Formal Languages. Recent research has demonstrated the expressive power of Transformers in relation to formal languages. Studies show that Transformers can learn regular and context-free languages, and implement algorithms like first-order logic with majority quantifiers (Hahn 2020 Merrill & Sabharwal 2022). Work by Giannou et al. (2023) further supports the view of Transformers as general-purpose computation devices. Weiss et al. (2021) established an initial connection between Transformer operations and program-like representations through the RASP language. Adaptive Transformer Programs build on this foundation, enhancing interpretability and programmatic representation align Transformers with human-understanddbble symbolic systems.
Interpretable Machine Learning Models. The field of interpretable machine learning has seen a surge in methods for understanding deep learning models. Post-hoc methods include attention visualization (Bahdanau et al., 2014). feature attribution (Ribeiro et al.) 2016; [Luddberg & Leep 2017), and concept activation vectors (Kim et al. 2018). Architectural modifications, such as sparse attention (Zhang et al., 202T) and inductive biases (Geiger et al.) 2024), attempt to enhance inter- pretability through model design. In contrast, intrinsically interpretable models offer direct access to underlying algorithms, improved transparency, and the potential for formal verification. Our Adap- tive Transformer Programs aim to learn inherently interpretable models, providing more faithful and complete explanations of decision-making processes. This approach addresses the limitations of post-hoc methods and architectural modifications by representing complex computations transpar- ently.
Uncertainty in Deep Learning. Uncertainty estimation plays a crucial role in developing reli- able and interpretable deep learning models. Quantifying uncertainty improves model reliability (Kendall & Gal] 2017). facilitates human-Al collaboration (Gal & Ghahramani 2016), and en- hances interpretability (Leibig et al.) (017).. Prominent techniques include Bayesian Neural Net- works (MacKay 1995), Monte Carlo Dropout (Gal & Ghahramani, 2016), and Ensemble Methods Lakshminarayanan et al. 2017). Our Uncertainty- Awaar Attention mechanism dynamically com- bines attention types based uncertainty estimates, leading more robust and interpretable mod- els. This approach uniquely integrates program synthesis, Transformer architectures, interpretabil- ity, and uncertainty estimation.

---

Published as a conference paper at ICLR 2025
6 CONCLUSION AND DISCUSSION
The study presents a novel framework for developing robust, expressive, and interpretable Trans- former models that can be translated into human-readable programs. Key contributions include Smooth Transition Mechanism for discrete optimization, an Uncertainy--ware Attention mecha- nism for adaptive attention blending, and a Position-Awar Attention module for numerical reason- ing. Empirical results on synthetic and real-world NLP tasks demonstrate superior performance compared benchmarks and provide concise, insightful interpretability analysis. This work ad- vances interpretable Al by improving performance and clarity through new adaptive mechanisms, facilitating the creation of transparent and reliable AI systems. Additionally, it explores the con- vergence of program synthesis, deep learning, and uncertainty estimation, promoting accountability and the societal benefits AI.
Integration of Contributions. The three enhancements introduced in this work synergistically contribute to the effectiveness of Adaptive Transformer Programs. The Smooth Transition Mecha- nism promotes program sparsity by gradually shifting from exploration-focused Gumbel-Softmax to applotaationffocusee
Uncertainty-Awart Attention dynamically adapts the at- tention strategy based on uncertainty estimates, enhancing expressiveness and robustness. Position- Aware Attention improves numerical reasoning and training stability through positional encodings.
Future Research Directions. Adaptive Transformer Programs exhibit potential but face chal- lenges such scaling to larger models and tasks, necessitating improved training methods com- pact representations. The complexity of these programs can impede human understanding, high- lighting the need for simplification, summarization, or visualization techniques. Extending their application to computer vision or robotics will require adapting knowledge representation and ex- traction processes. Future research might explore advanced structures like recursion and hierarchical composition for greater expressiveness and leverage interpretability to support human-Al collabora- tion through interactive program refinement tools.
REFERENCES
Dzmitry Bahdanau, KyungHyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning align and translate. arXiv preprint arXiv:1409.0473, 2014.
Yonatan Belinkov. Probing elassifiers: Promises, shortcomings, and advances. Computational
Linguistics, 48(1), March 2022.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901, 2020.
Jacob Devlin, Jonathan Uesato, Rishabh Singh, and Pushmeet Kohli. Semantic code repair using neuro-symbolic transformation networks. arXiv preprint arXiv: 1710. 1054, 2017.
Nelson Elhage, Neel Nanda, Catherine Olsson, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodee Danny Hernande7, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical frame- work for transformer circuits. Transformer Circuits Thread, 1(1):12, 2021.
Dan Friedman, Alexander Wettig, and Danqi Chen. Learning transformer programs. Advances in Newral Information Processing Systems, 36, 2024.
Yarin Gal and Zoubin Ghahramani. Dropout as bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pp. 1050-1059.
PMLR, 2016.
10

---

Published as a conference paper at ICLR 2025
Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann N Dauphin. Convolutional sequence to sequence learning. In International conference on machine learning, pp.
PMLR. 2017.
Atticus Geiger, Duligur lbeling, Amir Zur, Maheep Chaudhary, Sonakshi Chauhan, Jing Huang. Aryaman Arora, Zhengxuan Wu, Noah Goodman, Christopher Potts, and Thomas Icard. Causal abstraction: A theoretical foundation for mechanistic interpretability. 2024.
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. International Conference on Machine Learning. pp. 11398-1 11442. PMLR, 2023.
Sumit Gulwani, Oleksandr Polozov, Rishabh Singh, al. Program synthesis. Foundations and Trends@ in Programming Languages, 4(1-2):1-119, 2017.
Michael Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics, 8:156-18 171, 2020.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. In International Conference on Learning Representations, 2017.
Alex Kendall and Yarin Gal. What uncertainties do we need bayesian deep learning for computer vision? Advances in neural information processing systems, 30, 2017.
Been Kim, Martin Wattenberg, Justin Gilmer, Carrie Cai, James Wexler, Fernanda Viegas, and Rory Sayres. Interpretability beyond feature attribution: Quantitative testing with concept activation vectors (tcav). International conference on machine learning, pp. 2668-666880 1678- PMLR, 2018.
Solomon Kullback and Richard A Leibler. On information and sufficiency. The annals of mathe- matical statistics, 22(1):79-888 86, 1951.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predic- tive uncertainty estimation using deep ensembles. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Christian Leibig, Vaneeda Allken, Murat Seçkin Ayhan, Philipp Berens, and Siegfried Wahl. Lever- aging uncertainty information from deep neural networks for disease detection. Scientific reports, 7(1):1-14, 2017.
Jianhua Lin. Divergence measures based on the shannon entropy. IEEE Transactions Information theory, 37(1))145-15 1991.
David Lindner, János Kramár, Sebastian Farquhar, Matthew Rahtz, Tom McGrath, and Vladimir Compiled transformers as a laboratory for interpretability. Advances Neural Information Processing Systems, 36, 2024.
Mikulik. Tracr:
Scott M. Lundberg and Su-In Lee. A unified approach to interpreting model predictions. NIPS' 17, pp. 4768-4777, Red Hook, NY, USA, 2017. Curran Associates, Inc. ISBN 9781510860964
David JC MacKay. Bayesian neural networks and density networks. Nuclear Instruments and Methods in Physics Research Section A: Accelerators, Spectrometers, Detectors and Associated Equipment, 354(1):73-8))) 1995.
Zohar Manna and Richard Waldinger. deductive approach to program synthesis. ACM Trans.
Program. Lang. Syst. 2(1):90-121, 1980.
Andre Martins and Ramon Astudill.. From softmax to sparsemax: A sparse model of attention and multi-label classification. In International conference on machine learning, pp. 1614-1623. PMLR, 2016.
William Merrill and Ashish Sabharwal. Transformers implement first-order logic with majority quantifiers. 2022.
11

---

Published as a conference paper at ICLR 2025
Stephen Muggleton and Luc de Raedt. Inductive logic programming: Theory and methods. The Journal of Logic Programming, 19-20:629-677 6799 1994.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress mea- sures for grokking via mechanistic interpretability. In The Eleventh International Conference on Learning Representations, 2023.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, 7ac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. arXiv preprint arXiv:2209. 11895, 2022.
Bo Pang and Lillian Lee. A sentimental education: Sentiment analysis using subjectivity summa- rization based on minimum cuts. In Proceedingy of the 42nd Annual Meeting of the Association for Computational Linguistics (ACL-04), pp. 271-278, 2004.
Bo Pang and Lillian Lee. Seeing stars: Exploiting class relationships for sentiment categorization with respect to rating scales. In Proceedings of the 43rd Annual Meeting of the Association for Computational Linguisties (ACL '05), pp. 115- 124, 2005.
Jeffrey Pennington, Richard Socher, and Christopher D Manning. Glove: Global vectors for word representation. In Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP), pp. 1532-1543. 2014.
Scott E. Reed and Nando de Freitas. Neural programmerrineerpetter.. In 4th International Confer- ence on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Confer- ence Track Proceedings, 2016.
Marco Tulio Ribeiro, Sameer Singh, and Carlos Guestrin. " why should i trust you?" explaining the predictions of any classifier. In Proceedings of the 22nd ACM SIGKDD international conference on knowledge discovery and data mining, pp. 1135-1144, 2016.
Marco Tulio Ribeiro, Tongshuang Wu, Carlos Guestrin, and Sameer Singh. Beyond accuracy: Be- havioral testing of NLP models with CheckList. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 4902-4912, 2020.
Daniel Smilkov, Nikhil Thorat, Been Kim, Fernanda B. Viégas, and Martin Wattenberg. Smooth- grad: removing noise by adding noise. CORR, abs/1 706.7006..80 038250 2017.
Mukund Sundararajan, Ankur Taly, and Qiqi Yan. Axiomatic attribution for deep networks. In International conference on machine learning, pp. 3319-3328, PMLR, 2017.
Erik F. Tjong Kim Sang and Fien De Meulder. Introduction to the CONLL-2003 shared task: Language-independent named entity recognition. In Proceedings of the Seventh Conference on Natural Language Learning at HLT-NAACL 2003, pp. 142-147, 2003.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances Neural Infor- mation Processing Systems, volume 30, 2017.
Ellen M Voorhees and Dawn M Tice. Building a question answering test collection. In Proceed- ings of the 23rd annual international ACM SIGIR conference on Research and development in information retrieval, pp. 200-207, 2000.
Alex Warstadt, Alicia Parrish, Haokun Liu, Anhad Mohananey, Wei Peng. Sheng-Fu Wang, and Samuel R. Bowman. BLiMP: The benchmark of linguistic minimal pairs for English. Transac- tions of the Association for Computational Linguistics, 8, 2020.
Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Confer- ence on Machine Learning, pp. 11080-1 1090. PMLR, 2021.
Biao Zhang, Ivan Titov, and Rico Sennrich. Sparse attention with linear units, 2021.
Xiang Zhang, Junbo Zhao and Yann LeCun. Character-level convolutional networks for text clas- sification. Advances in neural information processing systems, 28, 2015.
12

---

Published as a conference paper at ICLR 2025
A PROGRAM EXTRACTION AND INTERPRETATION
A.1 DETAILED EXTRACTION PROCESS
A.1.1 DEFINITION OF "LINES OF CODE" METRIC
To quantify the complexity and interpretability of extracted programs, we define "lines code" (LOCC) as the number of active program components derived from trained Transformer model. The LOC metric serves as a quantitative measure program complexity and interpretability. In our framework, LOC is formally defined as the sum of active rules or operations in the symbolically extracted program representation. Specifically, LOC is calculated by summing the contributions from the following program components:
Predicate Functions: For each predicate function, we count the number of conditional branches (e.g., if,, el else) and the return statement. Each such statement consti- tutes one line code.
Attention Pattern Selection: Each invocation of functions like se select_cl osest or se lect to determine attention patterns is counted one line. These functions represent the core attention mechanism.
Each call to aggregation functions such as aggregate or aggregate_sum, which combine attention outputs, contributes one line to the LOC.
Aggregation Functions:
Output Score Calculation: The operations involved in accessing and processing the classifies fier.weight to compute output scores are also counted. Typically, indexing operations and subsequent processing contribute to the LOC.
A.1.2 PROGRAM EXTRACTION METHODOL.OGY
Our program extraction methodology translates learned model parameters into symbolic code, adapting to the Gumbel--ppasseax transition and uncertainty-aware attention mechanisms. The key steps are detailed below:
Gumbel-Sparsemax Transition. The Smooth Transition Mechanism, employing blend of Gumbel-Softman and Sparsemax, primarily influences the training process and the resulting learned parameters. It does not necessitate changes to the fundamental program extraction methodology it- self. The extraction process for predicate matrices, attention weights, and aggregation operations remains consistent with the original Transformer Programs. The effect of the Gumbel-Sparsenne transition observed indirectly in the extracted programs: the learned predicates tend to be sparser, and the attention patterns are more focused. These characteristics, which contribute to program sim- plicity and reduced "lines of code," are a consequence of the improved optimization facilitated by the Smooth Transition Mechanism, rather than a direct change in the extraction procedure.
Uncertainty-Aware Attention. The Uncertainty-Awart Attention mechanism introduces a dy- namic blending of categorical and score-based attention, guided by Jensen-Shannon Divergence (JSD)-based gating. While the extraction process for the individual components (categorical pred- icates, score functions, and aggregation) remains fundamentally similar, the key adaptation lies in how we interpret the extracted program and understand the influence of uncertainty awareness.
The core idea is that uncertainty- training shapes the learned parameters of the categorical predicates and score functions to be more robust and contextually appropriate. During extraction, we observe the resulting predicate logic and attention patterns which reflect the effects of this uncertainty-wware training. For instance, due uncertainty-aware gating:
More Generalized Predicates: Extracted predicates may exhibit greater generalization, being less sensitive to minor input variations, reflecting the influence score-base base- simi- larity under high uncertainty.
The program may demonstrate groupings of query positions with similar behaviors, even if their categorical features are not identical, indicating learn ing driven by contextual similarity guided by uncertainty.
Context-Aware Grouping:
13

---

Published as a conference paper at ICLR 2025
In summary, the extraction process for ATP retains the core principles of the original Transformer Programs method. The novelty lies in the enhanced training mechanisms that lead to qualitatively different and more interpretable extracted programs. The "lines of code" metric effectively captures the resulting reduction in program complexity stemming from these mechanisms.
A.1.3 IMPACT OF UNCERTAINTY-AWAR ATTENTION ON PROGRAM STRUCTURE
While the extraction methodology remains largely consistent, Uncertainty Aware Attention signifi- cantly impacts the structure and complexity of the extracted programs. This mechanism allows model to adapt its attention strategy dynamically based on uncertainty, leading to more robust and efficient program logic. During training, this dynamic gating, driven by the JSD-based uncertainty estimate, encourages the model to:
Learn Robust Predicates: By blending with score-based attention when categorical tention is uncertain (high JSD), the model learns categorical predicates that are less brittle and more generalized. This reduces the need for highly specific and numerous conditional clauses within predicates, thus simplifying predicate definitions and decreasing LOC.
Develop Context-Appropriate Attention: The adaptive blending allows the model to uti- lize categorical attention when confident (low JSD) and rely more on score- score-based attention in ambiguous noisy contexts (high JSD).. This context-dependent attention strategy leads to more efficient information processing and potentially reduces the complexity of subse- quent program modules.
Achieve Concise Program Logic: The overall effect of learning more robust predicates and coproot-appropriate attention is a more streamlined and concise program structure. This manifests as fewer "lines of code" due to simplified predicates, more efficient module interactions, and reduced redundancy in program logic.
In essence, Uncertainty-Awant Attention guides the learning process favor program structures that are not only accurate but also more robust and interpretable by dynamically adapting attention strategies based on the context-dependent uncertainty.
A.1.4 DIFFERENCES IN "LINES OF CODE" ACROSS TRAINED MODELS
The "lines of code" metric can vary significantly between trained models, even for the same task, reflecting differences in the learned program structures. This variability arises from several factors:
Optimization Path and Randomness: Neural network training involves stochastic op- timization. Different random initializations and training trajectories can lead to models converging different, yet functionally similar, solutions. Some solutions may be more concise and interpretable (lower LOC) than others due to variations in the optimization path taken.
Effectiveness of Learning Mechanisms: The effectiveness of mechanisms like Gumbel- Sparsemax and Uncertainty-Awart Attention in guiding the model towards simpler and more efficient solutions directly influences the final program length. Models trained with ATP, leveraging these mechanisms, tend to learn programs with lower LOC compared to baseline models.
Underlying Program Structure Efficiency: Different models might discover different algorithmic approaches to solve the same task. Some algorithms are inherently more com- plex (requiring more steps and rules) than others. ATP aims to facilitate the discovery of more efficient, algorithmically simpler solutions, leading to reduced LOC.
A.2 EXAMPLE PROGRAM COMPARISONS
This section presents side-by-side comparisons of extracted programs for the Double Histogram and Histogram tasks, demonstrating the improved conciseness and interpretability achieved by Adaptive Transformer Programs (ATP) compared to baseline Transformer Programs. We analyze the specific program structures and quantify the reduction in "lines of code" (L.OC) and accuracy improvements.
14

---

Published as a conference paper at ICLR 2025
A.2.1 DOUBLE HISTOGRAM TASK COMPARISON
We compare code snippets from Layer 2, Attention Head 0 (attn.1_0) and Attention Head (attn.1.1) for both baseline and ATP models trained on the Double Histogram task.
Baseline Transformer Program
Baseline code: attn_l_0
def predicate_l_0 (g.positon, k_position)- if q_position in (0):
return k_position =
elif q_position in (1):
return k_position = 2
elif g_position in (2):
return k_position == 3
elif position in (3):
return k_position == 5
elif g_position in (4, 5):
return k_position =
elif p_position in [6]:
return k_position == 1
elif q_position in [7):
return k_position
attn_1_0_pattern = select_closes (positions, positions, predicate_ 1_0)____________________________
Adaptive Transformer Program (ATP)
Our code: attn_1_0
def (
k_position);
# Grouped positions (0, 3, 4} mapped to k_position
if pososition in (0, 3, 4)
return p_position
# Grouped positions: (1, 2, 5) mapped to k_position
elif q_position in [1, 2, 5)
return k_position
elif p-position in (6):
return k position
elif q_position in {7):
return posoitio
attn_1_00 1_t_ter =
(positions, positions, predicate_1_0)
Figure 2: Side-by-Sice Comparison of attn_1_0 Predicates for Double Histogram Task. The ATP program groups query positions (e.g., {0. 3, and (1, 53) improve efficiency.
Analysis of attn.11 1..0 Comparison (Figure
Grouped Query Positions: The ATP program groups query positions within the if and elif conditions in predicate 1 O, as highlighted in q position in {o, 3, 4} and
position in {1, 2, 5]. The baseline program lacks this grouping, having separate ditions for each query position (e.g., e ilff q.positi iion in (0):, elif q.position in (3); elif g.position in {4}: in a fully expanded version).
Concise Predicate Logic: Grouping query positions significantly reduces the number of condi- tional checks, leading more concise logic in ATP's predicate. ATP's prediccate.ee. has fewer elif clauses compared to fully expanded baseline predicate.
Line Count Reduction:
Bascline predicate 1 0 (expanded): ss lines (if fully expanded based on separate conditions)
15

---

Published as a conference paper at ICLR 2025
ATP predicate.1.0: 7 lines
Baseline Transformer Program
Baseline code: attn_l_1
def predicate___1 (q_positoonn, k_position)
if q_position in (0):
return k_position =
elif q_position in (1, 3):
return k_position = elif ppposition in (2):
return k_position ==
elif q_position in (4):
return K_position = 2
elif g_position in :
return k_position =
elif q_position in (6, 7):
return k positioo ==
attn 1 pattern select_closed (positions, positions, predicate_1i 1_1)
Adaptive Transformer Program (ATP)
Our code: attn_1_1
def ( 0_ttpputt position):
# Conditions on attn_0_0_oon ______uuttttttttt
if attn_t_Q_outputt in (0): # Condition prev. output
return position
elif attn_0_0_outee in {1, 5, 7): # Condition prev. output
return position == 6
elif attn_0_0 ____t_u__ttuuttttt in (2, 3, 4) # Condition prev. output
return position ==
attr attr_l__paater
select_closest (positions, attn_0 predicate_1_]
Figure 3: Side-by-Side Comparison of at attn.l.1 Predicates for Double Histogram Task. Note that
ATP uses conditions based on att....0ooutuut.
Analysis of attn 1.1 Comparison (Figure
conditions attn_0_0 put (output of the previous attention head) instead of just positi t on as in the baseline. This is highlighted if attn.0.C.outpur in {0}: This indicates more hicrar- chical and efficient logic, where attention decisions are informed by the outputs of carlier processing stages.
Conditioning Previous Attention Output: ATP's predicate
Compact Logic via Output Conditioning: By conditioning on the output of attn_0_0, ATP achieves significantly more compact predicate 1 with fewer clauses. The baseline predi- cate requires more specific positional rules, whereas ATP leverages the information processed by at attn___0.
Line Count Reduction:
Baseline 1re 9 lines
ATP predicate.111: 7 lines
Overall Impact on Double Histogram Task:
The combined effeet of grouped query positions in attn_l_0 and output-conditions logic in attn 1 1 in ATP leads to a substantial reduction in program length and improved accuracy for the Double Histogram task:
16

---

Published as a conference paper at ICLR 2025
Line Count Reduction: 586 lines (Baseline) - 513 lines (ATP)
Accuracy Improvement: 66.78% % (Baseline) - 91.81% % (ATP)
This example clearly demonstrates how ATP learns more concise and efficient programs by leverag- ing its novel mechanisms. The grouped query positions and hierarchical, output-conditioned logic contribute to both improved interpretability (shorter program, easier to understand logic) and en- hanced performance.
B
METHOD DETAILS AND ANALYSIS
B.1
SMOOTH TRANSITION MECHANISM
B.1.1 VISUALIZATION OF a(T) EVOLUTION OVER TRAINING
Mathematical Formulation of the Smooth Transition Mechanism:
The Smooth Transition Mechanism blends Gumbel-Softmax and Sparsemax using the interpolation parameter (T), which is dynamically adjusted based the temperature
Gumbel-Sofftma Distribution:
exp(2;)
Vsott, softmax, (8)
Σg exp(2)
where & = (z; 0,)/T, 2, the raw logit, and gi is Gumbel noise.
Gumbel-Sparsemax Distribution:
Ysparse, = sparsemax(2) = min PEAR-1
lpp - 21/
where AK is the (K 1) dimensional probability simplex.
Smooth Transition Interpolation:
= (1 - (<(T)) Yoo + a((T) yqparse
Transition Function a(T):
i
((T)
if <T <T
where T₁ and T2 are predefined transition temperatures (71 T2). In practice we use simplified linear function described in the main paper,
a = -
Ti
for T2 < Tı, clamped for 2 Ti and 0 for S 72.
B.1.2 COMPARISON WITH FIXED GUMBEL-SOFTMAN AND SPARSEMAX
Table presents an ablation study comparing the performance of Adaptive Transformer Programs using the Smooth Transition Mechanism against models using fixed Gumbel-Softmax (no transition, a(r) = 0 effectively throughout training) and fixed Sparsemax (no transition, a(T) = effectively throughout training) optimization strategies the Most-Freg task.
Analysis of Ablation Study:
Adaptive (Smooth Transition) outperforms both fixed strategies in accuracy and significantly reduces program length. This highlights the benefit of dynamically balancing exploration and exploitation during training.
17

---

Published as a conference paper at ICLR 2025
Fixed Sparsemax achieves higher sparsity rates than fixed Gumbel-Softmax, but at the cost of accuracy. While Sparsemax promotes sparsity, using it throughout training without the explo- ration phase of Gumbel-Softman appears to hinder the discovery of optimal solutions, leading to lower accuracy.
The Smooth Transition Mechanism achieves the best balance, attaining high accuracy while also inducing significant sparsity, as reflected the lowest program length and highest sparsity rate.
B.2
UNCERTAINTY-AWA. ATTENTION
This section provides formal description of the reference attention, details of the JSD-based gating mechanism, and ablation studies comparing different attention strategies.
B.2.1 FORMAL DESCRIPTION OF REFERENCE ATTENTION Are/
The reference attention distribution Aref learned distribution that dynamically adapts during training to represent the expected categorical attention distribution. It is initialized as uniform distribution and annealed towards a one-hot distribution over training.
Initialization: A Are) uni dist, uniform distribution over keys.
Update Rule Training Step
A a (f) uni_ + (1 a ()) oot_hot.dist"
where one hot ( the one-hot distribution corresponding to the categorical attention A at step and (1) selfitmp(()
is the annealing parameter, with self. self.temp decreasing
3.0
from 3.0 to 0.01 over training.
B.2.2 JSD-BASED GATING MECHANISM DETAILS
The gating mechanism g Uncertainty-Awe Aware Attention is based on the Jensen gence (JSD) between the categorical attention Acat, and the reference attention Aref,i-
Jensen-Shannon Shannon Diver-
Jensen-Shannon Divergence (JSD):
JSD(PIIQ) 1 KL(P)MM KL(QIM)
where M (P+Q) and KT(PIIQ) = E; P. P.log eh is the Kullhack-Leibler Divergence. In our case, P A and Q Arefii
*
Gating Weight Calculation:
g MLP(JSD( ((Arer,))))))
where MLP is small multi-layer perceptron network that maps the JSD value to gating weight E [0, 1].
B.2.3 SCORE-BASER BASED VS. CATEGORICAL ATTENTION ABLATION
To evaluate the effectiveness of uncertainty-aware attention (UAA) mechanism, we compare it against purely categorical attention (Cat, baseline TP) and purely score-based attention (Score) approaches across diverse tasks. Table presents the results, detailing accuracy (%) and lines of code (LoC) for each method.
Our UAA approach consistently outperforms or matches the baseline categorical method in aceu- racy while significantly reducing LoC across most tasks. Compared to score-based attention, UAA achieves higher accuracy (e.g.. 91.81 % vs. 79.68% for Double Histogram) and lower LoC (e.g.. 513 vs. 927), demonstrating the efficacy of JSD-based gating. Notably, score-based attention increases LoC in complex tasks like Double Histogram (927) and Dyck-2 (1568) due to its dense, continuous logic, whereas UAA adapts by selecting categorical rules when uncertainty is low, yielding concise programs (e.g., 513 and 1065). This ablation highlights the novelty of our hybrid mechanism in balancing performance and interpretability.
18

---

Published as a conference paper at ICLR 2025
Table Ablation Study: Categorical (Cat), Score-Basee (Score), and Uncertainty-Aware Attention
(UAA)
Task
Cat (Baseline TP)
Score
UAA (Ours)
Acc (%)
LoC
Acc (%)
LoC
Acc (%)
LoC
Induction
100.00
107
100.00
105
100.00
101
Reverse
99.74
859
100.00
808
99.99
779
Histogram
99.94
199
99.95
185
99.95
189
Double Hist.
66.78
586
79.68
927
91.81
513
Sort
99.98
945
99.96
929
99.86
895
Most-Freq
76.44
1334
74.15
962
80.80
894
Dyck-1
99.69
1297
99.52
1289
99.93
1086
Dyck-2
97.98
1316
98.49
1568
98.14
1065
C EXPERIMENTAL DETAILS
C.I HYPERPARAMETERS AND IMPLEMENTATION
This section provides detailed information regarding the hyperparameters used, model configura- tions for different tasks, and training protocols.
C.1.1 OPTIMIZATION SETTINGS
Table summarizes the key hyperparameters used for training Adaptive Transformer Programs across all experiments, based on grid search and initial experiments on RASP tasks.
Table 6: Global Hyperparameter Settings
Hyperparameter
Value
Optimizer
Adam
Learning Rate
0.05
Learning Rate Scheduler
None (Geometric Annealing)
Weight Decay
0.01
Batch Size
512
Epochs
250
Gumbel Temperature Annealing
Geometrie, 3.0 to 0.01
Gumbel Samples
per step
Gradient Clipping Norm
1.0
Hyperparameter Search: We conducted a grid search over the number of layers (2, 3), number of attention heads (4, 8), and number of MLPs per layer (2, 4) for Adaptive Transformer Programs on RASP tasks to optimize performance. Attention heads and MLPs were evenly divided between categorical and numerical types, where applicable. The reported hyperparameters and model con- figurations reflect the best performing settings identified through this search. For standard Trans- formers (baseline), we followed a similar grid search for layers (2, 3) and attention heads (4, 8) but otherwise largely adopted the hyperparameters described by (Weiss et al. (021), including hidden dimension of 256, learning rate of 0.0003, batch size 50, and training for up to 100 epochs.
C.2 DATASET DETAILS
This section provides details about the datasets used our experiments, addressing the nature of "small-scale" datasets and providing relevant statisties.
19

---

Published as a conference paper at ICLR 2025
C.2.1 DEFINITION OF ""MALL-SCALLE DATASETS
In the context of algorithmic reasoning and Transformer Programs, "small-scale" datasets refer to datasets designed for focused analysis of model capabilities on specific algorithmic tasks. These datasets are characterized by:
Controlled Sequence Lengths: Sequence lengths are typically limited (e.g.. to 16 for Dyck languages, for other RASP tasks, 32-64 for NLP tasks) to manage computational complexity and enable detailed program extraction and analysis.
Vocabularies Matching Sequence Lengths (for RASP tasks): For RASP tasks, vocab- ularies are often kept small and matched to the sequence lengths isolate algorithmic reasoning without relying large-scale pretraining or extensive vocabulary knowledge.
Focus on Algorithmic Reasoning:
The datasets are specifically designed to evaluate the model's ability to learn and execute particular algorithms or reasoning patterns (e-g.. re- versing sequences, counting tokens, checking Dyck language validity, named entity recog- nition, text classification).
"Smmallscale" in this context is *not* indicative of limited complexity the algorithmic tasks them- selves, but rather design choice to enable focused interpretability analysis in controlled experi- mental setting.
C.2.2 DATASET STATISTICS AND PREPARATION
Table 7 provides key statistics for the datasets used in experiments.
Table 7: Dataset Statistics
Dataset
Task Type
Avg. Sequence Length
Vocabulary Size
Training Examples
Induction
RASP
8
8
10,000
Reverse
RASP
10,000
Histogram
RASP
8
8
10,000
Double hist
RASP
10,000
Sort
RASP
8
8
10,000
Most-Freq
RASP
10,000
Dyck-1
RASP
16
2
10,000
Dyck-2
RASP
10,000
CoNLL-2003
NER
20 (avg.)
10,000
14,985
TREC
Text Classification
10 (avg.)
10,000
5,452
MR
Text Classification
20 (avg.)
10,000
8,529
Subj
Text Classification
23 (avg.)
10,000
11,500
AG News
Text Classification
25 (avg.)
10,000
120,000
Dataset Preparation: For RASP tasks, datasets were generated following the specifications in
(Weiss et al., (021). GloVe embeddings (Pennington et al.] 2014) were used to initialize token embeddings for NLP tasks. No other specific preprocessing or data augmentation techniques were applied beyond standard tokenization and vocabulary building.
D FAILURE ANALYSIS AND FUTURE WORK
D. ERROR PATTERN ANALYSIS
This section analyzes common failure modes observed across different task categories, investigates program structures associated with these errors, and examines suboptimal attention mechanism choices in failure instances.
20

---

Published as a conference paper at ICLR 2025
D.1.1 COMMON FAILURE MODES ACROSS TASK CATEGORIES
Table categorizes common failure modes observed across RASP, NER, and Text Classification lasks, with example instances to illustrate each failure type.
Table 8: Common Failure Modes Across Task Categories
Task Category
Failure Mode
Example Instance
Incorrect Counting
Histogram: Miscounting tokens in sequences like "aabbe"
RASP
Logic Error Predicate
Dyck-2: Incorrectly classifying "(D]" valid
Positional Misunderstanding
Sort: Incorrectly ordering elements in longer sequences
Entity Type Confusion
Misclassifying "Apple Inc." LOC instead of ORG Incorrectly tagging multi-word entities (e.g., "New York City")
NER
Boundary Error
Text Classification
Misclassification of Nuance
Subj: Failing to identify subtle subjective language
RASP tasks, errors in counting often correlate with overly simplistic numerical attention modules predicate logic that fails to capture token frequency accurately. Logic errors in Dyck languages can be linked to predicate structures that do not fully capture the recursive nature of Dyck language grammar. Positional misunderstandings in Sort may arise from numerical attention modules not effectively leveraging positional encodings for long-range dependencies.
In NER, entity type confusions sometimes be traced back to categorical attention heads relying too heavily on local context and failing to integrate broader sentence-level information. Boundary errors often indicate limitations in the predicate logic responsible for identifying entity boundaries, potentially due to insufficient context window or overly strict predicate conditions.
D.2 TASK-SPECIFIC ANALYSIS
D.2.1 LIMITATIONS IN CURRENT TASK SCOPE
The current task scope primarily focuses on discriminative tasks (classification, tagging) and algo- rithmic reasoning tasks with deterministic outputs. Limitations of the current framework become apparent when considering generative tasks like machine translation or text summarization:
Decoding Mechanisms: Adaptive Transformer Programs, their current form, lack ex- plicit decoding mechanisms required for generating variable-length output sequences in generative tasks. Standard Transformer decoders with autoregressive generation capabili- ties would need to be integrated.
The current framework is primarily designed for tasks where output length either fixed predictable based on input length. Generative tasks require handling variable and potentially unbounded output lengths.
Handling Varithle-Length Outputs:
Training for Generation: Training paradigms for generative tasks (e.g., sequence-to- sequence training, reinforcement learning for generation) differ significantly from the lassifiiattonnfocusss training used for the current tasks. Adaptations to the training pro- cedure would be necessary.
Example Machine Translation Challenge: Consider translating "Hello world" to French. ATP would need to learn not just to classify the input but to $generate** the sequence "Bonjour le monde," word by word, while maintaining interpretability. This requires significant architectural and training modifications beyond the current framework.
D.2.2 ADAPTATION CHALLENGES FOR GENERATIVE TASKS
Adapting Adaptive Transformer Programs for generative tasks presents several key challenges:
Interpretable Decoding:
Designing interpretable decoding mechanisms that maintain the program-like transparency of ATP is a major challenge. Standard autoregressive decoders Transformers, while effective, can be complex and less directly interpretable.
21

---

Published as a conference paper at ICLR 2025
Maintaining Sparsity in Generation: Ensuring that the generation process itself remains sparse and interpretable, rather than becoming a complex black box decoder, crucial. Sparsity-inducing techniques would need to be extended to the decoding stage.
Balancing Generation Quality and Interpretability: Achieving high-quality genera- tion while preserving a high degree of interpretability may involve trade-offs. Research is needed to find optimal balances between these competing objectives.
D.3 FUTURE RESEARCH DIRECTIONS
This section outlines future research directions address the identified limitations, expand the task scope, and explore alternative uncertainty quantification approaches.
D.3.1
ADDRESSING IDENTIFIED LIMITATIONS
research agenda to address the limitations identified in our failure analysis and task-specific anal- ysis includes:
Enhanced Numerical Reasoning: Further improve numerical reasoning capabilities by exploring more sophisticated positional encoding schemes and numerical attention mecha- nisms. Investigate methods to reduce errors in complex counting and numerical manipula- tion tasks.
Contextual Predicate Learning: Develop techniques to enable categorical predicates to better leverage broader context and sentence-level information, particularly for tasks like NER where long-range dependencies important. Explore attention mechanisms that explicitly model contextual scope.
Improved Gating Mechanisms:
Investigate refinements to the JSD-based gating mecha- nism in uncertainty-awant attention. Explore alternative uncertainty measures and gating network architectures to optimize the dynamic blending categorical and score-based at- tention for different tasks and input contexts.
D.3.2 ARCHITECTURAL ADAPTATIONS FOR BROADER TASK COVERAGE
To extend Adaptive Transformer Programs to broader task coverage, especially generative tasks, key architectural adaptations include:
Integration of Interpretable Decoders: Develop or adapt interpretable decoding mech- anisms compatible with the ATP framework. Explore constrained decoding methods or program-guided generation approaches that maintain transparency.
Handling Variable-Length Outputs:
Design program structures and training procedures that can effectively handle variable-length output sequences, allowing ATP to be applied to tasks like machine translation and text summarization.
Modular Generative Program Components: Investigate the development of modular, in- terpretable program components specifically designed for generative tasks, such program modules for sequence generation, copy mechanisms, or hierarchical output construction.
D.3.3 ALTERNATIVE UNCERTAINTY QUANTIFICATION APPROACHES
Exploring alternative uncertainty quantification approaches beyond Jensen-Shannn Divergence could further enhance the robustness and adaptability of Uncertainty- Attention. Potential directions include:
Bayesian Uncertainty Estimation: Integrate Bayesian methods (e.g., Monte Carlo Dropout, Bayesian Neural Networks) to obtain more principled uncertainty estimates for categorical attention distributions.
Entropy-based Measures: Investigate using entropy or other coarmation-theorelee mea- sures to quantify uncertainty in categorical attention as altermative JSD. For example, using the Shannon Entropy H(P) = Σ, log P. of the categorical attention distribu- tion Acat.i-
22

---

Published as a conference paper at ICLR 2025
Learned Uncertainty Estimators: Explore training dedicated uncertainty estimator net- works that directly predict uncertainty measures based on input and intermediate represen- tations, potentially providing more task-specific and adaptive uncertainty quantification.
D.3.4 POTENTIAL INTEGRATION WITH OTHER INTERPRETABLE MODELS
Future research could explore integrating Adaptive Transformer Programs with other interpretable model architectures, such as:
Neuro-Symbolic Systems: Combine ATP with neuro-symbolic approaches to create hy- brid systems that leverage the strengths of both neural and symbolic reasoning for enhanced interpretability and robustness.
Program Synthesis Techniques: Explore tighter integration with program synthesis tech- niques to further constrain and guide the learning of interpretable programs, potentially using program synthesis to initialize or refine extracted ATP programs.
Contrastive Learning for Interpretability: Investigate contrastive learning frameworks to explicitly train ATP models to produce more disentangled and interpretable representa- tions, potentially by contrasting program representations for similar and dissimilar inputs.
These future research directions aim to push the boundaries of interpretable Ay addressing current limitations, expanding the applicability of Adaptive Transformer Programs to wider range of tasks, and exploring novel approaches to uncertainty quantification and model integration.
23


Reviews:
[{"review_id": "MskdUjU442", "reviewer": "Reviewer_U9Xh", "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. \n\n1. they replace Gumbel-Softmax with a smooth transition between Gumbel-Softmax and Sparsemax, a sparse version of softmax. \n2. 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. \n3. They add position embeddings.\n\nThe 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).", "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. \n\n2. 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)\n\n3. What is A_ref? it is \"introduced\" but I don't see how it is calculated during training.\n4.  \"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.\n\n5. \"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. \n\n6. 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?\n\n7. 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? \n\n8. 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. \n\n9. how can lines of code differ between trained models?\n\n# Suggestions\n1. 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.\n2. some typos in citations — line 109\n3. 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\" ?\n4. 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)\n5. typos in line 235-236.\n6. 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.\n7.line 310 \"we tested our models on small-scale datasets\" — what does small-scale mean? reference appendix for more details?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\nQuality 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). \n\nClarity: 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.\n\nAnother main weakness is that there is no appendix giving further details. \n\nI 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "x52OxW70Jj", "reviewer": "Reviewer_EZgs", "summary": "This paper introduces Adaptive Transformer Programs, an enhanced framework that improves upon existing Transformer Programs to create more interpretable AI models.\n\nSpecifically, 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.\n\nThen 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.", "questions": "(1) Why can reducing the \"lines of code\" be directly considered as improving interpretability?  \nTo 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.  \n\n(2) Regarding sparsity, how does the proposed method behave better than Gumbel-Softmax?  \nIn 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.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n- The experimental results are solid and the proposed method clearly surpasses the Transformers Programs baseline in both accuracy and \"lines of code.\"\n- The experimental setting is comprehensive, consisting of both RASP and NLP tasks.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2N5g1DfohR", "reviewer": "Reviewer_ewyT", "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.", "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?\n2.  In Table 2, accuracy appears to decrease after integrating the Smooth Transition Mechanism. How should this be interpreted?    \n3.  How does the proposed method perform on generative tasks beyond NER and text classification?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\nThese 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.\n\n**Minor issues:**\n\n1.  In line 96, the MHA computation omits the normalization factor $\\frac{1}{\\sqrt{n}}$, where $n$ is the hidden dimension.\n2.  In line 99, the double quotes before “Transformer” need correction.\n3.  In line 109, please insert a space between \"RASP\" and \"Weiss et al.\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "J5v1gxUGVp", "reviewer": "Reviewer_J1B6", "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). \n\nThe main architectural modification proposed by the paper are:\n1. Providing a method to smoothly switch between Gumbel-softmax and sparsemax as a method to control switching between exploration and exploitation.\n2. A method to weight between categorical attention and score-based attention\n3. A position-aware attention which allows numerical Value, while keeping Q, K categorical.\n\nThe 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.", "questions": "- Could more details be provided about the hyperparameters and optimizers used for the various experiments?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper considers an important problem: that of making intrinsically interpretable models more expressive.\n- The architectural modifications are well-motivated and largely well described in the paper.\n- Experimental results across a wide range of tasks provide evidence for the performance of the model.\n- A thorough ablation study is provided.\n- 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.\n\n- 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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)