Emergent Discrete Controller Modules for Symbolic Planning in Transformers
Authors: S M Rafiuddin
Muntaha Nujat Khan
Venue: ICLR 2026 Poster (2026)
Keywords: Transformers
symbolic planning
discrete controller modules
length generalization
algorithmic reasoning
Primary area: neurosymbolic & hybrid AI systems (physics-informed, logic & formal reasoning, etc.)
TL;DR: We equip Transformers with discrete controller modules for ASSIGN/ADD/COMPARE/BRANCH, proving bounded-step expressivity and showing strong length generalization with interpretable execution at minimal compute cost.

Abstract: Transformers struggle with tasks that require symbolic planning loops, variable updates, and conditional branching, especially under length extrapolation. We introduce discrete controller modules that insert a small set of program primitives (ASSIGN, ADD, COMPARE, BRANCH) into Transformer blocks via a Gumbel–Softmax selector over operations and a compact program state of registers, flags, and optional memory. We prove that the augmented model can simulate any bounded-step program by mapping each primitive step to one controller step, and we bound the deviation of relaxed execution from its discrete trace by $O(\tau+\kappa^{-1})$ (selection temperature $\tau$, comparison sharpness $\kappa$). Empirically, the controller-augmented Transformer achieves strong length generalization on algorithmic benchmarks (Sorting, Sum-of-List, BFS), improving longest-length accuracy by up to $20$–$40$ points over strong baselines, and yields consistent gains on symbolic QA (DROP) and program-synthesis-style tasks (RobustFill) with reduced compositionality drop-off. The learned execution is interpretable: operation traces align with ground truth, register roles are linearly decodable, and targeted knockouts cause localized accuracy losses. The approach adds only $\sim$5–7% FLOPs and can be applied sparsely (every $p$-th layer).

Published as a conference paper at ICLR 2026
EMERGENT DISCRETE CONTROLLER MODULES FOR SYMBOLIC PLANNING IN TRANSFORMERS
S M Rafiuddin
Muntaha Nujat Khan
Oklahoma State University sraf sraf Pafiud@okstate.e edu
Oklahoma State University
munkhan@okstate. edu
ABSTRACT
Transformers struggle with tasks that require symbolic planning loops, variable updates, and conditional branching, especially under length extrapolation. We in- troduce discrete controller modules that insert a small of program primitives (ASSIGN, ADD, COMPARE, BRANCH) into Transformer blocks via a Gunbel-Softmax selector over operations and compact program state of registers, flags, and op- tional memory. We prove that the augmented model can simulate any bounded- step program by mapping each primitive step to one controller step, and we bound the deviation of relaxed execution from its discrete trace by + (selection temperature т, comparison sharpness K)). Empirically, the controller-augmented Transformer achieves strong length generalization on algorithmic benchmarks (Sorting, Sum-of-List, BFS), improving longest-length accuracy by up to 20-40 points strong baselines, and yields consistent gains symbolic QA (DROP) and program-synthesis
tasks (RobustFill) with reduced compositionality drop-off. The learned execution interpretable: operation traces align with ground truth, register roles are linearly decodable, and targeted knockouts cause localized accuracy losses. The approach adds only ~5-7% FLOPs and can be applied sparsely (every p-th layer).
INTRODUCTION
Transformers (Vaswani et al. 2017) have revolutionized sequence modeling across natural language processing and beyond, yet they lack explicit mechanisms for symbolic planning constructs such as loops, variable assignments, and conditional branching. Concurrently, neuro-symbolic approaches have sought to integrate symbolic reasoning with deep learning, often by coupling external planners or eemorryaaggmnnt modules (Kaiser & Sutskever 2016 Graves et al.] 2016). However, these methods either treat symbolic reasoning post- prote-rrcessiig step or incur significant architectural complexity.
Recent work algorithmic reasoning benchmarks highlights this limitation. The CLRS bench- mark (Veličković et al., 2022) and ListOps tasks (Nangia & Bowman 2018) reveal that standard Transformers struggle generalize to longer problem instances, underscoring the need for induc- tive biases tailored to symbolic operations. Efforts to retrofit Transformers with neural modules, such Universal Transformers with adaptive computation time (Dehghani
2018)) or the Com- pressive Transformer (Rae et al., 2020), introduce dynamic depth but still lack discrete control flow semantics. Unlike adaptive-depth controllers that modulate computation time, our goal to endow layers with verifiable control-flow semantics (loops, branches, assignments) and produce step-wise- human-readable execution traces. Unlike external planners or generic memory modules, we main- tain program state internally rather than delegating reasoning to post-processing large external memorics.
Differentiable discrete operations have been enabled by advances like the Gumbel-SSftmax trick (aan et al., 2017 Maddison et al. 2017) and its continuous relaxations, yet their integration into large-scale architectures remains underexplored. Furthermore, mechanistic probing studies (Geval ct al. 2020) suggest that Transformers can incidentally learn symbolic behaviors, but without dedi- cated modules, such behaviors neither reliable interpretable.

---

Published as a conference paper at ICLR 2026
In this work, embed explicit, learnable discrete controller inside Transformer layers that selects from a finite opcode set with verifiable control-flow semantics, distinct from depth controllers (e.g., ACT) generic external memory. Our contributions are: (1) Gumbel- Softmax-parameterized controller that executes ASSIGN, ADD, COMPARE, and BRANCH with step-wise, human-readable traces; (2) tight integration with attention and feed-forward computations that maintains a program state vector without external memory; (3) formal expressivity result show- ing that O(K+B) instrumented steps (or total depth p(K+B) if every p-th block instrumented) suffice to emulate any program in bounded imperative class PK.B: straight- line code with condi- tionals and loops whose bodies are unrolled at most iterations, with total primitive steps per execution at most K+B;; approximation error can be made arbitrarily small under standard Gumbel- Softmax temperature annealing: and (4) empirical gains on algorithmic reasoning tasks (e.g., sorting, graph traversal) and symbolic QA benchmarks (e.g., DROP), including superior length generaliza- tion (c.g., 2-4x on tested tasks) ) and improved interpretability, supported by ablations isolating controller discreteness and opcode design.
2 BACKGROUND
2.1 TRANSFORMER LAYERS
The Transformer architecture introduced by Vaswani et al. (Vaswani al. 2017) models sequence of I" tokens (xI,..., IT by first mapping each token to D-dimensional embedding and adding positional encodings. Let X g²'x denote the resulting input matrix. Transformer block consists of two sub-layers: multi-head self-attention and position-wist wise feed-forward network. In the self attention sub-layer, query, key, and value projections are computed as-
Q XWQ
K = XWW XW
where Wº, W WV E DND and D₁ = D/H for heads. Attention weights are given by-
Attention (aanttinnQQ.K. Q. V) softmax
QK )
vDr
and the outputs of all heads are concatenated and projected back R.P. The heddfforwrrd sub-layer applies two linear transformations with ReLU-
FFN(Z) max(0, ZW; + b,) b2
where W1 E DDDD, W2 E DDrxD Each sub-layer is wrapped with residual connections and layer normalization-
Z' LayerNormmZ + SubTayer ())
Stacking L such blocks yields the full Transformer encoder.
2.2 GUMBEL-SOFTMAX TRICK
Training models with discrete choices poses a challenge for gradient-based optimization. The Gumbel-Softmax trick (Jang et al.] 2017] Maddison al. 2017) provides a continuous, differ- entiable approximation to sampling from a categorical distribution. Given unnormalized log- probabilities (logits) (log a; 1. draws i.i.d. Gumbel noise 9; log( log u;) with uq Uniform(O, 1) and computes-
exp((log a; + 9:)/T)
D/
j=1 exp (log aj + 9;)/T)
where 7 > a temperature parameter. T 0, approaches a one-hot sample; for larger т, y is smooth approximation that permits backpropagation. During training, may anneal r to encourage discreteness while preserving gradient flow. This mechanism enables end-to-end learning of modules that make hard categorical decisions within neural architectures.
2

---

Published as a conference paper at ICLR 2026
O STATE TOKM
INPUT
umbeed
oprerentatian eqjecion 9)
[M MULTI-H540
[ FEED
STATE 190)
<a>
spore) ell toloons
stade Trace
Begiiters C_N
GoOptional bay-raluu alots
Baga
frome
(scecuted every pith Traneformer layer)
[8.43 OPERATION SELECTOR) (Garsbe) Suronsa ove
Engpaid for
contrstor appeara
(9.0) repfeensteneerry saka enitter
ARGS]
pith
BACK)
(Dran)))
[S OPCOGE OPMANTICNS
COMPARE: ete
SUI RJI
(4- BRANCH GATE)
(C0 TASK HEAD
B TRACE LOGGES)
1 INTERPRETABLLITT Inocosat tult)
decidors (ertropy, brazchet)
B BBY TRAINING LOGGES & SCHEOULES)
Task CE/ALL
(optional traces)
Dibtt updabd)
& aarroulig
(udrrants/vellie))
OUTFUTS AUKILIARY
LAYER STACK INSTRUMENTATION TION (every p-th) Layer [2-1] (6-e) [a] [costralier active)
Layer 4 [2-5] + [6-9] + [8] (controlier active)
Layer L: (hapeat: chosse p:2 or control overhead -5-10%)
Controller-augmeentel Transformer overview. Tokens flow through while a compact program state (registers, optional memory, flags) inter- acts with the token stream via statet>token projections. discrete controller selects one of [ASSIGN, ADD, COMPARE, BRANCH, NOOP) and updates the state; the resulting change is written back into the residual stream. We instrument every p-th layer (multi-step execution across depth) and log traces (opcode, addresses, gate) only at instrumented layers; overhead is =557% FLOPs for 2=2. Training uses task loss with optional trace/entroopy/sematii regularizers.
Figure 1:
METHOD
3.1 CONTROLLER MODULE DESIGN
We augment each Transformer block with a small discrete controller C that selects and executes symbolic operations latent program state. Let ( E RTxD be the token representations entering block and the program state at step be-
1 (R MO,
(6)
where (d) E I N,.x, are N, register vectors, ( E N.. is an optional key-value memory, and (()) E {0,1) are binary flags. The controller reads a summary s(l) = d(H()) E Rd, via cermmutation-ivaaiian pooling (mean/max or learned attention) and emits (i) operation choice, (ii) its arguments (addresses and scalars), and (iii) an update to
Operation choice. Let the operation be O {ASSIGN, ADD, COMPARE, BRANCH, NOOP). The controller computes logits (!)) W, (!)) + bo E RRO01 and draws differentiable one-hot via

---

Published as a conference paper at ICLR 2026
Humbel-Softmase
n + g
(() softmax 8i
log( - log ux): ug Uniform(O.)
Temperature 7 is annealed during training (Sec. 3.4)
Addressing and arguments. The controller predicts read and write distributions over regis- ters and memory, and value vector: at (t,tmax(W,,s))) ((0) softmax WW ~ s(),,
v()) = W, $0, Reads and writes are differentiable mixtures: r()) = EE R R
EES a()) ( (i)
Differentiable operation semantics. Each operation induces deterministic map on (R, M, f).
mixed by y().
1(111) Σ No Sem r()) v())
(9)
oEO
We instantiate the primitives as follows (all elementwise where appropriate)
Sem
R (1+1) = (1 - X) R (1) + A B(a")) (()
(10)
Semwoo R (1+1) RR () + E(B()) ()) (()T
(11)
SeMCOHPARE
p((+)) -(s [AR BR BRR1) 11)
(12)
SemERWNc: 8(1+) =7 v (thel((()) )++ + - (0) v ()))
(13)
Semwoop
s((+1) S(()
(14)
Here E(G) RRNxx1 is the one-hot expected write address (soft in practice), (0, a write strength, a logistic with sharpness >> 1, A, B select two (soft) registers for comparison, and 70)) = (w, Tp()) + by) gates the branch. The branch subprograms Vehon, Velse are tiny MI.s sharing parameters across layers; they transform but do touch H.
3.2 INTEGRATION INTO TRANSFORMER
Let standard block be Block(.) = FFN MHA(.). We interleave controller updates with the residual pathway while exposing (()) to the token stream.
ii State-to-tokem:
b (!) A (S
E RR, A H +1 ( Block (H
iii Token update:
(iii) Token-to-state:
( oZ((), s(+1)) = C(S")
(iv) Residual write-back:
(++1) (1) + ((S(()))))))))) s(D) Wa,
where is learned state encoder, ø is pooling (mean/max attention with learned query), W. aligns state changes to token space, and P selects the register(s) whose changes should influence H. Practically, instrument every p-th block (e.g., p {1,2,3)), which controls compute while permitting multi-step programs across depth.
3.3 THEORETICAL EXPRESSIVITY ANALYSIS
We formalize fragment of imperative programs and show our architecture can simulate exactly in the limits --> 0) and к - cc. Notation. We use Luet for the number of Transformer layers/blocks (cf. Table 4), and L,p for the execution-trace length (program steps). In this section, bounds are stated in terms of K B and Ln < K +B; when relating depth we write Liet. Program class. Consider straight-line programs with conditionals and bounded loops over N, real-valued registers with primitives {ASSIGN, ADD, COMPARE, BRANCH). where COMPARE sets Boolean flags consumed by BRANCH, and loops are unrolled to at most B iterations. For completeness, the model's operation set includes NOOP, but the expressivity proof does not require it.
4

---

Published as a conference paper at ICLR 2026
Theorem (Expressivity for bounded imperative programs). For any program P in the class above with at most K primitive steps and loop bound B. there exist a depth Luet = O(K+B) (or to- tal depth pBK+B) if instrumenting every p-th block), widths (D, dr dm), and parameters of controller-augmentte Transformer such that, for input (X, S(U)), the sequence of states { J1 produced by the network matches the execution trace of exactly as T - Co. Co t and K 7 0c.
Here, Cg and Ca are scale factors applied to the operation and address logits, respectively. See Appendix[]jes the complete proof.
Error under finite temperatures. For finite scales, let ma be the minimum opeode- ma the minimum address- margen,, and my the minimum absolute predicate margin across steps. Then per step-
1 1 P(o: oî] - (10) 1)e G
(15)
|100 ey. 11300) -e-- 11 ma
(16)
(o(2) H()) se -
(17)
and the straight-through Gumbel-Softme smoothing contributes an additional ((r). union bound over L steps yields total deviation bounded by O(Le-como) (omm))+O(Le +O(Le me + km))) O(LT), which vanishes in the joint limit.
3.4 TRAINING OBJECTIVES AND LOSSES
Let 0 denote all parameters. The total loss is-
[(0) Luss Aopoop Asem + NeatCent + Cppar
(18)
Primary task loss. Lunk is cross-entropy (classification) or negative log likelihoodd (seq2seq) on model outputs. Operation supervision (optional). When operation traces are available (synthetic tasks), we include Cop & Σ/= CE(y""), (E()")),=")))) where y*()) is the ground primitive at step l. Semantic consistency. We penalize mismatch between predicted state transitions and their declara-
|((t)) -. you Sem,(S(0, ,r)0, v()) ||2
tive semantics: Erem Σio
On tasks with known ((T
invariants (e.g., sum conservation), we add constraint penalties, e.g. ||1 T R 2 1RB0 for ADD- free steps. Entropy annealing. To encourage discrete decisions we minimize controller entropy: frett 1, Et= Hy) r(t)= max(Fmin To e-))) with training step initial temperature To floor Tmin- and decay 7. Write sparsity. We regularize address distributions and write magnitudes: Cypar Ele (1 + (1v(() [13), which promotes near one-hot hot writes and small updates.
Optimization details. We use straight-througe estimation for hard branches at late training de- sired: during the forward pass yl one one_noooe [oog max x")); during the backward pass we substi- tute the Gumbel-Softman gradient. We jointly train A by AdamW with gradient clipping; auxiliary losses are ramped in over the first K Kwarmup steps to stabilize early learning.
Complexity. Each controller adds O(Dd, + N,d,) parameters and O(TDd, + N,d,)) FLOPs per instrumented block; with controllers inserted every p-th block and de, d, < D, the relative FLOP overhead to D-width Transformer is settings.
1 dd (see App. G.1 typically 5-7% for p=2 in our
2p
4 EXPERIMENTS
4.1 DATASETS AND TASKS
Synthetic algorithmic suite. We construct three controlled tasks to probe algorithmic generaliza- tion: (i) Sorting, (ii) Sumsof-List, and (iii) Graph Traversal. For Sorting and Sum-of-List, sequences are length-n lists of integers sampled i..dd from (1, 999} with duplicates allowed; inputs are se- rialized tokens 
 ; In 
 and outputs either the sorted sequence single integer
5

---

Published as a conference paper at ICLR 2026
token (the sum). We train n E {10, 20, 40, 80} and evaluate both in-distribution (ID) and length- extrapolation at n E {160, 320). For Graph Traversal we sample G(n, Erdös-Rényi graphs {10, 20, 40, 80} nodes, encode adjaceney lists and a start node, and require the model to emit a breadth-first traversal order (ties broken by node id). We generate 200k/20k/20k train/val/test examples for Sorting and Sum-of-List, and 120k/1 5k/15k for Graph Traversal (cf.
(p 0.2) with
Sec. 4.3 for generation details). Program synthesis and math reasoning. We include a string- to-program task derived from Robus(Fill-style benchmarks (Devlin et al. 2017), where the input is set of input/output string pairs and the output is a DSL program that satisfies them; following (Devlin et al.] 2017) we report exact match on held-out examples. We also use the Mathematics Dataset (Saxton et al.) 2019) (algebra/arithmetic subsets) to test multi-step numeric reasoning with compositional templates. Symbolic question answering. We evaluate on DROP (Dua al. 2019), which requires discrete operations (addition, counting, comparison) over paragraphs. Following prior work, we report Fl and Exact Match (EM). We construct a numeric-only subset (questions whose official program trace contains numeric operators) to target algorithmic reasoning.
4.2 BASELINES
We compare against strong families representative of the literature. (B1) Transformer: a standard
encoder-decoder (or encoder-only for classification) without controllers (Vaswani et al.
2017).
(B2)
Universal Transformer: adaptive-depth variant with dynamic halting (Dehghani et al.
2018
(B3)
Compressive Transformer: long-context memory via compressed history (Rae et al.
2020
(B4)
Transformer: an implementation following (Veličković et al.
2022) when
applicable. (B5) NPI/DNC: Neural Programmer Inoerreeee-Inteppreeeer and Differentiable Neural Computer
baselines (Kaiser& & Sutskever 22066; Graves et al.] 2016). (B6) Soft Modules: our architecture with the controller's 's operation selection replaced by temperature-I soft mixture (no discrete selection). We additionally include a privileged comparator, (B7) External Planner', which injects intermediate algorithmic hints into the token stream. Because B7 has access to supervision not provided to other models or ours, we report it separately and exclude it from any "vanilla baseline" aggregates and significance tests. For DROP, we additionally include (B8) FiD-Ts" (Fusion-in Frusion-in-Deooeer Deeoder) as a strong reference system. All baselines are matched for depth/width as closely as possible; when memory-oriented models require extra parameters (e.g., Compressive), we report both accuracy and compute (FLOPs) to contextualize results.
4.3 IMPLEMENTATION DETAILS
Model sizing. Unless otherwise stated, we use L= 12 layers, hidden size D 512, = heads, FFN width D,y = 2048, Controllers are inserted every p-th block with p = 2 (i.c., in 6 blocks).
Program-state registers: N, = 8 with d, = 64; optional key-value memory N,, = 64, dm = 64. Controller summary dimension d, = 128. Training. We train with AdamW (31 0.9, 82 = 0.98, weight decay 0.01), linear warm-up for 5k steps then cosine decay; base learning rate x 10-4 (seq2seq) 1 x 10-4 (classification). Batch size B 256 tokens-equivalent (total tokens per op- timizer step) for synthetic tasks and B = 128 for DROP/Robustill.. We apply label smoothing (e= 0.1) and gradient clipping at 1.0. The Gumbel temperature follows 7(t) max (Tmin The ")) with (To. Tmin, 7) = (1.0,0.1,5 5 x 10-5). Entropy penalty Aem is linearly ramped from 0 to 5 10-3 over the first 20k steps. Unless noted, we train for 50 epochs (synthetic) and 20 epochs (DROP/RobustFill) with early stopping on validation. Generation and tokenization. Synthetic corpora are re-generated with fixed seeds for each run; integer tokens are subword-free (<dNNN> per integer), and graphs are serialized as edge lists with delimiters. For RobustFill we use the public DSL and the standard tokenization of (Devlin et al. 2017). For DROP, we use the official prepro- cessing. Compute and reproducibility. Experiments run single modern GPU (40GB device). For algorithmic tasks we use x seeds, and for non-synthetic tasks we use model seeds; we report ttdd
FiD- -5 (Fusion-in in- Decoder) relies on document retrieval and generative scoring. Our synthetic evalua- tions are exact sequence-accuracy tasks without retrieval, FiD-T5 is omitted in Table T see Table 2a for DROP results.

---

Published as a conference paper at ICLR 2026
4.4 EVALUATION METRICS
Task accuracy. For Sorting and Graph Traversal, we report sequence acewacy:: a prediction is correct iff the entire output sequence matches the target. For Sum-of-List we report exact nu- meric match. For RobustFill we report program exact match on held-out I/O pairs. For DROP we report official FI and EM. Length generalization. We quantify out-of-distribution general- ization by the length gap (Peain:: Rees) Acc(N(st) Acc(Nuain). with Thaain {20, 40, 80) and Flest {160, 320}; we also report the area under the accuracy-lengee curve (AUL) over test lengths. Sample efficiency. We measure the minimal number of training examples (a) required to reach target accuracy a (e.g.. 95%) and the area under the data- terta-peffommnce curve (AUP) up to the full training set. Interpretability and controller usage. We report operation-selectio entropy +, E, H ()) branch confidence 7, Ei max, Yo and (when synthetic traces are available) trace alignment: Align t E { maxy(") y-(").. We include knockout tests that zero troller outputs at inference; the induced accuracy drop quantifies functional reliance on controllers. Compute cost. We report tokens/s and estimate FLOPs per example. Relative overhead is mea- sured FLOP FL.PPsadiee for our settings this is typically 5-10% (see Table 7 for full tables). All FL.OPsyaeite datasets follow their original licenses.
RESULTS
5.1 ALGORITHMIC REASONING PERFORMANCE
Setup. We evaluate Sorting, Sum-of-List, and Graph Traversal (BFS) at train lengths n E {10, 20, 40, 80} and test-time extrapolation n E E {160. 320) (Sec. 4.1b. All numbers are means with 95% bootstrup CIs over x runs (3 corpus seeds x model seeds). Baselines are (B1) Trans- former, (B2) Universal Transformer, (B4) CLRS-Alg. Transformer, and (B6) Soft Modules (our architecture with non-discrete operation selection). Our model is denoted Ctrl-Transformer Main findings. Table reports sequence-level exact accuracy. On ID lengths (n S 80) our model matches or slightly exceeds the best baseline. Crucially, on extrapolation (n 2 160) Ctrl-Transformer sustains high accuracy, while baselines degrade sharply. For Sorting at 320, B1 drops 35-45% whereas Crrl-Tranformer remains above 90%. BFS shows a similar pattern with a +18-25 point margin at large n. Length curves and significance. Accuracy-length curves (Fig- 5 show that our method's slope beyond the training regime is substantially flatter. For hypothesis Testing, we use Wilcoxon signed-rank tests on the paired runs (baseline vs. ours) at the longest length (n=320; also n=160 where reported), and apply Holm correction across baselines; results indicate statistically signifi- cant gains across all three tasks. Consistency with theory. We measure controller decision entropy 1 E, H (y) and find correlates with extrapolation accuracy (Pearson r -0.81 across runs): as the entropy anneals (Sec. (.4), teiperttepeerture mixing diminishes and results align with the discrete semanties predicted by Theorem
5.2 SYMBOLIC QA AND COMPOSITIONAL TASKS
DROP. Table 2a reports FI/EM on DROP full and the numeric-only subset. Cran-Tranffommer im- proves F1 by +6.8 on the numeric subset relative to parameter- matched Transformer (B1), and by +4.1 over Universal Transformer (B2). Program synthesis & mathematics. On RobustFill-style string-to-programm, we observe +6.1 point exact-match gain for compositional test sets where I0 pairs require multi-step rewrites TTablebb). On the Mathematics dataset (algebra/aritmmetic), Ctrl- Transformer is particularly strong on long-form addition and mixed-operator templates, with +4-7 absolute accuracy over B1.
5.3 ABLATION STUDIES
We conduct targeted ablations to isolate design choices. (A1) Discrete soft controllers. Replac-
ing Gumbel-Softmax with a temperature-1 mixture (B6) reduces extrapolation accuracy by 10-25 points 2 160 (Table [). confirming the necessity of discrete selection for control flow. (A2) Number of controller blocks. Varying the number of instrumented layers (every e (1,2, 3, 4)) shows compute-performmnce trade-off (Table Instrumenting every 2nd block attains > 95% of

---

Published as a conference paper at ICLR 2026
Table I: Algorithmic reasoning: sequence accuracy (%) length. Means with 95% bootstrap Cls over 3 = 9 runs (3 corpus x 3 model seeds); Wilcoxon signed-rank tests at the longest length with Holm correction. Ctrl-Transformer generalizes 2x-4x longer inputs with minimal degradation. Note: 'B7 (External Planner) injects intermediate hints and is privileged comparator reported separately; is excluded from non-privileged baseline aggregates and any significance tests. FiD-T5 relies on retrieval and generative scoring and is therefore omitted from these synthetic (non-retrieval) tasks.
Task / Model
n=20
n==80
11=160
n==320
Aten +00+)220
AUL
Sorting
B1 Transformer
99.1
8..4
6..2.119
38.7
46.7
0.69
B2 Universal Trf.
99.0
90.6
1.6
51.5112111
39.1
0.74
B3 Compressive Trf.
99.0++0.2
92.1+0.7
1.3
59.1+1.8
33.3
0.79
B4 CLRS-Alg. Trf.
99.1
93.8
7..
62.7+1.7
-31.1
0.80
B5 NPI/DNC
98.9
91.0=10
69.5
52.4+2.0
-88.6
0.75
B6 Soft Modules
99.0+-0.2
93.1
78.94:144
66.3+:19
26.8
0.81
Clrl-Transformer (ours)
99.1 0.2
97.4
94.1
91.6+0.9
5.8
0.95
Privileged oracle-aided
B7 External Planner!
09.
96.1=#0.6
90.2
85.0
11.1
0.92
Lis- BI Transformer
99.6+
88.9+1.1
63.7
41.2+22.2
17.7
0.70
Universal
99.6+0.2
92.8+0.8
7.44+1.4
56. 1.9
36.7
0.76
B3 Compressive Trf
99.6+0.1
93.6+0.7
7..9++1.3
61.0+188
32.6
0.78
B4 CLRS-Alg. Trf.
99.620.1
94.2=0.7
80.7
68.311.6
-2.99
0.83
NPI/DNC
99.4
92.110.9
71.414
55.812.0
33.3
0.77
B6 Soft Modules
99.6+0.1
95.2+0.8
83.7+11.2
72.9+1.6
22.3
0.84
Surl-Transformer (ours)
99.6-
98.1+0.4
95.6+0.6
93.8+0.8
44.3
0.96
Privileged oracle-aided
B7 External Planner!
99.6
97.4
93.610.88
89.5
-7.9
0.93
Graph Traversal (BFS)
B1 Transformer
98.210.4
86.3111
7..41..5
57.9119
-28.4
0.78
B2 Universal Trí.
98.3+0.3
89.5+1.0
7..1
62.8+1.7
26.7
0.80
Compressive Trf.
98.4+0.3
90.8+0.9
78.4+1.4
66.7+188
24.11
0.82
B4 CLRS-AIg. Trf.
98.610.3
90.210.9
77.81113
65.1
-2..1
0.81
B5 NPI/DNC
98.510.4
90.71111
74.611..6
60.312.1
30.4
0.79
Soft Modules
98.610.3
92.110.9
83.511.11
74.211.
17.9
0.86
Strl-Transformer (ours)
0.3
95.3+0.6
90.8+0.8
83.5+1.1
-1..8
0.90
Privileged oracle-aided
B7 External Planner!
98.610.3
94.0
88.41100
81.2114
-12.8
0.89
Table 2: Crrl-Tranffrrmer across tasks. Left: DROP (FI/EM), Right: program synthesis (Ro- bustFill exact match) and Mathematics accuracy. Mean over 3 seeds.
(a) DROP results. std over seeds.
(b) Program synthesis math reasoning, Exact match (%) on RobustFill; accuracy (%) on Mathemat- ics subsets.
Model
EM
B1 Transformer
81.2+0044
78.5
Model
RobustFill Math (arith.)
B2 Universal Trf.
83.9+0.3
80.6+0.5
B8 FID-T5 (ref.)
87.4
84.110.4
B1 Transformer B6 Soft Modules
0.7
76.1
Ctrl-Tranffommrr (ours)
88.0+0.3
85.140.4
71.3++0.6
78.2+0.6
Trrl-Transfommrr (ours)
73.9+
82.0+0.5
Numeric-only subset
B1 Transformer
78.3++0.5
71.2-10.6
B2 Universal Trf.
80.9+-044
76.3
Transfforrr (ours)
85.1
81.510.5
the maximal gain at only = 60% of the overhead. (A3) Temperature schedule & entropy. Faster
annealing leads to earlier discretization but occasional training instability; slower schedules yield smooth learning but slightly worse final extrapolation. A middle-ground (70, Tmin: 7) from Sec. 4.3
8

---

Published as a conference paper at ICLR 2026
Table 3: Ablation: controller insertion frequency. Sorting accuracy (%) vs. length; overhead is relative FLOPs to B1.
Every p-th block
n=80
n=
n=320
Overhead
p=1
(all blocks)
98.0
95.1
92.7
+10.1%
p=2
97.4
94.1
91.6
+6.4%
p=3 3
96.1
91.8
88.3
+4.4%
p-4
95.0
89.4
84.9
+3.2%
maximizes AUL while keeping seed variance small. (A4) Semantic consistency loss. Removing frem decreases trace alignment by points and hurts Sorting (n 320) by -4.5 points, indicating that explicit semantic regularization stabilizes execution. (A5) Knockout tests. Zeroing controller outputs at inference (keeping the rest of the network identical) collapses Sorting@ 320 from 91.6 to 44.8 and DROP numeric Fl from 85.1 to 79.0, demonstrating functional reliance on controllers rather than superficial correlation.
5.4 ADDITIONAL ANALYSIS OF MAIN RESULTS
The key effect is gap amplification with length. A useful way to read ia to track how the advantage over the strongest von-privieeged baseline (B6 Soft Modules) grows as we move from the training regime (n S 80) to extrapolation (n > 160). On Sorting, the margin over Soft Modules increases from +4.3 at n- (97.4 4 vs. 93.1) (3.1)to to +5.2 at 160 (94.1 vs. 78.9) and further to +25.3 n=320 (91.6 vs. 66.3). Lfmsof-List exhibits the same monotone expansion: +2.9 (at n=80)
+11.9 (at n= - +20.9 (at ==320). BFS shows a smaller but still consistent expansion: +3.2 (at 8=80) - +7.3 (at n=160) - +9.3 (at 3=320). This pattern indicates that the improvement not primarily uniform upward shift; it emerges most strongly under the length shift where baseline generalization degrades.
Length extrapolation can be summarized by decay compression. Using the length gap column in Table [T] Ctrl-Transformee substantially compresses the 80-+220 performance decay relative to both Soft Modules and the vanilla Transformer. For Sorting, the absolute decay is 5.8 points (ours) versus 26.8 (Soft Modules) and 46.7 (Transformer), i.e., our decay is F 0.22 x that of Soft Modules and 0.12x that of Transformer. For Sum-of-List, the decay is 4.3 (ours) versus 22.3 (Soft Mod- ules) and 47.7 (Transformer), i.e., s 0.19x and Rs 0.09x, respectively. BFS remains harder: 11.8 (ours) versus 17.9 (Soft Modules) and 28 4 (Transformer), i.e., 0.66x and = 0.42x. Thus, the strongest gains correspond to tasks where the controller most sharply reduces how quickly accuracy deteriorates beyond the training lengths (Sorting and Sum-of-Liste while BFS shows more modest decay compression.
Interpreting the improvements as error-rate reduction highlights practical effect size. At the hardest extrapolation point (n=320), the absolute accuracy gaps in Table translate into large ductions sequence-level error probability. Relative to Soft Modules at n=320, Ctal-Transformer
reduces error by se 75% on Sorting (error 33.7% 8.4%) and 77% on Sum-of-Liss List (er- ror 27.1% + 6.2%), while BFS yields a smaller but non-trivial reduction of z 36% (error 25.8% + 16.5%). Relative to the vanilla Transformer, the reductions are even larger: R 86% (Sort- ing), s 89% LSuss,f-List), and ss 61% (BFS). This framing makes explicit that the controller's benefit long lengths is not only statistically meaningful but also qualitatively large in terms of avoiding complete failures in whole-sequence correctness.
AUL corroborates that the advantage is not confined a single length. Because AUL in Ta- ble 1 aggregates performance across the evaluated lengths, it helps verify that improvements are sustained throughout the extrapolation range rather than appearing only at 320 The AUL gaps between Tran-frrnfrrrmer and Soft Modules are sizeable for Sorting (0.95 vs. 0.81) and Sum-of- List (0.96 vs. 0.84), and smaller but consistent for BFS (0.90 vs. 0.86). Combined with the slope behavior from n=166 to n=320 (Sorting: -2.5 12.66 Sum-of-List: - 1.8 - -00.;; BFS: -7.3 vs. 9.3), AUL supports the same conclusion: the controller changes the shape of the length generalization curve by slowing degradation throughout the OOD region.

---

Published as a conference paper at ICLR 2026
The privileged comparator contextualizes what can be achieved without intermediate hints. [ repports B7 External Planner separately as a privileged model. At n=320, Cran-Tranffommr matches exceeds it on all three tasks (Sorting: 91.6 vs. 85.0; Sum-of-List: 93.8 vs. 89.5; BFS: 83.5 vs. 81.2), and also exhibits smaller 80-+320 decay on Sorting and Sum-of-List (-5.8 vs. -11.1; -4.3 -7.9), with comparable decay on BFS (-11.8 vs. 12.8). While B7 is not di- rectly comparable to non non-privileged baselines, this juxtaposition underscores that the non-privileged controller mechanism achieves strong extrapolation without relying on injected intermediate hints at test time.
The same "harder-regime" advantage appears on non-synthetic benchmarks. Tables 2a and 2b show that the gains extend beyond synthetic length extrapolation to tasks emphasizing discrete and compositional operations. On DROP (full), Cull-Transformer improves over the parameter-matched Transformer by +6.8 F1 and +6.6 EM (81 (81.2/78.5 88.0/85.1; Table 2ap. On the numeric-only subset, the gains are +6.8 F1 and +7.3 EM (78.3/74.3 - 85.1/81.5; Table 2a), indicating that improvements persist (and slightly strengthen for EM) when evaluation emphasizes numeric operations. For RobustFill and Mathematics (Table 2b). Ctrl-Transformer also improves
over both Transformer and Soft Modules (RobustFill: 73.9 vs. 67.8 and 71.3; Math (arith.): 82.0
vs. 76.1 1 and 78.2). Across these benchmarks, the consistent pattern is that the controller yields its largest advantages precisely where multi-step structured computation most stressed by the evalu- ation setting.
5.5 INTERPRETABILITY ANALYSIS
We analyze whether the learned controllers execute human-interpretable plans. Operation traces. For synthetic tasks with oracle step types (Sorting and BFS), we compute trace alignment against ground-truth opcode sequences emitted by a deterministic simulator (FIFO BFS with ascending- id tie breaking; comparison-and-wwe inputs or targets and are used only for alignment and, where stated, a small auxiliary loss during training. Ctrl-Tranfformer achieves 92.3% alignment on Sorting and 88.7% on BFS, vs. 61.5% for B6 (soft). Auxiliary trace loss (no leakage). On synthetic tasks only, we optionally add Auace =0.05 to the training objective (See. (.4); no trace labels are available or used at test time, and no traces are used for DROP or other natural- natural-language benchmarks. Decision entropy. Entropy of operation selection declines steadily during training (Fig. 2) matching the annealing schedule and supporting the finite-temeerature analysis: runs that end with lower entropy attain higher extrapolation. Regis- ter semantics. Linear probes reveal consistent roles: e.g., a "running min" register in Sorting, and fromtier/parent-ponneer registers in BFS. Randomly permuting registers at test time degrades per- formance by -9 points, indicating learned specialization rather than redundancy. Failure modes. When annealing is too slow, the controller blends branches inon-discrete behavior), causing off- by- one errors at long lengths; when too fast, early hardening can lock in suboptimal strategies; entropy penalties and frem mitigate both.
schedule for Sorting, These traces are never serialized into
Compute cost. Empirical overhead is ~ 5-7% FLOPs for p-2 instrumentation, and throughput (tokens/s) decreases by 6.8% on average (Table which is acceptable given the +25-50 point gains at long lengths. Across synthetic and real tasks, Ctrl- Tran-fraanfrreer delivers (i) state-of-the-art length extrapolation, (ii) robust gains on symbolic QA and program synthesis, (iii) clear ablations isolating the role of discrete control, and (iv) interpretable execution traces consistent with the proposed semantics and expressivity analysis.
6 CONCLUSION
We introduce discrete controller modules, program-like ops (ASSIGN, ADD, COMPARE, BRANCH) em- bedded in Transformer blocks via Gumbel-Softmax, enable native control flow with stateful registers. The model is trained with task, semantic-consistennyy entropy- anneall.n and sparsity ob- jectives, and comes with an expressivity guarantee for bounded imperative programs. Empirically, improves length extrapolation on algorithmic tasks and boosts symbolic QA/program synthesis with only ~5-7% compute overhead and interpretable traces/register roles. Future work includes type- d/extensible primitives, distilling away controller cost, trace verification, and broader evaluations.
10

---

Published as a conference paper at ICLR 2026
ETHICS STATEMENT
Our work augments Transformers with discrete controller modules for symbolic planning and is evaluated on public or synthetic datasets (Sortng/SSm-offLis//h)
generators, RobustFill-style DSL tasks, DROP, and the Mathematics dataset). No experiments involve human subjects, PII, or sensitive attributes; IRB approval was therefore not required. We follow the licenses and usage terms of all datasets and provide data-generation code and fixed seeds for synthetic corpora.
Potential risks and mitigations. Stronger algorithmic reasoning and program-synthesis capabili- ties may be misused in automated decision-making or oode-generation settings. To mitigate this, we (i) include interpretability tools (trace logging. knockout tests) that expose execution behav- ior, (ii) discuss deployment caveats (human oversight, constraint checks, sandboxing for code) and known failure modes (temperature schedules, premature discretization), and (iii) encourage respon- sible release practices (dataset licenses, documentation of limitations). Our models not trained to targer protected attributes, and we did not observe discrimination--rlateed effects on our bench- marks; nonetheless, applications real-world decision pipelines should include fairness audits and domain-specific guardrails. We disclose that we have no conflicts of interest or sponsorship that could bias the work, and we adhere research-integrity norms (complete proofs and assumptions, ablations, and compute reporting), See Section B for additional discussion of ethical considerations, limitations, and potential broader impacts.
REPRODUCIBILITY STATEMENT
We aim to make all results reproducible: the paper specifies datasetss/generators (Section 4.1). model and training details (Sections 4.3 3.4, Appendix E evaluation metrics/protocols (Section 4.4)
periments/ablltions (Sections 4 and provides proofs (Appendix D). We fix seeds x3 runs for synthetic, three seeds elsewhere), report t and use single 40 GB GPU. The anonymous supplementary includes minimal reproducibility materials sufficient to reproduce our pipelines; some generators/baselnne are withheld for anonymity. Deterministic runs follow our documented seed and CUDA settings; see Section B for details.
LARGE LANGUAGE MODEL (LLM) USAGE DISCLOSURE
LLM in two limited capacities: (i) copy-editing for grammar, clarity, and tone (no technical content, claims, equations, or results were introduced or altered by the LLM); and (ii) assisting literature discovery by suggesting potentially relevant papers and dataset repos. All citations and facts were verified by the authors; no text was accepted verbatim without appropriate quotation citation; and the LLM was not used design experiments, analyze results, write proofs, generate code, or create datasets. See Section B for further details and limitations.
We used
REFERENCES
Jacob Andreas, Mareus Rohrbach, Trevor Darrell, and Dan Klein. Neural module net- works. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 39-48, Las Vegas, NV, USA, June 2016. IEEE Computer So- ciety. URL https://openacces thecvf.ccmm hatprteent_cvrrr211////hhee 2016/hhtmm/Andrea Module Networks Cc___CCR_____ (eii_2216__ppa 1ar00l6_pperrttmmm
Matej Balog, Alexander L. Gaunt, Marc Brockschmidt, Sebastian Nowozin, and Daniel Tarlow. Deepcoder: Learning write programs. In 5th International Conference Learning Represen- tations (ICLR 2017), 2017. htll:httss::/oot htttps::/opeereeieew www.net//ooommiiddByyy....
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating propagating gradients through stochastic neurons for conditional computation. CoRR, 0ad//308.3432. 2013. URL
https///axiv.orr//as//3...
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. CoRR, abs/1 807. 007.99899 2018. URL httttp:://axxiv.org/aee 03819

---

Published as a conference paper at ICLR 2026
Jacob Devlin, Jonathan Uesato, Surya Bhupatiraju, Rishabh Singh, Abdel-rahman Mohamed, and Pushmeet Kohli. Robustfill: Neural program learning under noisy I/O. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 990-998, Sydney, Australia, Au- gust 2017. PMLR.
mlr. pttpssss////////eeiii........ http:///wcccoonnnnoooooooooo//na
Dheeru Dua, Yizhong Wang, Pradeep Dasigi, Gabriel Stanovsky, Sameer Singh, and Matt Gardner. DROP: A reading comprehension benchmark requiring discrete reasoning over paragraphs. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Com- putational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp- Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: URR.hhttps://allanhhollo 1ttpaaannhhooooooooo///11/111-
10.18653/v1/N19-155
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23 (110))1--39 2022. doi: 10.485500/xXiv.2000 11396111.1..
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. CoRR, abs/2012.149133 2020. http:///jaaa..rrrrppapass/////// -222 html
Alex Graves, Adaptive computttion time for recurrent neural networks. CoRR, abs/1603.08983, 2016. hBL//tttss///ariiior////////// 08983
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. CoRR, Abs/1410,5401, October 2014.
wwwp.iwwww.//////1/// 5401 arXiv preprint.
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Natalia Grabska- Barwińska, Jesús A. Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, Adrià Puigdomènech Badia, Karl Moritz Hermann, Yori Zwols, Georg Ostrovski, Adam Cain, Helen King, Christopher Summerfield, Phil Blunsom, Koray Kavukeuoglu, and Demis Hass- abis. Hybrid computing using a neural network with dynamic external memory. Nature, 538 (7626):471-476, October 2016. ISSN 0028-0836. doi: 10. 10.1////at URL https:
7/www. htw.:r/cccco/rrrcccssstttrrre/
Edward Grefenstette, Karl Moritz Hermann, Mustafa Suleyman, and Phil Blunsom. Learn- ing to transduce with unbounded memory. In Advances in Neural Information Pro- cexsing Systems 28 (NeurIPS 2015), pp. 1828-1836, Montreal, Canada, 2015. Curran
Associates, Inc. URL https::/procceedings www.npe...c/pppppppeee http://cccc.app...eee/////////// 01839778009994222111222222226683
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. In
5th International Conference on Learning Representations (ICLR 2017), Conference Track Pro- ceedings, Toulon, France, April 2017. URL httt::///peerrreeeewe////rrrrrrn Published conference paper at ICLR 2017.
Armand Joulin and Tomas Mikolov. Inferring algorithmic patterns with stack-augmented recur- rent nets. In Advances in Neural Information Processing Systems 28 (NeurIPS 2015), pp. 190- 198, Montreal, Canada, 2015. Curran Associates, Inc. URL http:::::/ppaprre/ wwwsssssissscc.mm wwt.:nic.ppppppeeeeeeeep cc/paper/////// 5857
Lukasz Kaiser and Ilya Sutskever. Neural gpus learn algorithms. In 4th International Conference on Learning Representations (ICLR 2016), Conference Track Proceedings, San Juan, Puerto Rico, May 2016.
www.lv/wogg//bbbo wttp/////////r..ooo//.////// 08228 Published as a conference paper at ICLR
2016.
Daniel Keysers, Nathanael Schärli, Nathan Scales, Hylke Buisman, Daniel Furrer, Sergii Kashubin, Nikola Momchev, Danila Sinopalnikov, Lukasz Stafiniak, Tibor Tihon, Dmitry Tsarkov, Xiao Wang, Mare van Zee, and Olivier Bousquet. Measuring compositional generalization: A compre- hensive method on realistic data. In 8th International Conference on Learning Representations (ICLR 2020), 2020. URL htt.:t//ss//ppppomeeeeeeee/////s
12

---

Published as a conference paper at ICLR 2026
Brenden M. Lake and Marco Baroni. Generalization without systematicity; On the compositional skills of sequence-to- -eeeueeee--- recurrent networks. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th Intemationall Conference Machine Learning, volume 80 of Proceed- ings of Machine Learning Research, Stockholmsmmäsan, Stockholm, Sweden, 2018. PMLR. URI. nttps://procceedings..u mlr. wttp:/iiww..mm..pes///////////// wwwps//wee/aaeeeaaaeee html
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with condi- tional computation and automatic sharding. In 9th International Conference Learning Repre- sentations (ICLR 202/), Conference Track Proceedings, 2021. URL, wtw.hhttss//ppprriiieeee
forum?id=qrwe7XHTTne
Kanika Madan, Nan Rosemary Ke, Anirudh Goyal, Bernhard Schölkopf, and Yoshua Bengio. Fast and slow learning of recurrent independent mechanisms. In 9th International Confer- ence on Learning Representations (ICLR 2021), Virtual Conference, May 2021. URL https:
"'operreview.nnt//foummm http///r
Chris J. Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous
relaxation of discrete random variables. In 5th International Conference on Learning Rep- resentations (ICLR 2017), Conference Truck Proceedings, Toulon, France, April 2017. URL https://openreview. jet// Published as a conference paper at ICLR 2017.
Jayyaa Mao, Chuang Gan, Pushmeet Kohli, Joshua B. Tenenbaum, and Jiajun Wu. The neuro- symbolic concept learner: Interpreting scenes, words, and sentences from natural supervision. In 7th International Conference Learning Representations (ICLR 2019), 2019. doi: 10.48550/ httpri/eee/eeoorrrrrddrrrggiiene
arXiv. 12584. URL
Nikita Nangia and Samuel Bowman. Listops: A diagnostic dataset for latent tree learning. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Student Research Workshop, pp. 92-99, New Orleans, Louisiana, USA, June 2018. Association for Computational Linguistics. doi: 10.18653/ URL
https://accann hawp:holoclathooogyyyy www.oogg..oggg/// /errg/N1--4013
Jack W. Rac, Sebastian Borgeaud, Tim Harley Cai, Katie Millican, Jonathan Raiman, Giuseppe Campagna, Lukasz Kaiser, Helen King, Geoffrey Irving, and Timothy Lillicrap. Compressive transformers for long-range sequence modelling. In 8th International Conference on Learning Representaations (ICLR 2020), Conference Track Proceedings, Addis Ababa, Ethiopia, April 2020. Published as a conference paper at ICLR
URL httllittss///ppeeeiee...o. wwtp://wwwwweeooo/m//////ssiiaaa 2020.
Scott Reed and Nando de Freitas. Neural wrogrammerriineerpretee In 4th International Conference on Learning Representations (ICLR 2016), Conference Track Proceedings, San Juan, Puerto Rico, May 2016. htl.httpss///arvv.or//////////// 2016.
06279 Published as a conference paper at ICLR
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Lin- guistics, 9:53-68, 2021. t.c1-1.4
doi: 10. (0.1162/aacaa_00355 URL htttps://aclanthoolog. htwpshhlllg..oogg////// org/2021.
David Saxton, Edward Grefenstette, Felix Hill, and Pushmeet Kohli. Analysing mathematical rea- soning abilities of neural models. CoRR, abs/1904.01557, 2019. http:/tts:////xxiiorrr///
1904. 01557]
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Geoffrey E. Hinton, and Jeff Dean. Outrageously large neural networks: The pparselyygatte mixture-of-expert layer. In 5th International Conference on Learning Representations (ICLR 2017), Conference Track Proceedings, Toulon, France, April 2017. URL URR.httpss://pperrviiee. wwwpe/enrreiee..eeefrrruee wett/forummiddd-ickooddd
Andrew Trask, Felix Hill, Scott Reed, Jack Rae, Chris Dyer, and Phil Blunsom. Neu- ral arithmetic logic units. In Advances in Neural Information Processing Systems 31 (NeurIPS 2018), 2018 URL //ttps:// wewrrediinnnnnnuiis..... www..iis.cc/papprrrrrree..e files//aperr/2111/ hil.//64aaboccaai22eeaaa/2222223
13

---

Published as a conference paper at ICLR 2026
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Gamett (eds.), Advances in Neural Information Processing Systems 30 (NeurlPS 2017), pp. 5998-6008 Red Hook, NY, USA, 2017. Curran Associates, Inc. URL https: http::::///prrceeinggg...o www.ips.seeriir..... cc/ abstract..tmm
Petar Veličković, Adrià Puigdomènech Badia, David Budden, Razvan Pascanu, Andrea Banino, Misha Dashevskiy, Raia Hadsell, and Charles Blundell. The CLRS algorithmic reasoning benchmark. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Abhradeep Thakurta, and Sergio Verdú (eds.), Proceedings of the 39th International Conference on Ma- chine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 22084-22102, Baltimore, Maryland, USA, July 2022. PMLR. URL https::/prooceedings.s. wtt.eeddiin..mm....rr////// press/v162/ velickovic22a.he html
Oriol Vinyals, Meire Fortunato, and Navdeep Jaitly. Pointer networks. In Advances in Neural In- formation Processing Systems 28 (NeurIPS 2015), pp. 2692-2700, Montreal, Canada, 2015. Cur- ran Associates, Inc. URL httttpss//prrcceeinggs. neur ottaa//eeoooooaoooaoooaaaaaaaaaa
Wojeiech Zaremba and Ilya Sutskever. Learning to execute. CoRR, abs/1410.4615, October 2014. URL 1992100fFFF40cclaee24422999888--
A
RELATED WORK
A.1 EXTERNAL PLANNERS AND MEMORY-AUGMENTED MODELS
Classical neuro- symbolcc lines attach an explicit controller or planner to a neural network. Neural Turing Machines (NTMs) and Differentiable Neural Computers augment sequence models with an addressable external memory and learned read/write heads (Graves et al. 2016). Neural GPUs learn algorithm execution with grid-like recurrent convolution (Kaiser & Sutskever 2016). Stack- augmented and differentiable data-structure RNNs implement pushdown-style computation via soft stacks/queues (Joulin & Mikolov 2015 Grefenstette et al., 2015). Program-executing agents such as the Neural Programmer-Interpreee learn hierarchical subroutines with a learned step controller (Reed& de Freitas 2016). While powerful, these approaches keep symbolic control as a separate subsystem and often increase architectural complexity and training brittleness; our controller is in- place, lightweight, and shares the Transformer's optimization pipeline.
A.2 ALGORITHMIC BENCHMARKS, COMPOSITIONAL GENERALIZATION, AND TRANSFORMER LIMITATIONS
Benchmark suites have exposed gaps in length and compositional generalization, CLRS targets graph/data- algorithms (Veličković et al. (022), ListOps probes hierarchical computation on nested expressions (Nangia & Bowman 2018), and SCAN/CQQ stress compositional recom- bination in seguence-to-sqquemee settings LLake & Baroni 2018 Keysers et al., 2020). Vanilla Transformers fit in-distribution but degrade sharply under longer inputs or novel compositions ar- chitectural variants (e.g.. Universal/Compressit
Transformers (Dehghani et al. 2018 Rae et al.) 2020 help on context length but still lack explicit control-flow semantics. Our modules target this missing bias by instantiating opcode-level semantics (assigi/ad/commaarebbach))) inside the layers.
A.3 DIFFERNNTIALLE CONTROL FLOW, DISCRETE ROUTING, AND ADAPTIVE COMPUTATION
Learning discrete decisions inside deep nets is commonly approached via the straight-througe esti- mator (Bengio et al. 2013 or the Bannrrte/Gumbel-SSee -ofttball-of relaxation (Jang et al. 2017 Maddison et al. 2017). Adaptive Computation Time (ACT) introduces learned halting in RNNs and underlies the Universal Transformer's 's dynamic depth (Graves] 2016 Dehghani et al. 2018). Sparse Mixture- of of-Exxerts and routing families perform top-k booken-wisee dispatch with discrete gates (Shazeer et al.)
2017 Lepikhin et al. 202T [Fedus et al. [022]] Roy et al. (02T). Recurrent Independent Mecha- nisms encourage modular, sparsely interacting dynamical subsystems (Madan al. 2021). Unlike
14

---

Published as a conference paper at ICLR 2026
routing/halting, controller assigns program- like semanties decisions and exposes execution trace over a compact state, enabling proofs of expressivity and targeted interpretability.
A.4 NEURAL PROGRAM INDUCTION AND PROGRAM SYNTHESIS
Early neural program induction learned to execute pseudocode or DSLs directly from examples (/aremba & Sutskever 2014 Vinyals et al. 2015). Differentiable interpreters and program- induction systems (e.g.. Neural Programmer, NPI) sought modularity and verifiable subroutines
(Reed & de Freitas 2016). program synthesis from I/O, RobustFill introduced a DSL with pow- erful search neural scoring (Devlin et al. 2017), and DeepCoder learned to predict useful library components to guide symbolic search (Balog et al. 2017). Our program-style controller comple- ments these lines by embedding primitive semantics inside the Transformer rather than emitting code externally.
A.5 NEURO-SYMBOLIC REASONING AND MODULAR ARCHITECTURES
Neural Module Networks compose small operator modules into executable graphs for vi- sion-language reasoning (Andreas et al. 2016). while Neuro-Symbolic Concept Learners factor perception from logic over latent programs (Mao et al.) 2019). Contemporary LLM tool-use exe- cutes external programs (c.g., calculators), but still keeps control outside the model. Our approach keeps the control state internal and differentiable, producing traces that can be aligned with ground truth steps on algorithmic and symbolic- QA tasks.
A.6 NEURAL ARITHMETIC AND MECHANISTIC INTERPRETABILII
Neural Arithmetic Logic Units target exact extrapolative arithmetic by design (Trask et al. 2018), orthogonal to our control-flow focus and complementary in principle. Mechanistic studies show token-level circuits and key-value mechanisms emerging in Transformers (Geva al. (020); our controller makes such structure explicit, enabling direct trace alignment and knockout analyses tied opcode semanties.
B ADDITIONAL DISCUSSION
Strengths and contributions. Our results indicate that inserting discrete controller primitives into Transformer blocks provides a simple yet powerful inductive bias for algorithmic reasoning. Empir- ically ($5), the model sustains high sequence-level accuracy under 2 xx-4x length extrapolation with only a modest 5-7% compute overhead (. while also improving performance on symbolic QA and program-synthesis style tasks. Conceptually, the controller formalizes (oop/branch/updatee se- mantics inside the network; this yields (i) better generalization beyond the training regime, (ii) inter- pretable execution traces (operation sequences and register roles), and (iii) theoretically grounded view of why discretization helps (Theorem] [ [nd finite-temperature analysis). Together, these prop- erties make the architecture compelling, minimally invasive building block for models that must manipulate discrete structure.
Limitations and failure modes. First, training stability depends on the temperature/entrapy schedule. If the controller discretizes too late, branches blur and errors accumulate on long puts; too early, the model can lock into suboptimal programs. Our ablations show that entropy penalties and semantic consistency losses mitigate this, but do not eliminate the sensitivity. Sec- ond, our expressivity guarantee targets bounded loop unrollings; while depth can emulate many steps, extremely deep programs may require more controller placements or curriculum. Third, the present design assumes small set of primitive ops (ASSIGN/ADDCCMMPAA//BBACCCC (e.g., stacks, maps, heaps, arithmetic over big integers) likely need additional typed primitives and address mechanics. Fourth, although compute overhead is moderate with controller width and insertion frequency; large-seale deployments will benefit from sparsity (activating controllers only when needed) compilation. Finally, while report gains DROP and RobustFill-stye tasks, broader evaluation on code repair, table reasoning, and multi-hop QA will strengthen external validity.
richer domains
settings, it does grow
15

---

Published as a conference paper at ICLR 2026
Threats to validity and fairness. We match parameter counts and report FLOPs/tokens/s for all baselines, but residual differences (c.g., memory in Compressive Transformers) may still confound comparisons. Statistical significance based on few seeds is limited; therefore advocate reporting confidence intervals and non-parametric paired tests in future revisions and releasing generation seeds/configs for full reproducibility. Dataset artifacts (template regularities in synthetic corpora) can inflate gains; our numeric-only DROP subset partially addresses this but cannot preclude all shortcuts. Finally, we emphasize that "interpretable" traces show correlations between controller choices and ground-truth steps; they do not, by themselves, constitute causal explanations.
Potential applications. Discrete controllers are promising for tasks that naturally factor into program-like steps: numeric and tabular reasoning (SQL-style aggregation and comparison), pro- gram synthesis/repair in small DSLs, structured data transformation (e.g., spreadsheets), algorithm tutoring (showing stepwise traces), and robotics or planning where loop/branch patterns recur. Be- cause controllers expose explicit decision interface (operation type, addresses, values), they also open doors for hybrid verification: one can check controller traces against invariants or lightweight SMT constraints without instrumenting the entire network.
Future directions. Two near-term extensions are attractive. (i) Typed and extensible primitives: introduce a registry of typed ops (e.g., stack PUSH/POP, map GET/PUT, arithmetic with carry) and learn when import new primitives, akin library growth. (ii) Compile-time distillation: after training with differentiable controllers, distill the learned control flow into an executable symbolic program (or a compact neural program) and skip controller computation inference, reducing overhead and improving robustness. Beyond these, dynamic sparsification (event-driven controller activation), integration with retrieval/tool-uss (controllers as planners that call tools), and reinforcement learning for long-horizon tasks are natural next steps.
Ethical considerations. By improving reliability on discrete reasoning, the model may be de- ployed in decision pipelines (finance, education, code generation). We caution against fully au- tomated use in high-stakes settings without human oversight, audit logs of controller traces, and guardrails (e.g., constraint checks, anomaly alerts). For program synthesis, controllers could pro- duce insecure non-compliant code if objectives should be standard. Energy and compute costs are modest relative to baseline Transformers, but encourage reporting carbon metrics and preferring sparse/compiled variants when possible. Fi- nally, dataset curation must respect licenses and privacy; releasing generators and seeds supports community scrutiny and reduces hidden biases.
misspecified; sandboxing and static analysis
Discrete controllers provide principled route to planning inside Transformers: they improve length extrapolation, make intermediate reasoning legible, and connect cleanly theory, while remaining lightweight enough for practical use. Addressing the noted limitations (training schedules, richer primitives, broader evaluations) will further consolidate their value for robust, interpretable reason- ing.
C FORMAL DEFINITIONS AND NOTATION
This appendix fixes notation and gives a precise, self-contained description of the controller- augmented Transformer used in the main paper. Throughout, all random variables defined on an implicit probability space (2,F, F.FF.PP P);; bold uppercase denotes matrices, bold lowercase denotes vectors, and || ||a the Euclidean/Frobenius norm by context.
Sequences, embeddings, and layers. Let token sequence be 21.T nT with vocabulary V. Tokens are mapped to embeddings and (absolute or relative) positional encodings to yield X E RT RTXD A stack of L Transformer blocks produces hidden states H))) = X and H!!) E RTxD for l (I....,L} L); each baseline block is the standard-
Block(U) = FFN(MHA(UU)
(19)
with residual connections and layer normalizations as in Vaswani et al. (2017). Multi-head attention uses H heads with head width D3 D/H.
16

---

Published as a conference paper at ICLR 2026
Table Compact notation table (symbols used most often)
Symbol
Meaning
T
sequence length (tokens)
D. H. Lust
hidden width, # heads, # layers (network depth)
Nr,
registers and their width
Nm, dm
memory slots and their width (optional)
N
# relaxed Boolcan flags
H E TRR
hidden states at block
5O) (R.), M
program state at block
( = )(H())
controller summary (token - state)
D0 (()))
state bias (state token)
O
operation set (ASSIGN, ADD, COMPARE, operation logits and Gumbel-SSftmax selector read and write address distributions over registers read vector and candidate write value temperature (selection), sharpness (comparison) semantic map for operation acting state controller state width N.d, Nm.dm N
COMPARE, BRANCH, NOOP)
()) (
a 80
(9 (0
T, R Sem. d.
Controller-augmeented state space. At each instrumented block I, the model maintains a program
state-
( = RR (1) (()(()))
(20)
where (() E R-N-xdr are N, register vectors, M Nm xdm is an optional key-value memory (treated learned scratchpad; Nm, may be zero), and f(1) E [0, 1] flags that encode Boolean conditions in relaxed form. The controller reads summary (()) E Rdd. from token states via pooling map.
800 = (H
(21)
with R + Rd. (e.g.. learned attention pooling or mean pooling, with learned linear projection). state-to-token encoder N- x Nm xdm x (0, 11 + RDD produces a broadcast
bias ((0) (S())) that conditions the token pathway-
H (b())
(22)
T Block (A
(23)
Operation set and discrete selector. The controller chooses among a finite set of primitives-
[ASSIGN, ADD, COMPARE, BRANCH, NOOP)
(24)
Given s(O), logits n (1) ( +b, E Rlo are converted to a relaxed one-hot vector (!) E 101-11
via Gumbel-Softmat at temperature > 0-
(() softmax 7 (9 +g(")
(25)
g =o
(1) log( log u
"
Uniform((.)
(26)
As A converges almost surely to one-hot argmax (the discrete operation).
Differentiable addressing and arguments. The controller produces read and write address dis- tributions over registers and candidate value-
( softmax(W,,s (
(27)
( Seftmax(WW,ss)
(28)
(0) (.
(29)
17

---

Published as a conference paper at ICLR 2026
Soft read and write are-
N. ,
R Write(R g v())
=
(30)
R R
(31)
where Write mixes candidate updates into register rows according 3 (e.g., convex combination with strength A (0, 1)..
Primitive semantics as maps on state. Each operation E O induces deterministic map Sem, on the state, parameterized by learned weights (e.g., small MLPs for branch arms)-
Sem. (R, r, - (R', (R','',") F')
(32)
The overall state update at block / the mixture-
s(1+1)
Σn vou Semo(S r(), v(0)
(33)
DEO
We instantiate the primitives with the following relaxed semantics (elementwise where appropriate) -
ASSIGN: : R(++1) = (1 - A) R + +XE(B()) (V(1))
(34)
ADD : R(++1) 1) = R + E(B (VI))
(35)
COMPARE :
p((+)) (K [AR! (1) BR
(0]
(36)
BRANCH
s(+1) -0(0) v cheen
(ch (5 + (1 711 V (S()e
(37)
NOOP
S(11) (())
(38)
Here E(B) N.x1 denotes (soft) one-hot write address; À is a write strength; is a logistic with sharpness & >> 1; A, B are soft selectors over registers; and the branch gate is 7()) ow p() +
b.). The branch subprograms V then: Wesee are small shared MI.Ps that transform (R. M, f) without touching token states directly.
Token-state coupling (write-back). Changes the program state may written back to the token stream via an alignment matrix W,-
(L+1) = Z + ( s()) W.
(39)
where p maps state increments to a D-dimensional residual (e.g., concatenation of register means fed through a linear layer). In practice, only every p-th layer is instrumented control compute.
Temperatures, randomness, and annealing The Gumbel-Softnax temperature is annealed during training by-
T() max (Tmin: TD e --t)
(40)
with global step 1, initial To, floor Tmin. and decay 7. The comparison sharpness k may be increased similar schedule. At inference, either the relaxed selector is kept (soft execution) hard
(!) decision used via arg max Th (optionally with a straight gradient estimator during training).
Loss-aligned invariants (semantic regularization.. When task semantics imply invariants (e.g., value conservation for certain steps, monotonicity constraints sorting subroutines), the model employs penalties of the form-
Rine
Inv ||2
(41)
LET
for an index set of steps where an invariant Inv(-) should hold. This connects the relaxed semantics to intended discrete behavior and stabilizes training.
18

---

Published as a conference paper at ICLR 2026
Well-posedness. Under bounded parameters and Lipschitz continuous then: Velses the mixture update 1(++1) is well-defined for all r > 0 and yields measurable mappings of (H S ). In the limits - and i -> oc, the relaxed dynamics converge to a discrete execution trace consistent with the intended primitive semantics (formal statements and proof sketches appear in the proof' D).
This formalization isolates the controller's contracts: how it reads from and writes to a compact program state, how selects and executes discrete primitives in differentiable manner, and how the resulting state interacts with the token stream. All subsequent derivations, losses, and experiments reference these objects.
D PROOFS: EXPRESSIVITY AND FINTE-TEMPERATURE BOUNDS
Setup and assumptions. We use the formal objects from Let O (ASSIGN, ADD, COMPARE, BRANCH, NOOP) denote the primitive operations, each with semantic map Sem, acting on the program state S = (R, M, f). At instrumented block I, the controller forms logits π (1) over O, produces a relaxed one-hot y() via Gumbel-Softmax temperature т, and computes-
s((+1)
Σ v.a Sem,(S") 1 v() OEO
(42)
with read/write addressing and candidate updates defined C We will work under the following standard, verifiable conditions on the compact state domain D encountered during execution.
Assumption B.1 (Lipschitz semantics). Each Sem. D D is L--Lsscchit:: [[Sem.(S)
(e)l(S)] & LoI|s 31,, and max, L. S Lmmxx < oc.
Assumption B.2 (Bounded state and parameters). All states remain in compact D and con- toller/Transformer maps are bounded on D.
Assumption B.3 (Decision margins the trace). Along the intended (discrete) program trace, the correct operation at step has a uniform logit margin > 0 over all other operations (achievable by capacity/regularizatoon or weak supervision on synthetic tasks). Similarly, the COMPARE-ines COMPARE-inducing logits and address logits have margins cemp &adtir > O on D.
Definition (Program class PK,B). Fix K E N and a loop bound B E: N. PK, consists of straight- line programs over real registers with primitives O, conditionals whose guards are Boolean flags set by COMPARE, and loops whose bodies are drawn from the same primitives and unrolled at most B iterations. The total number of primitive steps in any execution of P Pk, most K+B.
Notation. Here L denotes the trace length Li (number of program steps), which satisfies Lu- S K+B; when we discuss network depth we write Luet-
Theorem B.1 (Expressivity for bounded imperative programs). Let P Pk. and let (Sx(0)
be its execution trace on input (X, S-()), with L S K + B. There exists a controller-augmented Transformer with L' instrumented blocks (e.g., instrument every block so L' L) and parameters (0, Q) such that, as r >> and the comparison sharpness x - oc,
|| - S000 a.s.
O
for all = 0, 1, ....., L
(43)
and the discrete operation chosen at each step equals the program primitive with probability in the limit. Moreover, the token pathway can be set to pass-through (or implement arbitrary continuous encodings) without affecting the state simulation.
Proof. We give a constructive simulation of P and prove by induction on the program step that the network state matches the execution trace in the stated limits. We assume a standard Pre-LN Transformer; when needed we set sublayer weights to zero so the residual path yields identity on the reserved stare-carrying subspace.
Setup and coding of state. Fix d, 2 and embed the i-th scalar register value in a dedicated coordi- nate of R,, keeping the remaining d,-1 coordinates Take M 0. Let the flag vector f E [0, 1)
19

---

Published as a conference paper at ICLR 2026
carry Boolean outcomes of COMPARE and one-hot program counter pe {1...., L). Reserve sin- gle token t"; write (x) into that token via the residual addition, set the block's sublayers to zero so it is identity, and let & read token t* with a fixed linear map U, yielding s()) = U vec(S() I Assume along the execution trace that every read/write address has unique maximizer with margin Ma > 0 and every branch predicate has margin my > 0.
Operation selection. Let of E O be the primitive prescribed by P at step I given (pe!), (()). There exist W., b, and a margin m, 0 such that
(Wos (1) + bb)
2 (Wos + bb) + mo Vo + of
(44)
Scale logits as (()) = co (Wos () + ba). Under Gumbel-Max, arg max. (To (1) + +o) is categorical with probabilities softmax(T (1)) and
P arg m (0) + [o) = 0j 2 1 (10) - 1)ee co
(45)
so selection is correct almost surely as c, to and T -
Addressing and arguments. For each desired register index jr at step I, choose a vector v(() E IR N, with unique maximizer jr and gap at least ma and set
a (1) softmax(c. B(0) softmax(c,
(46)
Then a (1,), ,0) en Co Set (() (,,(()) to the exact constant or linear form required at step 1; (1) contains ()) linear map suffices
Primitive semantics. With one-hot write addressing and A= 1, the updates
R ((+1) R (t) + E(3")) (()
(47)
p(++1) (1 [AR JR"))
(48)
implement exact ASSIGN or ADD (by choosing v(t) accordingly) and exact comparisons in the limit Ca > 00, K -> oo because A, B converge to one-hot selectors and (<z)) - Hz)). Specifically, ASSIGN is realized by setting (!) = v - RR for the selected write index j", or equivalently by viewing ASSIGN as (1 - R) (1) + E(G(U) (UT with 1. Choose A, B via the same softmax- addressing construction as a. 33 with scale ca. so the margin m. enforces unique selection. For BRANCH, pick wy, b, so that w. (1) + bb, equals the predicate bit and its pre-activation is bounded away from 0 by my; then (0) aw. f()) + b.) + {0, Let V then and Vesse be affine maps that update the program counter (and optionally write constants) exactly; affine maps are represented exactly by two-layer ReLU MLP.
Inductive simulation. Unroll bounded loops to at most B iterations, yielding L = O(K + B) steps. Suppose s(1) equals the program state after / steps. The constructions above ensure that the controller (i) selects the same opcode, (ii) reads/writes the same addresses and values, and (ii) updates flags and the program counter identically in the limits Co, + co, 7 + 0, and + co. Hence (((+1) matches the next program state. The base case holds by initialization, proving exact simulation along the full trace.
Token pathway. Set p p= (or reserve coordinates) so the token stream is a passive carrier of (S),,, which suffices since (1) is recoverable by construction. Thus the token computations do not limit state expressivity on the reachable set.
Lemma B. (Gumbel argmax limit). Let e RT and let g have i..dd standard Gumbel compo- nents. Define (T)) = softmax( (T g)/r) and
k* (g) arg m + gb),
Agg) (TT*(() + (k-(())
max (T, + gz)
J (g)
'In Pre-LN block, Y = x + SubLayer(LN(X)). Setting the sublayer weights to zero makes Y = X exactly, by reserving one token and writing (8)) into it via the residual addition, can recover the state linearly by attending to that token.
20

---

Published as a conference paper at ICLR 2026
Then (g) > 0 almost surely and
as. y(T)
C₁ (x)
(49)
T-20
Moreover, for almost every realization of g and all т > O,
A(g)
||(()) e |||
S 2 (m (m-1) exp
(50)
In partieular, if (g) is almost surely fixed index k* (i.e., PPkk* (g) k*] then y(r) - e/-
almost surely
Proof. Define the random score vector
a(g) == π T orggen Rm
(51)
and note that
(() softmax a(g)
(52)
Because the components of g are i.d. continuous, ties occur with probability zero, there is unique maximizer almost surely
k')g) = arg max (g(g)
(53)
ke().....
Denote the random top margin
A g = ag (g) (g)
max az(g) (x)*
a.s.
(54)
Applying Lemma B.3 to the fixed realization agg) with N, = m and j = kg) yields, for every
>>>
[|x(r) - ex-( S (m--) exp Ag))
(55)
Since Agg) 0 almost surely, the right-hand side of equation 55 converges T - 0, for almost every realization of g. Hence
a.s. (K)
(56)
T
For the addendum: if the argmax index is almost surely fixed ke (equivalently, P(k*(g) k+1
1), then equation reads y(r C₁ almost surely, as claimed.
read
Lemma B.2 (Sharp comparison). Let (x())) = a (R2) = (1 (1ke-_) Then for any E > O,
sup |a|2s
(xx(2) - 1(220)1 s enE
(57)
Hence 0x + [2200) uniformly on IR (-E,)) as a
Proof. Let 0x)=) = Ite mz and denote the Heaviside step by H(z) = 1 [220)) For 2 > 0 we have
K2
--)x((2 ox(2) Ite- I
(58)
and for z < 0,
[x)()) - H(z) =x(2) = 1 ell
(59)
Combining equation 58 equation 59 for all z = 0,
(ax(2) H(2) 1 |</
1
S -H|z
(60)
where used 1+er S 2 for 2 0
21

---

Published as a conference paper at ICLR 2026
Now fix any 0. Since
|2| + 1+e*|=] is strictly decreasing on [0, oc),
1
sup |=|
|ox(2) (x>00) 1 + enc I ee
(61)
Because the right-ham hand side of equation 61 tends to 0 as & - oc, we obtain uniform convergence:
sup =ER((-e,e)
[ox(2) - (x>03 X-+20
0
(62)
which proves the claim.
Lemma B.3 (Address concentration). Let a E R- have unique maximizer at index jr, with
margin = min (ag- - o4) > 0. For (T)) = (a/tm
[1 - 1 s 2(N, ((N--1)e A/r
0
(63)
Proof.
Write the softmax components
(T(T)
jE (1,..,,NN)
(64)
Σ ox//t
Let jr = arg max, aj be the unique maximizer and define the margins
ój == aj- aj > for j jJ"
A = min 63 > 0
(65)
j7j"
Then for j j*
(T)
(1j 13-0g. 0g*
S exp (- A
exp
exp
(66)
(T)
7
7
Define the tail mass
0j
8j = exp
Σ
8j
S (N, 1)e
(67)
From equation 64
equation 67 we obtain the normalized components
(T) I+s
B,(T)
8j for j + j
(68)
1+s
Hence the l distance to the one one-hot hot vector Cy+
2x
[|ß(T) -z- 1a-ll,=10 = -
Σ ( -1- j*1
(69)
1
Using * I+s S s and equation 67 we obtain the exponential bound
(T()) e,- |1, 2s 2s < 2 (N, - 1) 4/r
(70)
Since A > 0. the right-hand side of equation 70 tends to as -> 0, so
lim [|((()) e,- 11 0 7
(71)
which proves the claim.
Lemma B.4 (Branch selection equivalence). Let V Whhen V Volse D - D be L and define
the relaxed branch update (S) = then (5) + (1 11)) Weise (S), with y = a(w f b). If +
(resp. 7 + 0), then
(S) V Vehen S (1 1--7) 7) Co + 0 (resp. y Cv 0)
(72)
for C. = supsed V (S) - Weisa(S) < oc. Thus the relaxed branch converges to the discrete path.
22

---

Published as a conference paper at ICLR 2026
Proof. By definition,
ES(S) 7 V then then (S) + (1- (S), 7 7= f f
E (0,1)
(73)
Since v then and elsee are L Lipschitz D, they are continuous; if D is compact then the function
G(S)
v then (S) Weise(S)]]
(74)
attains its maximum on D. hence
C == sup SED
then (S) V Weise(S)]|
< oc
(75)
Now, subtract v "henn(S from equation 73
and use the triangle inequality:
||,((S)
v then (S)
[[(1 7) (S) v then (5))
s (1 7) (eeeee()) 1// chen (S)
(76)
S 7) Ce
An entirely symmetric argument yields
(.)()) Vesso 7 V Vzhen (8) - 1/ (S) S Co
(77)
Therefore, 7 + 1, the bound in equation 76 implies
[S) - W chen (S)]l $ (1 7- Cs
0
(78)
7-1
and if - 0, the bound in equation 77 implies
FF, (8) - (e)s
C
0
(79)
The bounds equation 76 equation 77 are uniform S because Cs is independent of S. Conse- quently, if - 1 (resp. 7 - 0) uniformly in S, the convergence in equation 78 (resp. equation 79 uniform; otherwise it is pointwise. This establishes equation 72 and shows the relaxed branch converges the corresponding discrete path.
Theorem B.2 (Finit-temmperature and sharpness error). Let S*(1) be the discrete trace pro- and comparison
gram in PK.B and let ( be the controller-augmenter state with temperature sharpness x 0, initialized at ((D) S-((), Under Assumptions B.I...3.
- - - C (
for all L < K + B,
for constant C' depending on (Lmax, K+B) and the decission/addtrccompare margins.
Proof. Setup and notation (uniform margins and regimes). Let
es = B||( Sx(0
(80)
Assumptions B. B.1 hold uniformly for all steps 1 < K+ B and all states in compact domain D: (i) each primitive semantic map Sem, is Tips-hipschitz on DD (ii) the decision margin Asp > separates the correct opcode o* from the others, (iii) the address margins A Awr 0 separate the correct read/write addresses, and (iv) the flag/predicate margin my > 0 separates comparisons from zero (used by BRANCH). Fix a small-temperature regime
S To and a large-sharpness
regime 2 Ko := 1/my. We also denote
Cup
sup max [[Sem.(S) - Sem.. (S) SED ofo"
(81)
which is finite by compactness of DD and continuity of Sem,, Finally, the read/write mechanisms are linear in the addressing weights and D is bounded, hence there exists
Caddr sup
sup ||
Sem..
Sem.- (S:a, < oo (82)
SED
so that addressing deviations control state-update deviations
23

---

Published as a conference paper at ICLR 2026
Step 1: One-step decomposition. Add and subtract Sem. (8(0):
-
- --11+11||| = E2
Σ yo Semo (()) Sem ((((1)
IF
wo Sem (5()) Sem. Sema-((
(83)
+
E Semo Semo-((
E
Sem. (S()) Sem (S*0)) - Linax
(84)
Step 2: Operation-selectionn error. Using Eo yo
and the bound equation 81
Σ y (e) Sem (5())
Σ
yo [Sem. Sem. (S())
(85)
Σ
(1)
Cop = [1 - You
(1) Cop
or2""
Taking expectations and invoking the Gumbel argmax limit (Lemma B.1 with uniform decision margin Aop > 0 yields, for O < To.
(101-1) e Dppt
E[1 [
vo0l < (|0) Am/T S Cop Cop
sup 0<r
<000
(86)
since A/T = o(T)
Step 3: Address-mixing error. When read/write addressing uses soft distributions β instead of
one-hot e;,; ejr, equation 82 implies
E addr error step |] < Caddr BiIla e,, 11 || ej- ]1]
(87)
By Address concentration (Lemma B.3 with uniform margins And, Awr
E|||a e;- |11 + 1|3 - -,,- [h] < 4(N,-1)e e Amat/r
Cadde T,
A Addr = min ( Ard Awr)
(88)
4(N, 1) e Aualt sup
Caddr
04T
T
Step 4: ComparisorDrarch-aating error. For predicate value with uniform flag margin
| 2 my > O, the relaxed gate "Yx = a (Kz) satisfies (Lemma B.2
hix H(2) < e kmy <
for all 2 Ko = 1/m,
(89)
emy K
By the branch selection equivalence (Lemma B.4), the resulting state update differs by at most
Co 1
E cmmpbrrncc error at step i] < Cy EE [[ H(:)]] <
Comp K
(90)
emf K
Step 5: Collecting step-wise errors. Combining equation
equation 84
equation
85 equa-
tion 86 equation 87 equation 88 and equation 90] we obtain
11+1 - Lmax Cop Cop + C'dtrr Caddr Cemp
-- Lmxx + + bx
(91)
with CopCap CuidrCaabtr and Cemp
24

---

Published as a conference paper at ICLR 2026
Step 6: Discrete Grönwall (unrolling). Since ( - we have eg = 0. Iterating equation 91 for L $ K+B,
4-1
L L
el
x
maxx (at+bk=t) = Lmax
(a (aT+bR-1)
(interpreted as when Lmmax 1)
1=0
(92)
Define
L
L Li if Lmax (93)
C
Lmax maxfa, b}
with the convention
ILLSK+
Lmax
which depends only (Lmax: K+B) and the stated margins/diameters. Then from equation 92
E E||s(() - s+(2) = el S C (++ for all L K+B, <<
<<T S To: R 2 Kg
(94)
which proves the theorem.
Corollary B.1 (Depth requirement). For any P E PK,B. controller-augni Transformer that instruments every block needs at most L = K+B instrumented steps to simulate P in the limit. only every p-th block is instrumented, the required total depth is p (K+B).
Proof. Execution-trace length. By definition of PK,B, the discrete execution trace has length
T(P) S K +
(95)
where K counts primitive (non-loop-iteration) steps and the total loop-unrolling bound across all loops. In the discrete limit, each primitive (ASSIGN, ADD, COMPARE, BRANCH) corresponds to exactly one controller step, so the number of required controller steps is
= T(P) S K+B
(96)
Case 1: Instrument every block. each Transformer block is instrumented, every block executes exactly one controller step. Hence after L blocks we have carried out L controller steps and, by equation 96) can realize the entire trace for any P E PK,n with
depth = L < K+B
(97)
Case 2: Instrument every p-th block. Let be the total number of blocks and suppose exactly the indices p, 2p, 3p.... ... are instrumented. The number of controller steps executed by depth D
N(D)
L
(98)
To simulate the program we need N(D) L. From equation 98 this requires
- 2 L
PL
(99)
Choosing D PL attains N(D) L. Combining with equation 96 we obtain the sufficient depth
depth pL - p pK+BB
(100)
Equations equation 97 and equation 100 establish the claim: a controller-augmentted Transformer needs at most L = K+B B instrumented steps if every block instrumented, and total depth p (K B) if only every p-th block is instrumented.
(i) The margin Assumption B.3 can be relaxed to hold only on compact neighborhood the trace; continuity then suffices. (ii) If supervision operation types unavailable. entropy annealing and Leem typically induce low-entropy decisions, empirically shrinking the constants in Theorem B.2 (iii The bounds non-asymptotic in L for fixed [K+B) and show linear dependence on (T+** consistent with observed correlations between low decision entropy and length generalization in &Results.
25

---

Published as a conference paper at ICLR 2026
E TRAINING DETAILS AND HYPERPARAMETER SWEEPS
This section specifies all optimizer settings, schedules, and ablation grids used our experiments, and reports systematic sweeps over lemperature/entropy schedules, controller frequency, register configuration, and state-encoder choices. Unless otherwise noted, we keep the base architecture and datasets as in $4.3 and $4.1
E.1 OPTIMIZATION AND REGULARIZATION
Optimization details. We train with AdamW (81=0.9, 12-00.99,22-0000 32-0.98, c-100-) and weight decay 0.01 applied only to non-embedding weights (biases and LayerNorrmeembedddng matrices are excluded via per-parameter groups); gradients under mixed precision). The learning rate follows a linear warmup for 5,000 steps, then cosine decay floor of 10-5; the peak (pre-decay) learning rate is task-dependent (seq2seq vs. classification) and is swept as reported in Table 5 Unless noted, we use label smoothing 6=0.1 for sequence- level cross-entropy, and teacher forcing ratio 1.0 on synthetic seq2seq tasks. Batching reported in tokens: B=256 tokens for synthetic tasks and B=128 for DROP/RobustFill;; sequences are padded to the maximum length in-batch with padding masks applied in attention and loss. Training uses FP16 (AMP) with dynamic loss scaling; we avoid gradient accumulation by default but enable it if memory-consttrinee.
clipped
global norm of 1.0 (applied after unscaling
Controller schedules and auxiliary losses. The controller temperature follows exponential an- nealing schedule (t) = max (Tmin, The-") with (To, Tmin: 7) swept in SE.3 this reduces stochastic mixing operations training progresses and encourages sharpness к (used in (xx2) for COMPARE/BRANCH is increased linearly from to 20 over the first 30% of training steps by default, trading early smooth gradients for late-phase hard decisions. To further promote discreteness, the entropy penalty weight ramps linearly as Aem(() 0 - A over the first 20k steps (we sweep Atm), while the consisseey weight Asmm held constant (swept) to align the learned transition with the declarative primitive semantics. The total objective C Luss + AupLop + Asemm-yem + Xem Aspur where Lup (operation supervision) used only when ground-truth traces are available (O on non-synthetic tasks), and Lapur encourages near one-hot addressing and small writes; unless noted we set Aupur = 10-5, All schedule hyper- parameters and loss weights are selected via validation sweeps, and (when applicable) reported in
one-hot selections. Comparison
E. and Table 5
E.2 LEARNING- RATE SWEEP
We sweep the peak LR on Sorting (ID: n-80) and extrapolation (OOD: n=320). Means + std computed over 3 seeds, same data generators.
Table 5: Peak learning rate performance. Sequence accuracy (%) Sorting train
length (ID) and extrapolation (OOD). Cull-Transfommn p=2, =, d d,.644 (1.0, 0.1, 5x 5), Aen 5x10-3 Asem
(To, Tmin: 7) =
1x10-3
Peak LR
ID (n==0)
OOD (n=320)
1x10
96.8
89.4
2x10-4
0.3
90.9 + 0.5
3x10-*
97.4+0.4
91.6 + 0.7
5x10 4
97.1
90.7 0.8
8x10-4
96.5 1 0.7
88.1 + 1.1
Learning-rate analysis. 5 a clear interior optimum at 3x 10-4, which attains the best ID accuracy (97.4 + 0.4) and the strongest OOD generalization (91.6 + 0.7). Moving below this point (1 x 10 4 2x10-4) modestly reduces ID and, more noticeably. OOD performance (down to 89.4%); this is consistent with slower optimization and a lagging entropy anneal that leaves trollers too soft at convergence. Moving above the optimum (5x10 8 x 10 -4) maintains competi- tive ID but degrades OOD (to 90.7% and 88. 1%), indicating that aggressive steps accelerate training
26

---

Published as a conference paper at ICLR 2026
over-shoot the discrete regime, causing premature commitment and poorer length extrapolation. Variance (the =) also increases slightly away from the optimum, suggesting reduced stability. In practice, 3x 10 offers the most reliable ID/OOD trade-off, while sweeps should be biased toward 2x 2x10- 10-1 for nearby tasks.
E.3 TEMPERATURE ENTROPY SCHEDULES
We evaluate three schedules: SLOW (larger To, smaller 7), MEDIUM (default), and FAST (smaller Tmin, larger 7). Below we plot decision-entropy over epochs.
Op-selection entropy
0.5
10
20
30
40
50
170=1.5, Tuin Tmin= 00.2, 2=2x10-0)
Medium (1.0,0.1,5xl0 ")
Fast (0.8, 0.05, 1x10-1)
Figure Operation selection Operation-selection entropy over epochs under three temperature schedules.
Interpretation and schedule choice. Figure 2 shows that all three schedules monotonically duce operation-selecteeee selitiiinnseetiionn entropy, but their rates and final floors differ ways that matter for OOD generalization. The FAST schedule drives early discreteness (steep drop by epochs 5-15), which can yield quicker optimization but risks premature commitment to suboptimal branches; this explains its slightly higher terminal entropy (zs 0.18) compared MEDIUM. The SLOW schedule preserves ex- ploration longer but retains the highest final entropy (ss 0.21), which correlates with weaker length extrapolation (controllers keep mixing ops inference). The MEDIUM schedule reaches the "el- bow" around epochs 10-20, then hardens steadily to the lowest entropy floor (= 0. 0.16), which in our experiments aligns with the strongest OOD accuracy. Practically, we therefore adopt MEDIUM as the default; if training shows early- irrllyinnttaillitylook-ii increase or decrease 7 (slower anneal), and entropy remains high convergence, raise A or lower Tmin to encourage crisper decisions.
E.4 2D SWEEP: (TO, HEATMAP
We report AUL (area under accuracy-length curve) on Sorting across (To.7) E (0.8, 1.0, 1.2, 1.5) (2 10) x 10-5 with Tiin 0.1.
Annealing sweep analysis. The 4x3 grid in Fig. 3 exhibits broad "ridge" of strong AUL centered around (7==1.0, 7=5x10-5) with the peak AUL of 0.95; its immediate neighbors (1.2, 5x10-9) and (1.0, 10x10000000 10-5) remain high at 0.94, indicating a reasonably flat optimum. Row- wise (fixing To), AUL generally improves from 7=2x 10 to 5x10-5 and then dips slightly at 10x10-5 suggesting that too-rapid annealing can trigger premature hardening. Column-wiss (fix- ing 7). performance non-monotone in 7o moving from 0.8 1.0 helps (more initial exploration), but pushing to 1.5 harms AUL, consistent with controllers staying too soft for too long. The overall spread is modest (about 0.90-0.95). 0.95,, but the mid-range schedules clearly dominate: extreme settings underperform by ss 0.03-0.05 AUL. Practically, ing shows early lock-in, increase To (e.g.. 1.2) or decrease (e.g.. 2x10 5) if controllers remain overly soft near convergence, increase y (to 10x10-5) reduce Tmin- These trends align with the entropy curves: schedules that balance early exploration with timely hardening yield the best length-extrapolation AUI..
set m--10, 5=5x10-5 by default; if train-
27

---

Published as a conference paper at ICLR 2026
AUL (higher is better)
AUL
0.95
10
0.94
< (xx10 5)
0.93
5
0.92
0.91
2
0.9
0.8
1.0
1.2
1.5
TO
Figure 3: Heatmap of (To, 7) vs. AUL on Sorting. The interior region around (ro-1.0, 7-5x10-9 attains the highest AUL.
E.5 ENTROPY AND SEMANTIC WEIGHTS
We sweep Aen E {0, 10-3 3x10-3 5x10-3 10- and Asem E {0, 5x10-1 10-3, 2x 2x10-3 on Sorting (n=320) and BFS (n=30 + 90 extrapolation). Three consistent trends emerge. (1) Entropy weight: Increasing Aen lowers operation-selection entropy at convergence and improves OOD, but is non-monotone: >= leaves controllers too soft (higher entropy, reduced trace align- ment); 10-3 to 3x10-= helps. and =5x00 yields the best trade-off across both tasks, matching the entropy floors in Fig. 2 Pushing to =00== occasionally induces premature hard- ening (spikes in validation loss around the anneal elbow) and small OOD drops.
(2) Semantic weight: Turning on Asem reliably improves execution faithfulness: Aeem=10- 10- boosts trace align- ment and adds +2-5 points the longest lengths on both Sorting and BFS. Larger values (2x 10 3) begin to over-regularize, slightly reducing ID and offering no further OOD gains (consistent with werrconstraiiing the relaxed dynamics early training). best region is moderately wide: Aem 10-3,
(3) Interactions and robustness: The and X E 10x10-1 10-3] are within 0.5% of the peak OOD across seeds. We observe mild interaction: when Aent small, larger Asem partially compensates by anchoring transitions; conversely, with Amm large, keeping Asem at 10-3 avoids over-constraining. Practically, we set =5m==5x10-3, >sem=10 10-3 by default. con- trollers stay soft late in training (entropy plateau), increase Aent reduce Twin if training exhibits early lock-in or brittle branches, reduce Aent or lower y (slower temperature anneal), and cap Agem at 10 As before, =0 on non-synthetic tasks.
E.6 CONTROLLER FREQUENCY P
We vary the insertion frequency p {1, 2, 4) (instrument every p-th block) and observe smooth accuracy-compute trade-off consistent with Table instrumenting every other block (p=2) retains 95% of the p=1 gains at only x 60% of the overhead. while p=3 remains competitive but loses 3 absolute points at the longest lengths. Concretely, the relative FLOP overheads for (,=1,2,3) 4} are (+10.1% +6.4%. +4.4%, +3.2%} (Table 3 mirroring throughput changes. From a depth standpoint, Corollary B. implies that instrumenting every p-th block multiplies the total block depth needed to realize
B-sttep program by i.e., depth Rs p (K+ thus, for large K+B or shallow models, smaller p is preferable. We also notice mild interactions with annealing: larger p (sparser controllers) benefits from slightly slower temperature schedules (smaller or larger T0) to allow each controller update tolerates faster anneals. In practice, we adopt p=2 as the default for strong OOD performance at modest cost; switch to p=1 for hard algorithmic tasks with tight step budgets (large K B) or
integrate evidence across more token updates; conversely, p=1
28

---

Published as a conference paper at ICLR 2026
when depth is plentiful, and to p=3-4 when compute/latency is the primary constraint and a ~ ~ point loss at extreme lengths is acceptable.
E.7 REGISTERS (Nr, dr) AND STATE ENCODER
We grid-sweep N, E {4, 12). d, E {32, 64, 128), and compare & E (mean, max, attn] (atten- tion pooling with learned query). Increasing controller capacity from (N, to (8, 64) yields consistent OOD gains +2-4 points on Sorting and BFS, with stable training and only a modest compute increase; pushing further (12, 128) shows diminishing returns (typically < +1 point) and adds ~ 2% FLOPs/latency due to larger state encoders and write-backs. We also observe that very small controllers (c.g.. 4 x232) tend retain higher terminal entropy (softer decisions), while very large ones (12x 128) can overfit early soft traces unless Am or Aspr is slightly increased. Across pooling choices, attn > mean > max: attention pooling improves trace alignment (cleaner opcode sequences) and boosts OOD by +1-2 points, presumably by letting the controller attend to task-relevant tokens rather than averaging signal across the sequence. Interactions are mild but notable: sparser insertion (p=3) benefits more from attn (since each controller update must inte- grate evidence gathered over more token updates), whereas dense insertion (p=1) narrows the gap between attn and mean. Taken together, and consistent with our main configuration, (N, d,,64, d=aatn) as a strong default; increase to (12, 128) only if the task requires longer multi-register programs and the small overhead is acceptable, or decrease (4, 32) for tight latency budgets with small OOD cost.
recommend
E.8 IMPLEMENTATION NOTES
Controller-side MLPs for the branch arms /eten/ Weise are initialized with Xavier-unifoma (fan- in/fan-out) and small output-scale; opcode logits are biased favor NOOP at the start (and write- strength À is initialized near 0) to prevent spurious early writes, while read/write addressers are started near uniform so that early gradients explore registers before annealing sharpens selec- tions. For late-stage hard decisions, we optionally enable straight-through (ST) estimator on y() after epoch EST E [0.6,0.8]: the forward pass uses one_ one_hot(arg max x0)) and the back- ward pass uses the Gumbel-Softmax gradient ("stopgrad (yhard M/sott + Ysoft " trick). We gate this switch by validation stability (no recent degradation) and an entropy threshold to avoid pre- mature hardening: with Aen active, ST further reduces terminal entropy without harming opti- mization. For reproducibilityy we fix seeds {1, [1 across Python random, NumPy, and Py Torch cwwpeennnddddddnn.nddoeoe cudnn. wwwedeeeiestieee110 PRNGs; set torch, backends. per-worker seeds via the Datalo to ensure identical shuffling, padding, and batch composi- tion. We also seed the Gumbel sampler used in the controller, tie dropout masks to the global seed, and record all sweep hyperparameters (LR, 7-schedule, k-schedule. À weights) in config files for exact reruns. Note that some GPU kernels (especially under AMP/FP16) may remain non- deterministic certain hardware/driver combos; when exact reproducibility is critical, we enable torch. use _ueeuusteeeueee www.ddetemmiiiscc_aggooooo thms (True) and fall back to deterministic ops at small through- put cost.
and benchmark=False; and propagate
0.1 Across sweeps, medium temperature annealing (70=1.0, Tmin 001, 7-551001))) moderate entropy
regularization > 10-3 semantic consistency Ave 10- controller frequency p=2, registers (8,=8, d4.)4), and attention pooling ø=attn deliver the most robust ID/OOD tradc-off with 5-10% compute overhead.
F SYNTHETIC GENERATORS AND PROTOCOLS
This appendix details how we generate the synthetic corpora used in 44.1 including pseudocode, seeds, distributions, serialization.
F.1 SORTING AND SUS-OF-LIST
Task definition, sampling, and reproducibility. For given sequence length n, either in- distribution (ID) n E {10,20, 20, 40, 80) or oft-o--diitriiutioon (OOD) n E {160, 320), we draw an in- teger list x (.1 - - pp) with i.i.d. entries x; Uniform 999) with replacement, duplicates
29

---

Published as a conference paper at ICLR 2026
common and stability is essential for determinism. The two targets are: (i) Sorting, where is the stable nondecreasing sort of x (ties preserve the original index order, i.., items with equal value output in increasing order of their input positions), and (ii) Sum-of-List, where y Σ z; a single integer. For each n. we construct Train/Val/Test splits totaling 200k /ook//20//20 examples across the four training lengths (evenly allocated, i.e., 50k/5k/ /5k per length), and generate OOD sets at E {160, 320} using the same protocol but only for test-time evaluation. Randomness is controlled via recorded seed per file, with default corpus seeds {100, 101, 1023 producing three independent datasets; we use disjoint seeds for Train/Val/Tess and for each length bucket to pre- vent leakage. Sequences are serialized with explicit start/end sentinels and integer tokens (e.g., 
) without subword decomposition to avoid confounds; padding and attention

 x_1 masks are applied per batch. This setup induces controlled length shift (ID - OOD) while keeping the value distribution fixed, enabling clean measurement of length extrapolation under stable sorting semantics and duplicated inputs.
Pseudocode:
SORT AND_SUA SUM
Algoritha
N, seed):
set_rng(seed)
D empty list
j N..N:
x random_ inrddom_int(C (0, 999) for i 1..n
with replacement
6
y-sort stable_sort(x)
# nondecreasing, stable
y_sum sum(x)
append (x, y_sort, y_sum)
return D
F.2 GRAPH TRAVERSAL (BFS)
Graph model and BFS task. We generate Frdös-Rényí graphs G G(n, p) with = 0.2, using {10, 20, 40, 80} for traininglevaluation and OOD lengths {160, 320} for test only: random- ness is controlled via recorded seeds. To ensure a unique BFS target, we enforce connectivity: if sampled graph is disconnected,
resample up 10 times; failing that, force-connect compo- nents by adding a minimal set of bridging edges in deterministic way (e.g., connect the smallest-id node of each component to the smallest-id node of the next component), after which adjacency lists are canonicalized. BFS runs from a designated start node 8 ~ Uniform that enqueues each node only at first discovery, and a fixed
with standard
FIFO-queue semantics, visited neighbor ordering: for each dequeued vertex, its adjacency list is sorted in ascending node id, and newly discovered neighbors in the same layer are enqueued increasing id. This tie-breaking guar- antees deterministic visitation even under duplicate frontier discoveries and yields unique target sequence (v,,..., U.) of visitation order. For serialization, we emit graphs as sorted edge lists (or adjacency lists) with explicit delimiters and output the BFS visitation as a token sequence; padding and attention masks are applied per batch to avoid length-induced biases.
Pseudocode:
GEN_BFS PAIRS
Algorithm LEN_BFS_PAIRRSE N, p, seed): set_rng(seens D = empty list for in N..N: G sample_Gnp(n, p) not (s_:o:necteddG):
G = (on)ecteec)))) random_int(1, n) for each sort(ad)[u))
% resample or add edges to connect components
ascending neighbor IDs FIFO queue
10
order BFS start=s)
30

---

Published as a conference paper at ICLR 2026
11
D.append( (G, order) )
F.3 SERIALIZATION AND TOKENIZATION
Canonical serialization and vocabulary. Each example is serialized on a single line (input; target) using chitespace-canonica format enable deduplication and hashing. Integers [0, 999] that appear as sequence elements are encoded atomically as <dxxx> (zero-padded), Sort-
ing inputs use 
 <len> </len> <seq? <d012>
<d999> </seq? 
 with targets </sorted> </t>. Sum targets use <t> <sum> Y </sum>> </t>
<t> <sorted> <d...>
where is the (possibly >999) base-10 integer sum without padding or tags. Graph inputs begin with a header <g>
<n> n </n> <p> 0.2 </p> <start> s </start> followed by <edges> (ul v1) (u2 v2)
</edges> where undirected pairs are sorted (u < v) and listed in </bfs> </t>. Numerals that are not part of <dxxx> token are emitted as plain base- 10 integers (no commas). We collapse runs of spaces and trim ends to forbid stray whitespace, and use this canonicalization for hashing. The vocabulary comprises the specials 
 
 <t> </t> <seq? </seq? <sorted> </sorted> <sum> </sum> <g? <n> </n> <p> </p> <start> </start> <edges> </edges> <bfs> </bfs>, plus the family of <dxxx> tokens for 000-999; we apply no subword splitting to <dxxx>, Sums are not constrained to <dxxx> tokens; they are serialized as plain integers.
lexicographic order; BFS targets are <t> <bfs> v1 v2
F.4 PYTHON FOR GENERATORS
import random, hashlib, json
def set__rn(seeed
random. seed (seed)
def ren_sort_and_sune N, seed):
set_rng(seed)
data []
for in range(N):
x Trandom.rrndintto 999) in range(n)]
y_sort sorted key=lambda i: (x[il, i)) # stable sort indices y_sort_vals = [x[i] for in y_sort]
12
= sum(x)
aata.appendd(x. y_sort_vals, y_sum))
return data
16
def (ns_order(n, edges, s):
adj = [[] for in range(n)]
for u, edges:
u v: continue
21
u, v (u, v) if u else (v, u) adjjul)appenddv); adj[v] (ppenddu)))
for range(n):
adj[u].sort()
25
seen [False]*n
26
[s]
seen[s] = True
order
29
head 0
30
while head < len(q):
q[head]; head += order. for in [dj[ if not seen[v]:
34
31

---

Published as a conference paper at ICLR 2026
35
seen[v] True q.append(v) return order
39
def serialize_sorie _orrtexs y_sort):
=
for i in x])
ys "'join(("<<(ii:ddd""" for in y_sort])
42
input_str <" <len> (len(x)) </len> <seq? (xs) </seq> 
"
43
target_str = f'<t> <sorted> {ys} </sorted> </t>"
44
return input_str, target_str
def serialize_sume y_sum):
47
# Inputs remain 3.000>..</d9999 form
48
xs = "_join(f"<d(i:0dd)> for i in x)
49
= str[int(y_sum
input_str f"
> <len> (len(x)) </len> <seq? (xs) </seq> 
" target_str f"<t> <sum> (ys) </sum> </t>"
51
52
return input_str, target_str
53
def serialize_graph(n, p, s, edges, order):
edge_str " ". ".join([f"(CC ([u] (v))" for (u, v) sorted(((min(u,)) max(u, v)) for (u,v
) in edges])[)
57
inp r"<g> <n> [n) </n> <p> {p) </p> <start> (s} </start> <edges> (edge_str) </
edges>"
tgt f"<t> <bfs> (' (mssssnm order))) </bfs> </t>" return inp, tgt
59
60
def sha256_line tgt):
", return iashlib.sha256 ancode("utf-8")) hexdigest()
f'(inn)\t(ttt)" (s
G BASELINE DETAILS AND COMPUTE ACCOUNTING
This appendix spells out how we matched baselines, how we compute/measure FLOPs and through- put, and the exact Fusion-in- -D-eodeer (FiD) configuration used DROP. Our goals are (i) archi- tectural parity, (ii) transparent compute accounting, and (iii) apples-to-apples token budgets.
G.1 PARAMETER MATCHING
Unless otherwise noted, all Transformer-like models share the same core hyperparameters:
T==12, D=512, H=8, D042048 (= 4D).
Embeddings and layer norms follow Vaswani al. (2017). For our controller model, controllers are inserted every p=2 blocks (i.c., 6 controller- augmente blocks). Registers: N,=8, d,,64;; controller summary d, 128; optional memory disabled by default (N,.-0).
Table 6: Architecture parity. Shared backbone unless specified. "Ctrl freq." denotes controller
insertion frequency p.
Model
D
H
Un
Ctrl freq.
Notes
B1 Transforner
512
8
2048
Encoder-only oll for Sorting/BFS; Encoder Decoder when needed ACT halting: expected depth matched
B2 Universal Transformer
512
8
2048
B3 Compressive Transformer
12
512
8
2048
Compression ratio c=3; compressed memory width D As in Veličkovic et al. 2202
B4 CI CL.S Transformer
512
8
2048
2023 sstyle implementation External contceller/menrry adapter loken space Our controller hut woft op mixture (no Gumbeel
B5 NPUDNC
512
8
2048
B6 Soft Modules
12
512
8
2048
p=2
Ctrl- Transformer (ours)
512
8
2048
2=2
N,=5, d,= d,=128
32

---

Published as a conference paper at ICLR 2026
Architecture parity. Table 6 holds the backbone fixed across baselines (L=12, D-512, H=8, Dx=2048=4D) so that differences reflect only control mechanisms. The plain Transformer (B1) is encoder-only for Sorting/BFS and decoder when needed. The Universal Transformer (B2) uses ACT with expected depth matched to 12. The Compressive Transformer (B3) applies com- pression ratio c= with compressed memory width D. The CLRS-AIg. Transformer (B4) follows a Veličković et al. -stye implementation. The NPI/DNC baseline (B5) relies on an external controller/memory with an adapter back to token space. The Soft Modules variant (B6) inserts con- trollers every p-2 blocks but uses a soft operator mixture (no Gumbel), Our Ctrl- Tranffranfrrmer also inserts controllers every p=2 blocks, with registers N,=8, -64 and controller summary d, 1221; optional external memory is disabled -N). "Ctrl freq." denotes the insertion frequency
Parameter counts (closed form).
Ignoring biases and norms, per Transformer block:
ParamSntock 3D (Q/K/V) D (out proj) + 2DDw (FFN) 4D + DD
With Dx=4D;
ParamSblock x 12D2. For D=512: 12D = 12 262, 144 = 3.15M; across List = 12: 8M. Embeddings add [V] (small for our synthetic vocab).
Controllers add (per instrumented block) small MLPs and projections. Let d, Nm denote the controller state width. Then a typical per-block parameterization
ParamScal
dg|O] + dgN, d,N, d,d, + MLP (empirically ~ 0.4-0.6M).
With p=2 (controller in every other block; 6 placements for a 12-block network), this yields a total overhead of roughly 2.4-3.6M parameters.
Compute overhead (FLOPs). Let be the instrumentation period (only every p-th block car- ries the controller). A standard Transformer block costs O(D²) per token (FFN-dominated), The controller adds two broadcasts (token+>state) and small state MLPs, costing O(Dd, + d)) per in- strumented block. Thus the dimensionless relative FLOP overhead is
overhead RV
DjDd, + ord?
1 -
(a
+
Chase D2
P
D2
with task- indeeeenennn constants a, 3 = ((). In our configs d, < D and FFN dominates, so the quadratic term governs and we report
1 dg 2
overhead
2p D
which matches the observed 5-7% for p=2.
G.2 FLOPs ESTIMATION AND THROUGHPUT
We report per-sequence forward FL.OPs at sequence length T using standard dense-kernel counts (multiply as 2 FLOPs). For a self-attention block:
+ + D
2TDDs 2T 4TD2 2TDD
Q/KV and proj
QK +VV
FFN
With Dx=4D: 27 + 12TD2 For our 1-12, 51212 D-512, T-128:
FLOPs/seq (2T D2T 12 (2 1282 512 12 128 512?) x 5.03 G.
Controller overhead per instrumented block is dominated by (i) é pooling (O(TDd,)). (ii small projections (O(d,d,)) )), and (iii) light MLPs on the compact state, together yielding a measured overhead of +5- 7% FL.PPs for p=2.
Throughput measurement. We measure tokens/s (training, forward+backward) as wall-clock throughput on a single A100-40GB, PyTorch 2.x, CUDA 12, FP16 (AMP), batch sized to satu- rate memory. Results are averaged over rums; variance reflects kernel scheduling and autotuning. Because absolute tokens/s depends on hardware and kernel fusion, we emphasize relative differ- ences.
33

---

Published as a conference paper at ICLR 2026
Table 7: Params, analytical FLOPs per sequence (T'=128), and measured training throughput (to- kens/s). ttd over 3 runs; throughput measured with identical token budgets and mixed preci- sion.
Model
Params (M)
FLOPs/seq (G)
Throughput (tok/s)
B1 Transformer
38.2
5.03
20,800 + 600
B2 Universal Trf. (ACT)
38.4
5.35
19,700 - 500
Compressive Trf.
41.0
5.60
= 19,100
B4 CLRs-Alg. Trf.
38.7
5.05
20,200 + 550
B5 NPI/DNC
40.2
6.10
17,400 + 450
B6 Soft Modules
41.3
5.33
19,600 = 400
Ctrl-Transe Tin--Transformer (ours)
41.3
5.35
19,400 = 400
Relative Bl: Ours +8.1% params, +6.4% FLOPs, -6.7% throughput
Latency vs. sequence length (inference). We also report per-sequence latency (ms) at batch size 1 to visualize the T attention scaling and our small additive overhead.
Notes on FLOPs vs. wall-clock. Analytical FLOPs correlate with, but do not fully determine, wall-clock time because kernel fusion, memory bandwidth, cache behavior, and launch overheads all matter. In Fig. 4 both B1 and Crrl-Transformer follow the same near-quadratic trend in sequence length T. while Transfrrasormmer incurs small, roughly constant relative overhead of & 7-8% across E {64, 128, 256, 512} at batch= 1 (e.g.. +0.2 ms at 7-64 and +1.2 ms at 7-512). This consistent with additional controller operations that do not alter the 722 term. To make efficiency comparable across hardware and implementations, we therefore report both analytic FLOPs and empirical lateneythrrouhhut.
G.3 FTD-T5 CONFIGURATION ON DROP
We include FiD primarily as reference architecture. To maintain fairness on DROP: No retrieval. DROP provides the question paragraph; do not add external retriever. Same token budget. We cap total input tokens equal to the Transformer budget. If the paragraph exceeds the limit, we use non-overlapping chunking into M spans such that _ [z((m) equals the baseline budget; the question prepended to each chunk. Model size. We use a small T5 settameter-mttcher within +5% of B1) to avoid capacity advantages. Decoder depth is adjusted to match parameter targets. Fusion. Standard FiD fusion in the decoder; no cross-chunk attention outside the decoder fusion step. Training. Same optimizer/sccecuuee as { early stopping dev F1; identical preprocess- ing and answer normalization. Token-budget sanity. We log, per example, the effective tokens consumed by FiD after chunking: mean and 95% CI match the baseline token budget (differences < 1%), ensuring comparable compute. Fairness checks. For all models, we report: (i) parame- ter counts, (ii) analytical FLOPs/seq T2- (iii) training tokens/s, (iv) inference latency T', and (v) identical token caps. This bundle avoids confounding from capacity or budget mismatches. Baselines share the same backbone; our controller adds 8% parameters and 6-7% FLOPs with a corresponding ~ 6-8% throughput drop, consistent with the modest overhead claimed in the main text. FiD is configured without retrieval and with the same token budget, serving as a strong but fair reference on DROP.
H EXTENDED ALGORITHMIC TASK RESULTS
This appendix expands the main-text results with full per-length tables, nonparametric uncertainty, and length-accuu curves. Unless otherwise noted, all numbers aggregate nine runs per config- uration (three corpus seeds from E three model seeds), and report (i) the mean, (ii) 95% bootstrap CI with 5,000 resamples, and (iii) Wilcoxon signed-rank p-values across seeds when comparing each baseline model the longest evaluation length.
34

---

Published as a conference paper at ICLR 2026
16
14
Latency (ms) batch=1
12
10
8
6
4
2
128
256 Sequence length T
512
- B1 Transformer
Ctrl- Transformer (ours)
Universal Transformer
B3 Compressive Transformer
B4 CLRS-Alg. Transformer
B6 Soft Modules
Figure 4: Inference latency vs. sequence length. Our model follows the same 722 trend with a small
additive overhead.
H.1 SIGNIFICANCE AND UNCERTAINTY
For each configuration and length we run 3 x 3 = 9 trials (3 corpus seeds x 3 model seeds). We report 95% bootstrap CIs by resampling the 9 per-run accuracies with replacement (5,000 draws: percentile CI). For pairwise significance at the longest length (Sorting/Sun/BBS: n=320), we apply the Wilcoxon signed-rank test on the 9 paired runs (baseline vs. ours) and report Holm-corrected p-values across baselines.
H.2 SORTING
Across list lengths all models are near-ceiling on ID n=200, but accuracy decreases as grows, with our Ctrl-Transformer consistently strongest and gaps widening out of distribution: at =
(ID), ours reaches 97.4% (96.9, 9..97 versus B6 93.0% [92.3, 93. 7]. B2 90.6% (89.8, 91.5], and B1 85.3% [84.0, 66.6]: at 160 (OOD), ours attains 94.2 2% [93.5, 5, 94.9) versus 78.8 78.8/72.5/61. /61.4% for B6/B2/B1: and at ==320 (OOD), ours delivers 91.6% [90.8, 92.4) versus 66.4/51.4 /3..6% with
35

---

Published as a conference paper at ICLR 2026
Table 8: Sorting, sequence accuracy (%) vs. list length n. Means and 95% bootstrap confidence intervals over runs. Wilcoxon p-values are computed vs. our model at n=320.
Model
n= (ID)
80 (ID)
n= 160 (OOD)
== (OOD)
B1 Transformer
99.1 198.8,98 99.3]
85.3 [84.0, 86.6]
61.4 [60.0, 62.9)
38.6 [36.8, 40.4]
vs. ours @320
10
B2 Universal Trf.
99.0 (98.8, 99.2]
90.6 (89.8,91 91.5
72.5 [71.2, 73.8]
51.4 [49.6, 53.1]
@320
10 66.4 [65.1, 67.8)
B6 Soft Modules
99.0 (98.7,90 99.2]
93.0 (92.3, 93.7]
78.8 [77.7, 79.9]
vs. ours @320
< 10
Sirl-Tranffomer (ours)
99.1 98.9, ,, 99.3
97. 99..9, 97.9
94.2 198.5, 94.9]
91.6 190.8, 92 .4
Wilcoxon tests at n=320 showing p<10-3 against each baseline. Confidence intervals are 95% bootstrap over runs; the advantage of our model larger remains highly significant.
H.3 SUM-OF-LIST
Table Sum-of-List, exact numeric match (%) vs. length n. Means and 95% bootstrap CIs over runs. Wilcoxon p-values ours at n=320.
Model
n= (ID)
T2 =z=80 80 (ID)
n= 160 (OOD)
Yz=220 (OOD)
Transformer
99.6 [99.5, 99.7]
88.9 (88.0, 89.9]
63.6 [62.2, 64.9]
41.1 [39.5, 42.8)
vs. ours @320
< 10
B2 Universal Trf.
99.6 (99.5,99.1]
92.7 [92.0,90 93.55
75.5 [74.4, 76.7]
56.0 [54.4, 57.6)
vs. ours @320
< 10
Soft Modules
99.6 [99.4,90 99.7]
95.2 [94.6, 95.8]
83.8 [82.9, 84.8]
72.8 [71.5, 74.1] < 10
ours @ 320
Ctrl-Transte Trrr-Transfrmer (ours) 99.6 99.55, 99.7]
98.1 197.7, 7, 98.5]
95.6 [95.0, 96.2]
93.8 [93.1, 94.4]
Across lengths n, all methods are at ceiling on ID (n=20;; R 99.6%), but accuracy declines with for baselines while our Ttrl-Transformer maintains high performance deep into OOD: at n=80 (ID) we obtain 98.1% (97.7. 98.55] vs. B6 95.2% [94.6, 95.8]. B2 92.7% [92.0, 93.5], and B1 88.9% [88.0, 89.9]; ; at n= n=160 160 (OOD) we reach 95.6% [95.0, 96.2) vs. 83.8 83.8/75 /3///555//33 5/3.6% (B6/B2/B1) 1); and at n=320 (OOD) we deliver 93.8% (93.1 94.4) vs. 72.8/56.0/ 06..0/41..% 1%, with
Wilcoxon tests n=320 showing p<10- -3 against each baseline; 95% bootstrap Cls are computed over 9 runs
H.4 GRAPH TRAVERSAL (BFS)
Table 10: BFS: sequence accuracy (%) vs. graph size n (nodes). Train sizes n E {20, 40, 80); OOD sizes E (160, 320). Means and 95% bootstrap CIs over runs. Wilcoxon p-values vs. ours at
n=320.
Model
2= (ID)
1==10 (ID)
160 (OOD)
320 (OOD)
BI Transfomer
98.2 (,7.8,98.8 98.6]
92.1 [91.1, 93.0]
86.3 [85.2, 87.4]
71.4 [69.9, 72.9]
57.9 [56.0, 59.8)
vs. ours 320
10
B4 CLRS-AIg. Trf.
98.6 [98.3,90 98.9)
93.8 [03.0, 94.7]
90.2 [89.3, 91.1]
77.8 [76.5, 79.
65.1 [63.4, 66.8)
ys. ours 320
B6 Soft Modules
98.6 [98.3, 98.9)
94.6 [93.8, 95.3)
92.1 [91.2, 93.0]
83.5 (82.4, 84.6]
74.2 [72.7, 75.7]
vs. ours 6 320
10
Trr-Trransormer (ours) 98.7 (98.4, 99.0)
95.3 (94.6, 95.9)
95.3 (94.7, 95.9)
90.8 (90.0, 91.6)
83.5 |82 4, 84.6)
Across graph sizes n, all models are near-ceiling at n=20 (ID), As n increases, the gap between Tran-rrasforrrrr and baselines widens, matching the trend in Table [ At n-80 (ID), we ob- tain 95.3 95.3% 3% (9.77, 95.9) vs. 92.1/90.2 86.3% (B6/B4/B1). Under extrapolation, the gap becomes reach 90.8 [90.0, 91.6) vs. 2.5/77.8/71.499 and n=3200 (OOD) Wilcoxon signed-rank tests at 3==220 yield 1<1003 against each baseline; 95% bootstrap CIs computed over 9 runs.
larger: at n= 160 (OOD)
we achieve 83. 5% (82.4, 84.6) vs.
36

---

Published as a conference paper at ICLR 2026
H.5 ACCURACY-LEDS CURVES AND AUL
We visualize length-accuracy curves for all three tasks. Shaded regions denote the area under the curve (AUL) for our model; baseline AUL values are reported in the main text and mirrored here.
Sorting
Sum-of-List
BFS
100
100
100
Accuracy (%)
Accuracy (%)
Accuracy (%)
80
80
80
60
60
60
40
40
40
20
80
160
320
20
80
160
320
20080
160
320
Length n
Length n
Length n
B1
B2
B6
Ours
Figure 5: Accuracy vs. length for Sorting, Sum-of-List, and BFS. Shaded region indicates AUL for our model (higher is better).
Methodology, plots, and takeaways. In Fig. 5 accuracy- curves for Sorting, Sum-of-Liss and BFS show that baselines degrade as n grows while our controller-augmentee Transformer main- tains substantially higher accuracy with noticeably shallower OOD slope; the shaded band under our curve visualizes AUL (larger area indicates better length robustness). For each configuration and length, we compute accuracy per run and then (i) construct bootstrap 95% CIs from 5,000 resamples of the 9-run bag (balanced across corpus/model seeds), reporting the 2.5/97.5 percentiles, and (ii) run Wilcoxon signed-rank tests that compare per-seed accuracies each baseline against our model at the longest length CoottingSSum/BFS/ n=320), applying Holm correction across tested baselines; reported p- values are post-correction. We treat n E {20, 40, 80} as in-distribution for BFS and E {160, 320} OOD; for Sorting and Sum-of-List we evaluate at E {20, 80, 160, 320} and regard n Z 160 as OOD. Taken together with the figure, these procedures show that our model preserves near-ID accuracy deep into OOD lengths, with tight CIs and strong nonparametric sig- nificance at the longest lengths; the consistent gap to Soft Modules further supports the need for discrete control to achieve robust generalization rather than gains attributable to capacity token- budget artifacts.
I EXTENDED RESULTS ON SYMBOLIC QA AND PROGRAM SYNTHESIS
This appendix expands the DROP, RobustFill-style program synthesis, and Mathematics (arith- metic) evaluations with per-category breakdowns, compositional analyses, and per-template accura- cies. Unless stated otherwise, results aggregate nine runs per configuration (three corpus seeds x three model seeds); we report means and 95% bootstrap CIs (5k resamples).
1.1 DROP CATEGORY BREAKDOWN
We follow the official DROP taxonomy and report FI/EM on four categories that require discrete reasoning:
Addition, Count, Compare, and Date. Compared to parameter-matched Transformers (B1) and Universal Transformers (B2), our controller-angmentte model shows consistent gains, es- pecially on Addition and Compare, which directly benefit from the learned ADD and COMPARE/BRANCH primitives.
Per-category analysis. Across DROP categories, our Ctrl-Transformer attains the highest F1/EM with non-overlapping 95% CIs relative to the strongest baseline (B2) in every column. The
37

---

Published as a conference paper at ICLR 2026
Table 11: DROP per-category results. FI EM (%) with 95% bootstrap CIs over runs.
blodd
Addition
Count
Compare
Dute
Trausooreer
76.1 445,20
71.0 00.6 72.4
88. 8,808
70
Si
51
76. To. 78.7
27.0 75.0
largest absolute gains appear on Addition (+6.4 +7.8 EM vs. B2: x +8.0% R +10.5% rel- ative) and Compare (+5.5 / +5.5; +6.6% +6.9%), aligning with the controller's discrete arithmetic/relational branching. Count also benefits (+5.4/ = +6.4% +6.6%% consistent with low-entropy ADD/ASSIGN traces. Date improves by +6.3 +6.3 (= +7.9% +8.3%) despite no specialized date operator, suggesting that explicit control flow helps compose comparisons and offsets. Gains are mirrored in EM, indicating improvements are not an artifact of softened span scoring. Wilcoxon signed-rank tests our model are significant at p < 0.01 after Holm correction across baselines, reinforcing the visual gaps in Table
1.2 ROBUSTFILL COMPOSITIONAL SPLITS
We partition the test set by program length (number of DSL primitives the minimal consistent program): Simple ($2 ops), Medium (3-4), and Hard (5-6). We also summarize DSL-operator coverage and report accuracy vs. program length (Fig.
Compositional exact match (%). Weights across splits are approximately 0.4 (Simple), 0.4 (Medium), 0.2 (Hard). Numbers align with the overall scores in the main text.
Table 12:
Compositional exact match (%) by program length on the DSL benchmark. We partition by the number of DSL primitives in the minimal consistent program: Simple (< 2), Medium (3-4), Hard (5-6). Split weights are ze 0.4/0.4/0.. /0.2, so row-wise averages aligm with overall scores. Values mean [95% CI]; bold marks the best per split. See Fig. 6 for accuracy length and operator
coverage.
Model
Simple (5 2)
Medium (3-4)
Hard (5-6)
B1 Transformer
78.0 [76.6, 79.4
67.0 (65.5, 68.5)
50.0 1.8.1,51.9
Soft Modules
81.0 [79.7, 7, 82.3]
70.0 68.6,, 71.4]
55.0 53.2, 56.8
Tirl-Tranfformer (ours)
82.0 [80.7, 83.3]
72.0 70.6, 73.4]
60.0 58.2, 61.8)
DSL coverage. Most frequent primitives: CONCAT (33%), SLICE (28%), REPLACE (21%), LOWER/UPPER (10%), TRIM (8%); common two-op chains: SLICE>REPLACE (13%), CONCATOSLICE (11%). Our largest relative gain (+7-10 pts) appears on SLICEOREPLACE and CONCATOSLICE reflect- ing benefits from explicit COMPARE/BRANCH and stateful registers in multi-step string rewrites.
RobustFill length sensitivity. As shown in Fig. exact-match accuracy declines with program length L. but our Tran-Trranformer degrades more slowly (84 + 60, -24 pts) than B6 (83 + 55, -28 pts) and B1 (80 + 50, pts). The advantage is small at short programs (I.=1-2; +1-2 pts) and widens with composition depth (e.g., L-6: 60 vs. 55/50 for B6/B1), indicating reduced compositionality cliff for our model.
I.3 MATHEMATICS DATASET TEMPLATES
We report per-template accuracy on a subset the Mathematics dataset (arithmetic). Templates are grouped by operation type and difficulty (carty/borrow mixed operators). Our controller model shows larger gains on long-carry addition and mixed-operator templates, mirroring algorithmic length generalization.
Failure modes (heatmap summary). Error analysis reveals two dominant failure clusters for baselines: (i) carry/borrow propagation across long spans (leading to off- off-by-10 errors), and (ii) precedence violations in mixed arithmetie; our model reduces both clusters by favoring low-entropy
38

---

Published as a conference paper at ICLR 2026
90
80
Exact match (%
70
60
50
40
2
Program length L
Transformer
Soft Modules
Cinl-Transfommrr (ours)
Figure 6: RobustFill accuracy vs. minimal program length positionality cliff larger L.
Discrete controllers reduce the com-
Table 13: Per-template accuracy (%) the Mathematics (arithmetic) subset. Templates are grouped by operation and difficulty (carry/borrow, mixed operators). Trl-Trrasffrmer achieves the best re- sults overall, with the largest gains long-carry addition and mixed-operator templates, consistent with length generalization. Values are mean [95% CI]; bold indicates the best per row.
Template
B1 Transformer
B6 Soft Modules
Tral-Traasformer (ours)
Addition (2x (333-dig no carry)
88.2, 90.6
90.1
89.0, 91.2)
92.8 91.9, 93..7
Addition 4-dgii,,,, with carry)
72.0
70.5, 73.5
75.6
74.1, 77.1
81.7 80.3, 83.1
Subtraction (borrow)
74.5
73.0., 76.0
76.1
77.5)
79.9 78.6, 81.2
Multiplication (2-digit x 2-digit)
70.8
69.3, 72.3
72.6
71.2, 74.0
76.4 75.1, 77.7
Mixed arithmetic (3 ops)
63.2 61.6, 64.8
66.9
66..3, 68.5
71.0
69.5, 72.5
Order of operations w/ parentheses
68.7
67.2, 70.2]
71.0
(69.6, 72.4
74.8
73.5, 76.1
COMPARE BRANCH traces and stable accumulator updates (a full heatmap over error types X tem- plate families is included in the supplemental repository). Summary. Discrete controllers improve DROP categories that require numerical comparison//ddditon,, reduce the compositionality cliff in RobustFill at higher program lengths, and deliver stronger accuracy on arithmetic templates that demand multi-step carry/borrow or precedence handling, consistent with the semantics and finite- temperature analysis in the main paper.
J INTERPRETABILITY ANALYSES
This appendix probes whether the controller executes legible plans. We (i) align operation traces ground truth on synthetic tasks, (ii) train lightweight probes for register roles and quantify linear decodability and mutual information (MI), (iii) run knockout and randomization tests to measure causal reliance on controllers, and (iv) visualize the controller state across steps with t-SNE/UMAP.
J.1 TRACE ALIGNMENT EXAMPLES
Protocol and results. For tasks with step-type annotations (Sorting, BFS), generate ground- truth step sequence (y)()) | by instrumenting the reference algorithm, where y (()() E O, and the model produces per-layer operation posteriors (() A 101-1 The alignment score is
Align
Σ
arg max, ( = y y-1}
(101)
L
1=1
39

---

Published as a conference paper at ICLR 2026
We report Align on held-out sequences and stratify by imput length On Sorting, Align 92.3% at n=80 and 90.5% at n=320; on BFS, 88.77 7% at train sizes and 85.2% OOD. Soft Modules (no discrete selector) achieve 61.5%/57. .%, indicating that hard selection sharpens step semantics, Fig- ure visualizes operation usage over depth/time.
Operation usage intensity (soft counts)
Usage
0.8
2
3
0.6
5
Layer
6
7
0.4
9
10
0.2
11
3
5
9
11
13
15
17
19
Time step
Figure Heatmap of per-layer operation usage over time on Sorting; early layers emphasize COMPARE/BRANCH, mid-layers ASSIGNADD, late layers taper to NOOP.
J.2 REGISTER PROBES AND ROLES
Linear probes. For each register NI,...,NNy} we learn a linear probe (logistic regression) mapping R E RRd to labels reflecting hypothesized roles (e.g.. "rumning min", "pivot", "frontier membership"), training on train splits and evaluating on OOD lengths; we report mean balanced
accuracy (BA) and mutual information MT( (") Z) with discrete role Z. To estimate MI, we dis- cretize probe logits via B-20 equiprobable bins to obtain p(= bin) and then compute
Σ p(b, 2) log p(b, 2)
Mi(R; Z)
(102)
p(b)p(z)
which we found more stable than kNN estimators small sample sizes.
Table 14: Register probes Sorting @ 320 and BFS@33 Balanced accuracy (BA %) MI (bits).
Task
Register
Role
Soft Modules
Cral-Transformer
Sorting
Rı
Running min
12.1/0.08
89.4 89.4/0.22
Sorting
R3
Accumulator
70.5/0.06
87 87.2 0..2
BFS
Ra
Frontier flag
68.3/0.07
84.6/0 0.2.
BFS
Parent pointer
65.9/0.05
81.1/0.18
Register specialization and disentanglement. Registers in the discrete controller exhibit consis- tent specialization across depths and OOD lengths (Table [4], whereas Soft Modules remain par- tially entangled. On Sorting @ 320, Rı (running min) improves from 72.1% 0.08 89.4% 0.29 (+17.3 BA points; +0.21 bits; 3.6x MI), and R3 (accumulator) from 70.5% 0.06 to
87.2 97. 0.24 (+16.7 BA; +0.18 bits; 4x MI). On BFS@320, Rg (frontier flag) rises from 68.3% / 0.07 to 84.6% 6% 0.22 (+16.3 BA; 0.15 bits; 3.1x MI), and R3 (parent pointer)
from 65.9% / 0.05 to 81.1% 0.18 (+15.2 BA; +0.13 bits; ~ 3.6x MI). The tandem gains in BA and MI (estimated via Eq. equation 102) indicate stronger,
disentangled role encoding rather than probe overfitting to superficial cues, supporting the claim that discrete selection sharpens register semantics and generalizes under length shift.
40

---

Published as a conference paper at ICLR 2026
J.3 KNOCKOUT AND RANDOMIZATION TESTS
Setup. We perform three interventions at inference: (K1) Controller knockoul (set y() = NOOP in a target layer range), (K2) Register shuffle (permute rows of R at each step), and (K3) Operation randomization (replace arg max with random op among the top-2); we report the resulting absolute accuracy drops A Acc at the longest lengths.
Table 15: Intervention effects at longest lengths (Sorting @ 320, Lum@320, BFS@ 320). Numbers are AAce in points (mean over 3 seeds).
Intervention
Sorting
Sum-of-List
BFS
(K1) Knockout (layers 7-12)
-46.8
-44.1
-27.9
(K2) Register shuffle
-12.5
-9.7
-11.2
(K3) Randomize op (top-2)
-18.9
-16.3
-14.6
Intervention analysis. As summarized in Table 15 (measured at the longest lengths: Sort- ing m20, Sum 3200, BFS @ 320), late-layer controller knockout (K1, layers 7-12) triggers the largest collapses, -66.8 pts on Sorting, -44.1 on Sum-of-List, and -27.9 on BFS, showing that correct execution of the learned plan in deeper layers is critical, not merely its early encod- ing. The smaller (but still substantial) BFS drop suggests partial redundancy from graph struc- ture (e.g., frontier dynamics) relative to strictly sequential algorithms. Register shuffle (K2) yields moderate yet consistent losses (-12.5/ 9.7/ 11.2 for Sorting/Sum/BFS), corroborating the register-role specialization observed in ş.22 disrupting stable assignments degrades performance even without changing token budgets. Top-2 randomization (K3) produces intermediate damage (-18.9/ 16.3/ 14.6), indicating that near-tic controller logits still encode semantically distinct choices, i.c., the argmax is not arbitrary. Overall, the pattern across KI-K3 supports the view that discrete control and persistent register alignment are both necessary for reliable long-length gener- alization.
J.4 STATE-SPACE VISCALIZATIONS
We collect register vectors (1. across layers time steps t, and registers i, then embed them via t-SNE (perplexity 30, 1,000 iterations) UMAP (15 neighbors, min dist = 0.1), coloring each point by the selected operation at that step so that clusters reflect functional states (e.g.. COMPARE, ADD regions). In practice, discrete controllers yield well-separated clusters aligned with operations, whereas Soft Modules exhibit overlapping manifolds; within ADD, trajectories progress smoothly as accumulators update, while BRANCH splits produce distinct lobes, consistent with explicit control flow.
Protocol.
Caveats and summary. High alignment/MI indicates predictive structure, not necessarily causal- ity: our knockout tests strengthen the case but remain within-distne within-distribution, and truly out-of-support manipulations (e.g., adversarial controller perturbations) are deferred to future work. Consistent with these caveats, the visualization in Fig. 8 shows operation-coherent clusters of register states, COMPARE and ADD occupy distinct regions, BRANCH forms split lobes (reflecting alternative control paths), and NOOP appears separate, low-variance island (white-filled entropy, semantically stable regimes rather than diffuse manifolds. Taken together, the evidence indicates that controllers learn interpretable plans: operation traces align with algorithm steps; reg- isters specialize and are linearly decodable; performance degrades under targeted interventions; and state embeddings form operation-coherent clusters, supporting the discrete-semantie analysis and the empirical gains reported in the main paper.
low-
K STABILITY AND FAILURE CASE STUDIES
We analyze where training is brittle and how to mitigate it. We focus on temperature/entropy sched- ules, ofe- and blended-brande errors at long lengths. and diagnostics that anticipate instabili- ties. Unless stated otherwise, results aggregate three seeds and use the same data/compute settings
41

---

Published as a conference paper at ICLR 2026
O ASSIGN O ADD O COMPARE O BRANCH O NOOP
t-SNE dim 2
-1
-2
-3
-4
-5
-6
-5
-4
33
-2
-1
0
2
3
4
5
6
t-SNE dim
Figure t-SNE-style seatter register states with 10 points per operation (n=50 total).
Case studies and diagnostics. With SLOW (ro=1.5, Tmin 0.2 5), late discretization leaves branches partially mixed well into mid mid-training, OOD accuracy rises steadily but underfits by 2-4 pts at convergence. With FAST (0.8, 0.05, 10 *). entropy collapses early but exhibits occa- sional loss spikes and settles at slightly higher final entropy than MEDIUM (cf. Fig. A3), consis- tent with premature commitment to suboptimal routines. The MEDIUM schedule (1.0, 0.1, 5x.5xx10 10-5) balances both, yielding the best AUL (Fig. 9. On Sorting/Sum at n=320, failure traces commonly show COMPARE-*ADD sequences where the write address advances before the comparison flag hard- ens. This yields off- aff-by-one placement accumulator updates, especially when the register write strength is large while x (comparison sharpness) lags. Increasing ramp speed and adding a small semantic monotonicity penalty reduces incidence. When entropy decays too slowly, BRANCH updates average then/else states, creating states that never appear the discrete program. Symptoms in- clude inconsistent frontier sets in BFS and small yet compounding numeric drift in Sum-of-Listt. Entropy penalties with mid-course bump (at 30-50% training) and straight-through (ST) harden- ing the last 20-40% of training substantially reduce blending. We log: (i) op-selection entropy Hop 1, E, (O); (ii) gradient norms gr ||VoLll2 with EMA; (iii) loss spikes (5% up- per tail of AL;); (iv) temperature traces (l) and sharpness s(l). Stable runs display monotone entropy decay 20.1550.20. bounded gt. and loss spikes confined early training.
K.1 MITIGATIONS
Entropy penalty schedules. A two-phase Aum (ramp-up 5x by 20k steps: brief mid-training bump) lowers H.p without premature hardening. Semantic consistency (Aum)). Encourages in- variants (e.g., accumulator monotonicity). reducing blended-branch states and cutting off aff-by-one
42

---

Published as a conference paper at ICLR 2026
Loss
Entropy
Accuracy (OOD)
80
Val loss
1.5
0.5
Ace
60
0.5
40
20
40
0
20
40
20
40
Epoch
Epoch
Epoch
MEDIUM FAST
Figure 9: Training curves for MEDIUM (blue, dashed) vs. FAST (orange, dotted) annealing with shared legend. FAST discretizes early (lower entropy sooner) but converges to slightly worse OOD accuracy; MEDIUM reaches lower final entropy and better OOD.
frequency by ~30%. Early curriculum. Start from shorter lengths (n (10,20), for 10-20% of training, then mix in full lengths; stabilizes comparisons before long-range writes. Late ST harden- arg max(y) in the last 20-40% of training; improves decisiveness
ing, Switch trraight-througe and reduces branch blending. Address damping. Decrease write strength À when address entropy high; prevents noisy multi-register writes.
K.2 FAILURE TAXONOMY AND FREQUENCIES
Table 16 summarizes observed failures at the longest lengths; frequencies are measured as the per- centage of test examples displaying the symptom under the stated schedule (means over 3 seeds).
Table 16: Failure taxonomy longest lengths (Sorting @ 320, Sum @300, BFS (90).
Failure
Symptom Trigger
Sorting
Sum
BFS
Slow anneal (blended branches)
Mixed then/else states; lingering high Hap under-discretized BRANCH. Early commitment to suboptimal routine: small spikes; higher final H.p than Medium.
3.8%
3.2%
2.5%
Fast anneal (early hard- ening)
2.9%
2.6%
2.1%
Off-by-one writes
Write index advances before flag hardens:
4.6%
4.1%
1.2%
accumulator/inex shift by 1.
Address diffusion
Multi-register writes when centrated; seattered updates.
con-
1.8%
1.5%
1.7%
Controller uscillation
Entrupy uscillations with periodic loss spikes; usually resolves after mid- training.
0.9%
0.8%
1.0%
Mitigation most effective
Amt + ST
K1 + Asem ST curriculum
Failure summary and mitigations. As detailed in Table 16 the dominant errors at the longest lengths are off-by-one writes (Sorting 4.6%, Sum 4.1%, BFS 1.2%), followed by slow-anneal blended branches (3.8/3.2/2.55% and fast-anneal early hardening (2.9/2.6/2.1%) Address dif- fusion modest but nontrivial (1.8%.)/1..%%) while controller oscillation comparatively rare (x00.-1.0%) and often transient. These patterns highlight the interplay between selection temper- ature т, comparison sharpness K, and address concentration 8: late discretization leaves BRANCH mixed; premature hardening locks in suboptimal routines; lagging produces off-by-ones; and dif- fuse 3 scatters writes across registers. The task-specifi mitigations the table align with these mechanisms: raising Aen and applying late ttraught- (ST) hardening most effectively reduce
43

---

Published as a conference paper at ICLR 2026
blended branches in Sorting; increasing and adding semantic penalty Asmm curtail off- onf-by-one updates in Sum; and ST combined with a mild curriculum stabilizes frontier dynamics in BFS. In practice, monitoring Hap E.A.-smoothe' gradient norms, and temperature/Sharpnss traces reliably flags emerging issues; applying mid-course entropy shaping and the listed remedies lowers failure rates and improves OOD) accuracy the longest lengths.
44


Reviews:
[{"review_id": "RCDz3tyIFX", "reviewer": "Reviewer_cvky", "summary": "This paper presents an augmented version of a Transformer that includes controller blocks. These blocks select and execute symbolic primitives. The program state and Transformer activations influence each other. The controller module is made differentiable via the Gumbel-Softmax trick. This Transformer variant outperforms other Transformer variant on length generalization over several synthetic tasks. The Transformer also outperforms other variants on DROP and some program synthetic tasks.", "questions": "See weaknesses above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* The paper appears to be well executed. The authors should be applauded, as it seems like a considerable effort to construct such a model and get it to converge to reasonable solutions across the tasks studied.\n* The empirical results demonstrate strong length generalization on synthetic tasks, as well as strong performance on more real-world tasks such as DROP.\n* The paper provides some formal results related to the expressivity of the proposed controller module.", "weaknesses": "* The authors differentiate their approach from prior work on integrating programming primitives with neural models (e.g. neural GPUs, neural turing machines) by claiming that prior work has incurred \"significant architectural complexity\" or add \"training brittleness\". However, it's not clear that the proposed approach significantly mitigates these limitations.\n* My understanding of prior work was that a main challenge of such methods is that it can remain difficult to learn complex programs from weak supervision. While discrete operations can technically be made differentiable using continuous relaxations (e.g. Gumbel-Softmax) this doesn't necessarily make such models trainable in practice. Indeed, the authors design various additional objectives and annealing schedules. Additionally, the synthetic tasks seem to rely on trace supervision to learn complex compositional programs. I understand the non-synthetic tasks such as DROP don't use trace supervision, but also don't seem to require learning very compositional programs, and it wasn't clear what is actually being learned by the neural controller for DROP despite the claims related to interpretability. I don't want to propose running more experiments, but curious to hear the authors response to this. Can this approach scale to learning complex programs from weak supervision? Or is the approach useful even without this ability?\n\nNits:\n\n* \"Emergent\" seems like a strong claim in the title, given the strong architectural bias.\n* Some of the exposition and notation could potentially be improved, e.g. `d_{ff}` is presumably FFN hidden dimension but not explicitly defined (line ~88) and the relation between `\\mathcal{S}^{(l)}` and `s^{(l)}` was not immediately clear (line ~155) despite the potentially confusingly similar notation.\n* It might be interesting to discuss recent work on RASP, Tracr, and ALTA that demonstrate that program primitives can be directly compiled to and represented by standard Transformers.", "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": "4GRQ0gTanV", "reviewer": "Reviewer_uGAD", "summary": "This paper proposes adding some programming primitives after every 2 \ntransformer blocks to give the neural architecture an ability to learn\nfixed length programs.\n\nThe paper describes this extension in detail. Specifically, the \"controller\"\nblock that goes after a transformer block has 8 registers (each storing a 64-dim vector), \nmemory (64x64) and some binary flags -- these define the state of the controller. \nThe operation set consists of Assign, Add, Compare, Branch and Noop operators. \n\nIf H_l = token representations entering controller block l\n1. controller first reads a summary s_l = \\phi( H_l ) \n2. s_l is used to draw an operator using Gumbel softmax using logits W_o s_l + b_o \nwhere o is one of the 5 operators.\n3. each operator has a semantics that is used to update the registers, memory and\nflags.  Assign and Add update the registers, compare updates the flags, branch \nupdates the full state. The final updated state is a linear combination of the\nupdates from the different operators - weighted by the values from (2).\n4. the output H_{l+1} of the controller block is H_l + \\rho( change in state ) W_s\n\nWhen this feeds into the next transformer block, we map the state-to-token \nand add to H_{l+1}, apply the transformer block on it, and feed that to the next\ncontroller block or transformer block.\n\nThe evaluation compares the new Ctrl-Transformer with pure transformers and their\nvariants. The benchmarks include sorting, sum-of-list, graph traversal, along\nwith some program synthesis, math reasoning and symbolic question answering. \nSome of the algorithmic problems are parameterized -- so one would train on some\nvalues of the parameters, say n = 10, 20, 40, 80 and then test on same or higher\nvalues of the parameters. The main finding is that Ctrl-Transformer maintains\nits performance for higher parameters, whereas baselines degrade in performance.\nEven on other non-parametric tasks, there is a visible improvement.\nThe paper also has detailed ablations and error analysis to justify the choices.", "questions": "I do not have a very specific question. Based on the detailed comments in the Weaknesses section above, if there is anything that can help clarify those points, then it will help me better understand the details of the paper.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "Strengths:\n1. The Ctrl-Transformer architecture is an interesting extension to transformers\n2. The experiments are fairly conclusive in showing the value of Ctrl-transformers\n3. There is extensive evaluation and discussion on how the various components are\ncontributing and how things are failing, showing that the investigation here is \nrigorous.", "weaknesses": "Weaknesses:\n1. My main concern is around the writing, but I believe it can be potentially fixed \nby a careful pass on the paper.  However, I could not figure out a lot of the details\neven after going back and forth between the different sections and the Appendix, and\nthat is the reason for my low confidence in the review. Detailed comments below.\n\nDetailed Comments:\nl87: Equation (3) defines FFN and Equation (4) uses SubLayer, which is undefined\nl173: The last part there where \\beta are used to combine u's is not clear. What \nare the u's ? Where is that R-tilde used?\nl177: In equation (9) the value read from the registers, r, is the second argument to\nSem function, but where is it used in Equations (10)-(14)?\nl191 mentions \"sharing parameters across layers\" -- how is reflected in the calculation\nof the total number of parameters that is present in the Appendix.\nl195: What is the function MHA( )? Where is it defined?\n - Is it true that the memory is only updated by branch? I don't see the memory getting\n  explicitly updated in the semantics of any of the operators.\n\nTheorem 1: The statement of the theorem lacks rigor - first the class of programs is only\ndescribed informally; moreover, the semantics of programs is unspecified. Looking ahead\nto Appendix D, there appear to be plenty of assumptions mentioned there, if they are indeed\nneeded for the proof, then the statement of the Theorem should also include those.\n\nl225-235 on \"Error under finite temperatures\" is also difficult to understand since it \nintroduces terms like minimum opcode-margin that are undefined.\n\nl307 says d_s = 128, but d_s on Line 1636 in the appendix is defined as N_r d_r + N_m d_m + N_f,\nwhich is clearly larger than 128 for the particular choices of the parameters.\n\nI found the parameter count calculation on Page 31 helpful. It would help further if you\ncould further clarify Params_ctrl calculation on Line 1638.", "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": "vi4WzSQgho", "reviewer": "Reviewer_HbvX", "summary": "The paper presents the idea of a discrete controller module. A discrete controller module adds to a transformer block the ability to learn 5 specific instructions, by using the Gumbel-softmax optimization trick to handle the discrete nature of the learning task.\n\nAn empirical evaluation is presented for 3 small benchmarks: a synthetic one, some string-to-program tasks inspired from RobustFill, and the DROP benchmark. In all of the presented results in Table 1, the method has a noticeable drop in accuracy as the input size n is scaled from 80 to 320. It is less pronounced than the baselines compared with, but it is a statistically significant drop. Table 2 and 3 reports that about upto 5% average gain in accuracy.\n\nThe added performance cost is about 5-7%.", "questions": "What is the scientific contribution that extends to program reasoning, and is applicable to tasks that need more than those 5 opcodes?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "- The proposed method and evaluation criteria is clearly described.\n\n- The performance cost is modest, 3 to 10% depending on often the controller is inserted\n\n- The traces are human interpretable, in the sense that one can inspect the opcodes printed out.", "weaknesses": "- The method is very specialized to the choice of 5 opcodes and the applications it can support. This severely limits the scientific value of the work.\n\n- No justification is given for how the benchmarks are chosen. They appear to be chosen to fit the 5 opcodes and are rather small.\n\n- We see a significant loss in accuracy, albeit lesser than other baselines, when n is scaled from 80 to 320. As such, the method doesn't scale well for the presented example applications, and will probably fair poorly as we increase n beyond what is reported. \n\n- It is difficult to interpret the alignment study results in Section 5.4. Taking the example of sorting, why are we checking trace alignment with a specific way of sorting, e.g., compare-and-swap? Doesn't the algorithm matter, eg. quicksort vs. bubble sort? \n\nOverall, one could perhaps see the scientific value better if the paper's method extends to a much larger set of opcodes. For example, the authors can consider taking the opcodes of a real processor (e.g. RISC-V, x86) and show if the method works generically for that. This would considerably broaden the applications for which the idea of discrete controllers can be used.", "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": "jWpUtC88xW", "reviewer": "Reviewer_H2RF", "summary": "The paper designs, controller-augmented transformer, using Gumbel-Softmax trick to integrate diccrete operations and differentiable learning. The approch in the paper is notable for its strong theoretical grounding, providing an expressivity guarantee for a class of bounded imperative programs. Empirical results confirm the effectiveness of controller-augmented transformers.", "questions": "I know there are many approaches that combine statisical selction and differential learning. Why do you choose  to combine Gumbel-Softmax and transformers? Is it because it is easy to make theortical proofs for Gumbel-Softmax?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The combination of Gumbel-Softmax and transformers leverages both advantages of symbolic learning and transformers\n2. The theortical proofs are elegant.\n3. Empirical results shows considerable improvements compared to previous approaches.", "weaknesses": "I am unsure whether the paper have conducted enough comparsions between other approaches. CLRS-Algorithmic Transformer, the newest method in 4.2 BASELINES section, is developed in 2022, which maybe out of date. It hard to say whether the paper has really improved the margin of the state", "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)