Compositional Generalization through Gradient Search in Nonparametric Latent Space
Authors: Haruki Shirakami
James Henderson
Venue: ICLR 2026 Poster (2026)
Keywords: compositional generalization
variational Bayesian methods
meta-learning
abstract reasoning
nonparametric representations
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We propose test-time gradient search over nonparametric variational Bayesian latent representations to achieve improved compositional generalization in abstract reasoning tasks.

Abstract: Many state-of-the-art methods in deep learning fail at systematic reasoning in settings which require compositional generalization. To address this, we propose a novel architecture which uses a nonparametric latent space, information-theoretic regularization of this space, and test-time gradient-based search to achieve strong performance on compositional meta-learning tasks such as program induction, Raven's progressive matrices, and linguistic systematicity tasks.  Our proposed architecture, Abduction Transformer, uses nonparametric mixture distributions to represent inferred hidden causes of few-shot meta-learning instances. These representations are refined at test-time via gradient descent to better account for the observed few-shot examples, a form of variational posterior inference which allows Abduction Transformer to solve meta-learning tasks that require novel recombinations of knowledge acquired during training. Our method outperforms standard transformer architectures and a contemporary test-time adaptive variational approach, indicating a promising new direction for neural networks capable of systematic generalization.

Published as a conference paper at ICLR 2026
COMPOSITIONAL GENERALIZATION THROUGH GRA- DIENT SEARCH IN NONPARAMETRIC LATENT SPACE
Haruki Shirakami EPFL, Idiap Research Institute ha varuki... shi rakami @epfl.c.
James Henderson
Idiap Research Institute
james. henderson dd aap
ABSTRACT
Many state-of-the-art methods in deep learning fail systematic reasoning in tings which require compositional generalization. To address this, we propose a novel architecture which uses a nonparametric latent space, information- theerrtti regularization of this space, and test-time gradient-base based search to achieve strong performance on compositional meta-learning tasks such as program induction, Raven's 's progressive matrices, and linguistic systematicity tasks. Our proposed architecture, Abduction Transformer, uses nonparametric mixture distributions to represent inferred hidden causes of few-shot meta-learning instances. These rep- resentations are refined at test-time via gradient descent to better account for the observed few-shot examples, form of variational posterior inference which al- lows Abduction Transformer to solve meta-learning tasks that require novel re- combinations of knowledge acquired during training. Our method outperforms standard transformer architectures and a contemporary test-time adaptive varia- tional approach, indicating a promising new direction for neural networks capable of systematic generalization
INTRODUCTION
The ability to solve novel tasks by recombining known primitives, often referred to as compositional generalization, is commonly understood as prerequisite for general intelligence (Fodor] 1975 Cat-
[ell 1963 Chomsky 1965 Gick & Holyoak 1980). In settings where exhaustive memorization of
the solutions to problem instances is impossible, the ability to learn and compose reusable concepts necessary for solving problems in general; any system which fails do will exhibit catastrophic failure on problems which are out of distribution (Chollet, 2019). Historically, it has been debated whether or not connectionist architectures such as neural networks possess the capacity to learn representations and circuits which allow for this form of systematic recombination of knowledge to infer new concepts (Rumelhart & McClelland, 1986] Smolensky 1990] Fodor & Pylyshyn| 1988). While some progress has been made to refute this claim (Lake & Baroni 2023). ), and despite the rise of LLMs which seemingly display some characteristics of general intel Inteliience and compositional reasoning (Bubeck et al. 2023 An al. 2023 Hosseini al. 2022). a plethora of negative results suggest that current architectural paradigms are inadequate and so further investigation is warranted
(Dziri et al. 2023 Mirzadeh al. 2025 Opedal al. 2025 Shojaee al. 2025).
In this work, we propose novel architecture with strong compositional generalization ability. We demonstrate the presence of these abilities in our proposed methods by considering meta-learning lasks which necessitate outoof-distribution (OOD) adaptation to novel test instances, requiring some form of knowledge recombination. In particular, we study program induction (Summers)) 1977 Biermann, 1978) and grammar induction (Lake & Baroni (023), abstract meta-learning tasks which standard transformer architectures perform poorly on due to their compositional nature. While some combinations of compositional generalization (Schug et al. 2025 Chen et al. (220)) test-time adap- tation (Hübotter et al.] [025 Dong et al. [025 Gladstone et al. 2025| Mathur et al.] (025), meta- learning (Vettoruzzo et al.) 2025 Yao et al.) (2022), and abstract reasoning (Wang et al. 2025 Li ct al.) 2024) have each been studied in isolation, we are, to our knowledge, the first to propose and
evaluate method covering all of these aspects.
"Code will he published httpss///tttti..ccm//o/////eetbt

---

Published as a conference paper at ICLR 2026
Problems of this type require some form of search over hypotheses inferred from observations and prior knowledge (Macfarlane & Bonnet 2025 Chollet et al., 2025). In our case, these hypotheses include possible combinations of concepts and abstractions learned during training. By representing these hypotheses in some latent space, the problem of knowledge recombination (and by extension compositional reasoning) becomes matter of discovering some latent representation of the novel input which allows it to be explained using the generative model acquired during training, giving us a form of hypothesis testing.
Inferring such a latent cause of the observed input is abductive inference, so we call our proposed ar- chitecture Abduction Transformer. Abduction Transformer is a deep variational Bayesian model, like Variational Autoencoders (Kingma et al., 2013) and Latent Program Networks (Macfarlam & Bonnet 2025). but unlike these previous models it takes full advantage of the power of the trans- former architecture's set-of-vector embeddings by encoding the hidden causes as nonparametric latent representations (Henderson Fehr 2023). Nonparametric representations have the advan- tage that they generalize well across situations of varying complexity, such generalizing from simple concepts learned during training to their more complex compositions. Crucially, our method includes not only amortized inference but also test-time gradient-basee search over its latent space to find the most plausible hypotheses which account for problem instances. We show that this rest-time search procedure, combined with our information-theonee regularization over our latent space, en- ables discovery of minimal latent representations of inputs and consequently allows our model to perform well on novel tasks.
Our main contributions are summarized below:
We find that set-of-vector representations with variable cardinality lead to better compo- sitional generalization. We represent inputs as nonparametric (variable-sized) discrete mix- ture distributions. as opposed parametric (fixed-sized) vectors commonly used in previous Bayesian methods.
ii) We show that test-time gradient search over latent representations leads to improved gen- cralization extreme OOD regimes. Test-time search enables our models to solve problems containing vast number of unseen concept combinations, which standard transformers with the same training data are unable to solve.
iii) We demonstrate that stochastic sampling of representations at training-time leads to searchable latent space. By encoding inputs into parameters of Dirichlet processes from which we sample discrete mixtures, our latent space benefits from cormattion-thoretti regu- larization, resulting in a smooth and searchable space.
iv) Our models significantly outperform standard transformer architectures and previous test-time adaptive methods on OOD abstract reasoning tasks. In addition, nominally outperform GPT-5 Thinking (OpenAI 2025) (w.o. fine-tuning) on Raven's 's progressive matri- ces (Raven 1962) and perform comparatively on 1-D ARC problems (both in OOD settings), while using only ~ 1.2M parameters.
Overall, our contributions make significant progress towards developing test-time adaptive neural networks which are capable of knowledge recombination in novel situations.
2 META-LEARNING AS INFERRING HIDDEN MAPPINGS
Problem definition. We consider few-shot meta-learning problems where the objective is to infer some hidden mapping which explains the few-shot examples. The few-shot examples are defined as a set of input/output pairs, which form a problem specification:
X
{21 31). (22, Uz), ..., (Zn, }n))
A mapping H* is said to solve the problem if H (xx) = yss Vi E [1, n). In practice, we test whether the correct prediction of the ground truth test output y" some test query Equery is made, namely that y H* (Fiqury))
Hypotheses and compositions. We refer to a candidate mapping H as a hypothesis coming from some hypothesis space H. A problem specification X is considered function compositional if its solution H can be expressed as a composition of two mappings H,, H2 e H. such that H* = H, H2- In addition, problem specification with test query s'ueryy and ground truth test output
2

---

Published as a conference paper at ICLR 2026
Drichtet Processs Postaria
Teer-Shet Examples
Projection
will X) DPV o").o) Maln, Mimn ay Σa
Oeservation Vectors
Sample Latent Mature from Dériniet Process
H
Dacader
Abcuction Transfarmrr Encodhe
Prediction
Ground Truth
Fans: Shot Ecampi
Test Query Output
H
11144411114411
Text CLPPY
Prei Tradicter
Decoder
Figure 1: Overview of Abduction Transformer. Left::
Few-shot example pairs are tokenized and encoded into set of vectors from which the parameters of posterior Dirichlet process are inferred. Right: A latent mixture distribution is sampled from the DP and used to decode the inputs of few-shot examples into predicted outputs, from which cross-entropy loss is computed against the true outputs. We update our latent representation H w.rt. this loss to refine our hypothesis to better account for the few-shot examples. After repeating this process fixed number of times, refined latent representation to decode the test query into the final prediction.
y" of the same form Eq. is considered production compositional if there exist mappings H1, H2 E H such that:
y/ (x:) Vo H2(2)) V" = H1 (H2((ee))) i.j E [1,n]
We refer to the individual mappings constituting the compositions constituent hypotheses. A tuple (X, Equery, y²) is said be function compositional relative to a set mappings T if contains its constituent hypotheses, and the same can be said for production compositional problem instances. system A which has been trained to recognize set of hypotheses T C H (H*) is able to solve a problem with solution H which is function compositional or production compositional relative to T, we say that system A exhibits compositional generalization.
3 ABDUCTION TRANSFORMER
Few-shot learning as variational inference. The objective of the few-shot meta-learning task presented in the previous section be viewed posterior inference of the hypothesis / which best accounts for the observations specified by X. In other words, the task amounts to inferring the posterior distribution p(H] X).
Fig. gives an overview of the Abduction Transformer architecture. Similarly to VAEs (Kingma al. 2013), we define parameterized neural networks which approximate the distributions qa ( X) (encoder) and X H) (decoder), and train these networks to minimize variational free energy, an upper bound on the KL-divergence between the true posterior p(H X) and (H X x
F(d.0) -L( X) (((()) - (HIp Po pa X H)]
nonparametric latent space. Our architecture uses transformer encoder infer qo q which has the crucial property that the number of vectors in the transformer's output grows propor- tionally to number of tokens in its input. This ability to model situations of variable complexity is analogous to the way Bayesian nonparametrics is able to model mixture distributions of variable complexity. We take advantage of this ability to model mappings H with variable complexity, so that models trained simple mappings generalize naturally to their more-complex compositions.
X),
To do so, we use an approach proposed by Henderson & Fehr 2023) for general purpose transformer VAEs. We project the transformer output into the parameters of nonparametric distribution, namely
In practice, this loss is implemented with modi modifications; the exact form of the training loss is given in Eq.
3

---

Published as a conference paper at ICLR 2026
Dirichlet Process (DP). Each vector in the set output by the transformer is projected to pseudo- observation for the DP with distribution (Vps of) and pseudo-count a; In particular, we define
(H | X) as a DP such that
n
go X) = DP(No(H, (2), au) Nolu, = Σ N(H au
ai
aj
(4)
where ja RRR o² transformer output V E Rexp.
Raxx, I2 parameters linearly projected from the d-parameterized
As result, our latent representations are samples from DPs, ie. mixture distributions whose ef- fective number of components are unbounded and sensitive to input observations By treating our latent space distribution over distributions, we hope to learn a space which is smooth and disen- tangled, taking inspiration from VAE variants (Burgess et al. 2018). This approach can be viewed as a natural generalization of VAEs from the vector space regime the set-of-vector space regime, which is appropriate for transformer architectures.
A Abduction Transformer encoder. Our encoder transformer is trained approximate 4o(H Te Vi). where (Ig, 14) is a pair contained in problem specification. To give our transformer en- coder this probabilistic interpretation, its output is stochastic. The encoder takes a pair from the problem specification X. tokenizes it into a set of input vectors, and encodes it into the set of pa- rameters of a DP: lt, a², and a. Then discrete mixture sampled from this inferred DP using a factorized sampling method from Henderson & Fehr (2023). The mean of these samples across pairs in is taken as the encoder's s output hypothesis: == >i= His H; ~ qo(H z4,91). See App. A for an overview of the operations involved in inferring DP parameters.
Adduction Transformer decoder. The decoder specifies the distribution ý ~ po (yt x4, H). is implemented as transformer decoder that autoregressively generates its prediction by cross- attending to the latent representation H. and self-attending to the context z₁. Since is represented as mixture distribution, cross-attention H involves the denoising-attention operation, which Henderson & Fehr 2023) show theoretically subsumes regular attention.
Thus, because the decoder generates its prediction yi conditioned 2; and H, it can be viewed as a mechanism which computes H (x:). In practice, the latent hypothesis H is often the inferred average hypothesis H over the pairs in X, and x; is the input query Tquery- Further details regarding cenoising-attention and its equívalence regular attention can be found App.
Gradient search over latent hypotheses. Similarly to Macfarlane & Bonnet (2025), our archi- tecture allows for refinement of latent representations produced by the encoder to further improve congruence with the few-shot examples given in the problem specification. Given some hypothesis H, we allow its gradient-based refinement by minimizing - log Poyp shot examples [2,, } with respect to H.
x,, H) across few
This process can be thought of as gradient search over latent space to find a hypothesis that better accounts for potentially novel observations, where the search is initialized by a forward pass of the encoder. By decoding each candidate hypothesis into predictions over y; at each step, our procedure can be viewed as form of iterative hypothesis testing and solution verification against few-shot examples. A refined hypothesis Ho is then used to decode some test query Tquery in order to predict the ground truth y² both of which are not seen during the gradient search process.
Training procedure. As mentioned $3 the training loss for these networks variational free energy, standard objective used by architectures such as VAEs. We assume dataset of meta- learning episodes, each containing a problem specification X, a test query Equery and ground truth test output y² For each episode, we compute the loss:
1 L( >XI. n
Σ KL(g Ii, vi) ((H) log.poy" (aquery H)
OOrr implementation incudes isotropic Gaussian prior component in the DP's base distribution, which omitted from the expression given here for brevity. "Samples from DPs are theoretically infinite. We truncate samples by only considering Ko + com- ponents, where is the number of input vectors with an added prior component.
4

---

Published as a conference paper at ICLR 2026
A. 1-D ARC
C. Raven's Progressive Matrices
TexeShot Example
Test Query
Traniing Sec
O&D
I
?
2
Selutton Rule Ser
B. Composition in ARC
Compestionsl Gerersiization
Constant
Tresing 5Mc
3
?
Progression
Trensiiss
TITSEEN CencHa
Schition Rule
Cumpoos
Progression
6
Adation
Infer unseen compoolions
Learn to infor printive transforrations
Learn lnineer oubset carnisationss
D. Linguistic Systematicity Task
000 Compostional Gereralzation
Few Shor Esaricie
Test Quary
Dax
Lug Fep Kiki Wif
Dax Blicket Lug Kiki Wif
3
Solution Ruie
Wif
Dax Kiki Lug
S B
Gonsiont
Lug
Outpur
Additoo
Lug Blicket Wif
Lug Fep
Figure 2: Illustration of compositional generalization tasks. A (1-D ARC) A set of 3 few-shot examples are given the form of input/output pairs, and the objective is to predict the correct output some test query. B (Composition in ARC) We train a subset of possible transformations. At test time, we evaluate on unseen transformations which compositions of transformations seen during training. C (Raven's Progressive Matrices) Our model is trained on problems containing some subset of possible combinations of rules for feature progressions A| test time, we evaluate on problems with unseen rule combinations. D (Linguistic Systematicity Task) Our model trains on meta-leaming episodes containing sentence/interpreaation pairs. At test-time, the model is evaluated on meta-learning episodes with unseen grammars govern- ing the sentence/interpretatioo pairs.
model problems containing
and update our parameters & and according to back-propagated gradients, where Axi. is a scalar hyperparameter.
The KL-divergence term acts as an ieformmtionn-thoreee regularizer on the posterior DP parameters, encouraging sparsity in the mixture weights given to mixture components (the of vectors accessi- ble by cross-attention), and encouraging a smoother latent-space due to noisy training-ting sampling of component vectors. The prior we use in the KL-divergence term is DP with an isotropic Gaus- sian with unit variance as its base distribution, and concentration parameter set to one.
also possible train Abduction Transformer while including intermediate gradient search ap- plied to the average hypothesis H before decoding. In this case, the second term in Eq. becomes - log (y quuery; H"), where H" is the refined hypothesis after gradient search to optimize the few-shot examples in X. Details and pseudo-code for the training procedure, as well as the particular KL-divergence loss we use for DPs, are given in App.
4 PROGRAM INDUCTION IN 1-D ARC
For our first meta-learning task, we consider 1-D abstract spatial reasoning problems inspired by the ARC-AGI benchmark (Chollet et al. (025). The task consists of problem specification X contain- ing input/ouput pairs of pixel sequences which are all governed by some common transformation H*, and the objective to predict the ground truth output pixel sequence y² given Tquery- The task thus contains perceptual component of discovering entities within sequence of pixel values, as well as inferring the transformations being applied to them, and is useful as a benchmark for abstract reasoning ability (Chollet) 2019).
Problems of this nature are difficult for neural networks to solve, including architectures incorporat- ing pre-trained L.MSs (Xu et al. 2024 Dimitriadis & Samothrakis 2025 Chollet et al. 2025). In
5

---

Published as a conference paper at ICLR 2026
this section, we seek to investigate the ability of Abduction Transformer to solve ARC-like prob- lems whose solutions are not seen during training. but consist of compositions of transformations seen during training. In other words, investigate whether Abduction Transformer exhibits com- positional generalization with respect to function composition according to the definition given in $2
1-D ARC dataset. For these experiments, we utilize an open-source dataset arc-like (neurallambda
2024 which allows procedurally generate 1-D ARC-like problems with composable combi- nator functions. We train Abduction Transformer and other baseline architectures on a training set of ARC-like problems with solutions H E with [T| 36. The set of training hypotheses cludes primitive transformations such as translation, color-shift, expansion, sorting, reflection, and denoising, as well as a subset of possible compositions of these primitives. Training examples are generated by sampling a random transformation H from T, generating a random input sequence conditioned on H, and then computing the output sequence by applying the predefined combinator functions associated with H.
Our training set can expressed as set tuples Dunin {(X, T Equery y )= with problem specifications taking the form X = [(x;, (=)) We define Fí, Vi, Equery, 1f as sequences of pixel values such that I (z:)) y4 and !1 (Iquory) = y".
Experimental setup. To test OOD compositional generalization, we evaluate our trained models on test set with solutions H E V with |V| 14 such that V only contains transformations which are compositions of those found in T. but not themselves found in T. For example, if t translat and denois se are found in T. then V may contain translate dencise. In other words, we design test set of tuples representing meta-learning episodes Dee = {(x, qquery y² y²), ز=ر such that for all (X, Equery y²) Drest (X, Equery, )") is function compositional relative to T. See Fig. 2 for a visual illustration of the training and test sets. The test set contains 2000 such meta-learning episodes which are each generated in the same way those in the training set. See App. D.1 for complete enumeration of training and test set transformations.
Table 1: Performance on 1-D ARC OOD Composition Task. Zero-shot test set performance for GPT-5 Thinking and GPT-4.1 (OpenAl, 2024) are provided to contextualize task difficulty (see Appendix. H for details). We report standard error over 3 seeds.
Model
Solve Rate (%)
Gradient Search Steps
Train
Eval
Abduction Transformer (Ours)
25.1 +
100
LPN (Macfarlane & Bonnet 2025
1.9 + 1.0
100
Encoder-decoder Transtormer Baseline
0.1 1.00
1 None
100
Decoder-only Transformer Baseline
5.2 + 1.3
None
Ablations
Abduction Transformer (No KL-regularization)
16.7
100
Abduction Transformer (No gradient search)
0.1
None
None
GPT-5 Thinking (w.o. fine-tuning)
29.0
N/A
N/A
GPT-4.1 (w.o. fine-tuning)
11.0
N/A
N/A
Results. Table compares perfect solve rates across various architectures. We see that Abduction Transformer performs significantly better than any of the baseline architectures, solving 25.1% % of problems perfectly. Our results show that gradient search is effective for solving 1D-ARC prob- lems requiring compositional generalization. In addition, representations outperforms previous test-time gradient search approaches which utilize single vec- tor latent representations, namely Latent Program Network (Macfarlane Bonnet, 2025). We also see that our information-theone regularization plays an important role, with our ablations on KL- regularization performing worse.
find that using nonparametric latent
Althhouhh not illustrated, in order to drastically increase the size of the hypothesis space the feature vector components are randomly permuted column- column-wise, meaning the row (1, (1, (1, (1,4) may be modified to (1, 2), (3, 1). (1, 4) (where feature vectors in the second column of other rows have their components permuted in the same way).

---

Published as a conference paper at ICLR 2026
The baseline models with no test-time gradient search both show significantly diminished perfor- mance, demonstrating that compositional generalization on ID-ARC tasks is difficult without test- time adaptation. The effect of scaling the number of gradient search steps is discussed in App. D.3 D.3 Finally, we highlight that the deterministic encoder-decoder transformer baseline solves only 0. % of problems, even when trained with gradient search, indicating that without a stochastic treatment of the encoder the training process does not lead to searchable latent space. Model and hyperpa- rameter details can be found in App. D.2 Further discussion on the geometry of our learned latent space can be found in App. D.4
Verifying non-compositional abilities. To clarify the role that our nonparametric latent space plays in the improved performance for compositional generalization, we further evaluate both Abduction Transformer and I.PN on problems without unseen compositions, i.e., novel instances of problems whose transformations are contained in the training data and thus do not require compositional gen- eralization. Results are presented in App. F We find that on these problems, both Abduction Trans- former and LPN exhibit almost identical performance. This indicates that, in our experiments, the improvements from using a nonparametric latent space are specific to compositional generalization.
5 RAVEN'S PROGRESSIVE MATRICES
Our second experiment considers a symbolic variation of Raven's Progressive Matrices, an abstract reasoning task commonly used as a measure of intelligence in humans (Raven 1962). In its original formulation, a problem instance takes the form of a matrix of panels whose contents evolve row- wise according to some common set of rules, and the objective is to predict the contents of query panel usually placed on the last column of the last row. The task requires the discovery of a set of rules which explain the progressions seen in the problem, and therefore requires search over possible hypotheses (Carpenter et al., 1990). As such, we frame Raven's 's Progressive Matrices as a meta-learning problem similar to previous tasks considered, where we take each row of the problem as few-shot meta- meta-leaming example.
SRAVEN. We utilize the open source SRAVEN dataset, which is a symbolic variation of Raven's Progressive Matrices specifically developed to test compositional generalization in neural network architectures (SchuuurrrSSSgggggg et al. 2025). Instead of the graphical format which Raven's 's Progressive Matri- ces classically presented in, SRAVEN encodes each panel of problem instance into feature vector with feature values from a fixed vocabulary. Unlike the previous ARC-like reasoning task which contains large perceptual component, SRAVEN is designed to test purely symbolic compo- sitional reasoning, giving us fine-grained control over the rule combinations (i.e. compositions) that evaluate.
Specifically, SRAVEN encodes each problem panel as a K-ddimensional vector, where each feature can take on one of F-many integer values. Each feature is governed by progression rule which takes as input two integers and outputs a single integer. Thus, the first two columns of each row in combination with the set of K rules governing each feature determine the right-most panel of each row. See App. E. for an overview of the list of progression rules contained in SRAVEN.
Problems are generated by first sampling a of K rules (out of N total), and producing each row by sampling random inputs to the rules. The final column of the final row is treated as the test query and is hidden to the model. This leads a 3 x matrix R of K --dimensional vectors, described as a problem specification X {2 y: 2 where x; = 1 R222) and yı R... In addition, Tquery defined as (Rs,1, R3.2) and y* as R3.3- To increase the difficulty of our problems, the components of feature vectors are randomly permuted in the same way for each column in the problem across few-shot examples, Identically to the previous experiments, training sets take the form a set of tuples D = {(X, qquayy 1² >ز}ر and the same format applies to our test sets.
Experimental setup. In order to evaluate OOD compositional generalization ability, we generate two datasets with N = 8, K- 4 and F = 8 such that the rule combinations seen in problems of each dataset are disjoint[] For instance, if the training set contains a problem instance with rule set {A, B), then the test does not contain any problem with that rule set. See Fig. 2 for an illustration of our dataset split. Out of all possible combinations of K rules, we randomly sample
This gives (R+K-) ")-())))) -KK!² - K) 188, 760 many possible SRAVEN tasks (possible hypotheses) in total, taking feature permutations into account.
7

---

Published as a conference paper at ICLR 2026
fraction of these for training set; for this experiment, train models 1% of possible rule combinations, as well as on 90% of possible rule combinations, and evaluate on problems containing the held out proportion of rule combinations. In our few-shot meta-learning framework, this amounts to performing evaluations on test sets whose problems are function compositional relative to the training set. Our test sets consist of 2000 generated OOD SRAVEN problems each.
Table SRAVEN OOD Composition Task: Performance when training on 1% and 90% of possible rule combinations For each training method, we evaluate on problems with rule combinations taken from the remaining unseen proportion of combinations. We additionally present zero-shot perfor- mance on the test set containing 99% of rule combinations for GPT-5 Thinking and GPT-4.1.
Model
Train 1 %
Train on 90%
Solve Rate (%)
Gradient Steps
Solve Rate (%)
Gradient Steps
Train
Eval
Train
Eval
Abduction Transformer (Ours)
46.1 4.2
100
96.4 0,4
10
LPN (Macfarlane Bonnet, 2025
37.1 2000
100
93.5 1.0
10
Encodcr decoder Transtormer Baseline
10.8 + 22
100
27.2= 3.0
10
Decoder-only Transformer Baseline
28.8 4: 1.3
None
None
95.3 4 1.1
None
None
Ablations
Abduction Transformer (No (e-rreularizzaioon
16.8
100
33.7
10
Abduction Transformer (No gradient search)
20.9
None
None
96.7
None
None
GPT-5 Thinking (w.o, fine-tuning)
41.0
N/A
N/A
N/A
N/A
N/A
GPT-4.1 (w.o. fine-tuning)
34.0
NEA
N/A
N/A
N/A
N/A
Results. Our results are summarized in Table [ We find that when trained on 90% possible
compositions, Abduction Transformer and the decoder-only transformer baseline perform compa- rably, showing close to perfect performance on the OOD test set. When trained on 1 1 of possible compositions, find that the Abduction Transformer significantly outperforms all other architec- tures we tested. We observe that the decoder-only baseline which performed well in the previous setting now degrades significantly in performance, indicating poor generalization in more extreme OOD regimes. The LPN similarly equipped with test time gradient search performs better than the decoder-only baseline but not as well as Abduction Transformer, indicating the advantage of using nonparametric latent representations.
Our ablations show that 1) regulaallaionti during training is necessary for effective latent space search, 2) gradient search is needed for more extreme OOD compositional generalization and 3) gradient search fails when applied to a deterministic encoder-decoder architecture, indicating that probabilistic treatment is necessary. Model and hyperparamter details can be found in App. E.2 We discuss the effect of scaling gradient search steps App. D.3
Verifying non-compositional abilities. Similarly to our previous experiments, perform addi- tional evaluations on non-compoiitional SRAVEN problems. We again find that Abduction Trans- former and LPN perform almost identically in this regime. Results are presented in App. F
LINGUISTIC SYSTEMATICITY
Finally, we consider grammar induction task proposed by Lake & Baroni (2023) to test linguis- tic systematicity and compositional generalization ability in neural networks. The task is meta- learning task where each meta-lenrning episode contains 14 few-shot examples of sentence/interpre- tations pairs, namely pairs consisting of sequence of lexical symbols and sequence of colored circles. The objective of the task is to infer the underlying interpretation grammar from these exam- ples and ascertain the correct interpretation of some query sentence. See Fig. 2 for an illustration of the problem setup.
Experimental setup. We use the open source training and test sets created by Lake & Baroni 2023) to train and evaluate our models. The training set is a set of meta-learning episodes Dtrain [X, Equery- y²) ) each generated from a randomly sampled interpretation grammar, with X {(x,, y4 11 consisting of pairs containing lexical symbol sequence x; and a color sequence yi- The pairs in X are generated such that given the ground truth interpretation grammar H" we have
Reported standard error is computed over 3 seeds.
8

---

Published as a conference paper at ICLR 2026
that (x+) y/- The test set is generated in the same way from interpretation grammars that guaranteed to be different from those used the training set. Thus, in this experiment we evaluate models' ability to perform production compositional generalization. See App. G for details on the sampling process for interpretation grammars.
We compare the performance of Abduction Transformer against the encoder-decoder transformer architecture baseline while de- creasing the number of few-shot examples in test prob-
1.
from Lake & Baroni] (2023)
Perfect Solve Rate
lems. With this setup, we aim evaluate OOD compo- sitional generalization ability as few-shot examples be- come scarce to the point where they may not be sufficient to imply a unique interpretation grammar. In our exper- iments, we vary the number of few-shot examples given at test-time in the range [1, 14] and evaluate how well our models generalize to these situations.
0.6
0.4
0.2
Abduction Transformer
Results. As shown in Fig. 3 we find that compared to the encoder-decoder baseline, Abduction Transformer exhibits perfect accuracy down to 10 few-shot ex- amples before degrading, maintaining 50% accuracy at 5 few-shot examples. This is in contrast to the baseline architecture which shows continuous degradation in per- formance as the number of few-shot examples decreases. This indicates Abduction Transformer's strong compositional generalization ability in the face of in- complete information, considering the fact that the ground truth interpretation grammars in the test set contain exactly 7 rewrite rules and require the observation of at least 7 examples to be uniquely identified. Our results highlight Abduction Transformer's robustness against OOD regimes, both in terms of the meta-learning task itself and the scarcity of few-shot examples.
Enroder decorerr Baseline
0.0
10 Num. Few-Shot Examples
Figure 3: Performance on linguistic sys- tematicity task when varying the num- ber of few-shot examples given in meta- learning episodes at test time.
7 RELATED WORK
Test-time adaptation methods. Performing gradient updates at test-time has become popular approach in order to solve difficult inference problems. common approach is Test-Time-Fing Tuning (TTFT), method for improving inference performance by fine-tuning models some curated dataset conditioned on the test-input (Hübotter et al., 2025; Krause et al., 2018 Hardt &
Sun 2024: Sun al. 2020). Similar approaches have been used in successful architectures for the
ARC-AGI benchmark (Chollet et al.) 2025 Akyürek et al. 2025 Franzen et al., (025). Approaches
using updates model activations enable test-time adaptive performance increases while avoiding
model parameter updates (Macfarlane & Bonnet 2025 [] et al. 2025).
Compositional generalization. The ability for neural networks solve 0OD compositional rea- soning problems has been well studied in the context of LLMs (Kudo et al. 2023 An et al. 2023 Liu al. 2023 Mészáros et al. 2024).. Hosseini et al. (2022); Furrer
(202 indicate the presence of compositional generalization ability in LLMs for parsing tasks, and Schug et al. (2025) investigate the mechanistic origins of compositional generalization
transformer architectures by
framing attention as a hypernetorrk. Despite this, many studies have shown that LLMs, including those trained specifically for reasoning, fail on OOD problems that require compositional recombi- nation of learned knowledge and subroutines (Dziri et al. 2023 Mirzadeh et al. 2025 Opedal et al.)
2025 Shojaee et al., (025). Thus, the problem of whether transformer architectures (and neural networks in general) can learn to compositionally generalize remains open.
8 CONCLUSION
We introduced Abduction Transformer as novel architecture capable of OOD compositional generalization on reasoning tasks, using nonparametric mixture distribution latent representations, cotmmation-thhorett regularization, and test-time gradient search. Our architecture learns to infer initial distributions over mixture distributions in a smooth space of hypotheses which then supports the iterative refinement of these hypotheses at test-time to better account for fho- meta-learning

---

Published as a conference paper at ICLR 2026
examples, thereby solving problems involving unseen compositions of knowledge obtained during training.
Our experiments show that Abduction Transformer is capable of solving problems involving such unseen compositions ARC-like reasoning tasks, symbolic Raven's Progressive Matrices, and grammar induction tasks, even in extreme OOD settings, where standard transformer architectures and previous test-time latent space search methods fail. We show that performance drops signif- icantly without our lest-time gradient search procedure, and that the search process becomes less effective without information-thcoreti regularization, nonparametric representations, or stochastic training. Our method serves as a new direction in designing neural network architectures that are capable of complex OOD generalization in reasoning domains.
ACKNOWLEDGMENTS
We thank Fabio Fehr, Andrei Coman, Philip Garner, and Itsaso Olasagasti for helpful discussion and feedback. This work was funded in part by the Swiss National Science Foundation under the NCCR grant Evolving Language, Swiss National Science Foundation Agreement #51NF40 246.
Ethics statement. This work is largely foundational, demonstrating the viability of new architec- tural paradigm. We foresee no harmful applications of our methods, nor potential for discrimination, bias or fairness concerns. No LLMs were used in the writing or proofreading of the main text. See App.[]fo details on LLM use.
Reproducibilityy statement. Code will be published at AbductionTranste and App.
We provide information to reproduce our experiments in App. D
REFERENCES
Ekin Akyürek, Mehul Damani, Adam Zweiger, Linlu Qiu, Han Guo, Jyothish Pari, Yoon Kim, and Jacob Andreas. The surprising effectiveness of test-time training for few-shot learning, 2025.
URL https: //atpps://rxiv..ao/ org. /btp///.////1122775
Shengnan An, Zeqi Lin, Qiang Fu, Bei Chen, Nanning Zheng, Jian-Guang Lou, and Dongmei Zhang. How do in-context examples affect compositional generalization? In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki (eds.), Proceedings of the 61 Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 11027-11052, Toronto, Canada, 2023. Association for Computational Linguistics. doi: 10. 1/20233/12
Alan W. Biermann. The inference of regular lisp programs from examples. IEEE Transactions on Systems, Man, and Cybernetics, 88):555-6000 1978. doi: (9109//TMC.198.44331011
Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, Harsha Nori, Hamid Palangi, Marco Tulio Ribeiro, and Yi Zhang. Sparks artificial general intelligence: Early experiments with gpt-4, 2023. DRLhttpss /ttttps///riivi org abs 2303. 12712
Christopher P Burgess, Irina Higgins, Arka Pal, Loie Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in 8-vae, arXiv e-prints, pp. arXiv-1804, 2018.
Patricia Carpenter, Marcel Just, and Peter Shell. What one intelligence test measures: A theoretical account of the processing in the raven progressive matrices test. Psychological review, 97:404-31, 07 1990. doi: 10.1037/0033. 2900003333999999333933 97..34044
Raymond Cattell. Theory of fluid and crystallized intelligence: A critical experiment. Journal of Educational Psychology, 54:1-22, 02 1963, doi: 10.1037/10046743.
Xinyun Chen, Chen Liang, Adams Wei Yu, Dawn Song, and Denny Zhou. Composi- tional generalization via neural-symbolic stack machines. In H. Larochelle, M. Ran- zato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information
10

---

Published as a conference paper at ICLR 2026
Processing Systems, volume 33, pp. 1690-1701. Curran Associates, Inc., 2020. URL https://ppooeeiunnsnnneisssss..s httpp//erraeaa//pp///////////// 2221e222227762222211522200813163
pdf
Francois Chollet, Mike Knoop, Gregory Kamradt, and Bryan Landers. Are prize 2024: Technical report, 2025. URLhttps:: httpp////aaaxxi.oo//abb/////////
04604
François Chollet. On the measure of intelligence, 2019. URL http::://arlvvv.rr//aatt.
1911. 01547
Noam Chomsky, Aspects of the Theory of Syntax. MIT Press, Cambridge, MA, USA, 1965.
George Dimitriadis and Spyridon Samothrakis. Out-of-distributiion generalisation gent hard: evidence from are-like tasks, 2025. URL /ORLLa/ppaa///xx..oo///////////
09716
Hao Dong, Eleni Chatzi, and Olga Fink. Towards robust multimodal open-set test-time adaptation entropy-aware entropp optimization. In The Thirteenth International Conference on Learning net
via adaptive Representations, 2025. http::tt/pee...eeeeeeeee http:/rrrrreee..ettee/oooommdaod /tt/://mmmmmnmn/maoooaooooo
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Sean Welleck, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D. Hwang, Soumya Sanyal, Xiang Ren, Allyson Ettinger, Zaid Harchaoui, and Yejin Choi. Faith and fate: Limits of trans- formers on compositionality. In Thirty-seventh Conference on Neural Information Processing kckkrr
Systems, 2023.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural networks, 107:3-11, 2018.
Jerry Fodor. The Language of Thought. Harvard University Press, 1975.
Jerry A. Fodor and Zenon W. Pylyshyn. Connectionism and cognitive architecture: A critical anal- ysis. Cognition, 28(1-2):3-71, 1988. doi: 10.00160011000077(8)0000--
Daniel Franzen, Jan Disselhoff, and David Hartmann. Product of experts with LLMs: Boosting performance on ARC is a matter of perspective. In Forty-second International Conference on Machine Learning. 2025. URL / hwtp:/wwwww...ee.e...eeeeeeeeeee
Daniel Furrer, Mare van Zee, Nathan Scales, and Nathanael Schärli. Compositional generalization in semantic parsing: Pre-training vs. specialized architectures, 2021. URL
org/abs/2007 08970
Mary L. Gick and Keith J. Holyoak. Analogical problem solving. Psycholoyy, 12(3):306- 355, 1980. ISSN 0010-0285, doi: http:///ccaaaaaaa.aaaa////////// // //wp:/ccience..eee
URL h
com/
Alexi Gladstone, Ganesh Nanduru, Md Mofijul Islam, Peixuan Han, Hyeonjeong Ha, Aman Chadha, Yilun Du, Heng Ji, Jundong Li, and Tariq Iqbal. Energy-based transformers scalable learners and thinkers. arXiv preprint trXiv:2507.02092 2025.
Moritz Hardt and Yu Sun. Test-time training on nearest neighbors for large language models. In The Twelfth International Conference Learning Representations, 2024. URL https:
7/ //ppenzeveew.
James Henderson and Fabio James Fehr. A VAE for transformers with nonparametric variational information bottleneck. In The Eleventh International Conference on Learning Representattions 2023. URL WRRLtttp::/// / wwtp:/eetttooooeoooooeeee ottp://www......oo.oooooooooaaaa cs03x
Arian Hosseini, Ankit Vani, Dzmitry Bahdanau, Alessandro Sordoni, and Aaron Courville. On the Compositional Generalization Gap of In-Context Learning. In Jasmijn Bastings, Yonatan Belinkov, Yanai Elazar, Dieuwke Hupkes, Naomi Saphra, and Sarah Wiegreffe (eds.), Pro- ceedings of the Fifth BlackboxNLP Workshop on Analyzing and Interpreting Neural Nenworks for NLP, pp. 272-280, Abu Dhabi, United Arab Emirates (Hybrid). December 2022. Associa- tion for Computational Linguistics. doi: 1/2.223w1122222aacccooooo 1.22. URL https:
7/acl ogy. hww.aattthoollooor////222/2222 1ttp://0000..bacccooooo111111111 ackboxnlp-1.22

---

Published as a conference paper at ICLR 2026
Jonas Hübotter, Sascha Bongni, Ido Hakimi, and Andreas Krause. Efficiently learning at test-time: Active fine-tuning of LLMs. In The Thirteenth International Conference on Learning Represen- tations, 2025.
Diederik P Kingma, Max Welling, et al. Auto-encoding variational bayes, 2013.
Ben Krause, Emmanuel Kahembwe, lain Murray, and Steve Renals. Dynamic evaluation of neural sequence models. In Jennifer Dy and Andreas Krause (eds.), Proceedings ofe 35th International Conference on Machine Learning. volume 80 of Proceedings of Machine Learning Research, pp- 2766-2775. PMLR, 10-15 Jul 2018. URL https: httthhtt:://pooceec. htt.://ngg.......eseeeeeeeeeee
krause 18a. html
Keito Kudo, Yoichi Aoki, Tatsuki Kuribayashi, Ana Brassard, Masashi Yoshikawa, Keisuke Sak- aguchi, and Kentaro Inui. Do deep neural networks capture compositionality in arithmetic rea- soning? In Andreas Vlachos and Isabelle Augenstein (eds.), Proceedings of the 17th Confer- ence of the European Chapter of the Association for Computational Linguistics, pp. 1351-1362, Dubrovnik, Croatia, May 2023. Association for Computational Linguistics. doi: 10. 18653vvlv ea22.aall-min.98 ott.tt/pp///eaaaaaoooooo.ooo./o.
98/
Brenden M. Lake and Marco Baroni. Human-like systematic generalization through a meta-learning neural network. Nature, 623:115- 121, 2023. URL
oxq/CozpusID: 264489248
Hengli Li, Chenxi Li, Tong Wu, Xuekai Zhu, Yuxuan Wang, Zhaoxin Yu, Eric Hanchen Jiang, Song-Chun Zhu, Zixia Jia, Ying Nian Wu, and Zilong Zheng. Seek in the dark: Reasoning via test-time instance-level policy gradient in latent space, 2025. URL htttss://aaxi..oor
Wen-Ding Li, Keya Hu, Carter Larsen, Yuqing Wu, Simon Alford, Caleb Woo, Spencer M Dunn, Hao Tang. Wei-Long Zheng, Yewen Pu, et Combining induction and transduction for abstract reasoning. In The Thirteenth International Conference Learning Representations, 2024.
Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang, Transformers learn shortcuts to automata. The Eleventh International Conference on Learning Representa- lions, 2023. URL
Matthew Macfarlane and Clément Bonnet. Searching latent program spaces, In ICML 2025 Workshop on Programmatic Representations for Agent Learning, 2025. URL https://
openreview www.:e///ffoo.o mett://orrumm.dvaa3333222
Mrinal Mathur, Barak A. Pearlmutter, and Sergey M. Plis. MIND over body: Adaptive thinking using dynamic computation. In The Thirteenth International Conference Learning Represen- tations, 2025. https:: /wwpnrevieeeeeeeeeeoeeeeee..eeee net. www.mmiiddda.... JJCD
Anna Mészáros, Szilvia Ujváry, Wieland Brendel, Patrik Reizinger, and Ferenc Huszár. Rule ex- trapolation in language models: A study of compositional generalization on OOD prompts. In NeurlPS 2024 Workshop on Compositional Learning: Perspectives. Methods, and Paths Forward, 2024. URL http::ttt////teette htt.://wwwweeeteeeeorrroeooeeeee
Seyed Iman Mirzadeh, Keivan Alizadeh, Hooman Shahrokhi, Oncel Tuzel, Samy Bengio, and Mehrdad Farajtabar. GSM-symbolic: Understanding the limitations of mathematical reasoning in large language models. The Thirteenth International Conference on Learning Representations, 2025.
Vinod Nair and Geoffrey E. Hinton. Rectified linear units improve restricted boltzmann ma- chines. In International Conference Machine Learning, 2010. URL hhtpps//app. semanti cernttiicchoaass owg. oesectooler.orrr//oooppst1111333
neurallambda. Like arc, but code to generate visual puzzles., 2024. URL httts:
com/neurel 1 1atisraallmbba//a/o cambbb/aarcc1aa Accessed: 2025-06-30.
12

---

Published as a conference paper at ICLR 2026
Andreas Opedal, Haruki Shirakami, Bernhard Schölkopf, Abulhair Saparov, and Mrinmaya Sachan. evaluation on problems with arbitrarily complex proofs. In The Thirteenth International Conference Learning Representations, 2025. URL https:
MathGAP:
//openreveww net./ d=5ck9PIrTTp
OpenAL Introducing gpt-4.1 in the api, 2024.
URL http:://ppenal.comm//1
gpt 4-1/ Accessed: 2025-08-11.
OpenAL. Introducing GPT-5, 2025.
URL
https://penni..omm
introducing-gpe 5/ Accessed: 2025-08-11.
John C. Raven. Advanced Progressive Matrices, Set II. H. K. Lewis, London, 1962.
D. E. Rumelhart and J. L. McClelland. On learning the past tenses of english verbs. In Parallel Dis- tributed Processing, Volume Explorations in the Microstructure of Cognition: Psychological and Biological Models. The MIT Press, 07 1986. ISBN 9780262291262, doi: 10.7551 5237.003 000.0000.008. URL httthtts:///oll http::://ddoooooooo/ org http://.1...511111 / 5237. 003. 0008
Simon Schug, Seijin Kobayashi, Yassir Akram, João Sacramento, and Razvan Pascanu. Atten- tion as hypernetwork. In ICLR, 2025. URL
V4K9hiqNxE
Parshin Shojace, Iman Mirzadeh, Keivan Alizadeh, Maxwell Horton, Samy Bengio, and Mehrdad Farajtabar. The illusion of thinking: Understanding the strengths and limitations reasoning models via the lens of problem complexity, 2025.
06941
Paul Smolensky. Connectionism, constituency and the language of thought. In Barry M. Loewer (ed.), Meaning in Mind: Fodor and His Critics. Blackwell, 1990.
Phillip D. Summers. A methodology for lisp program construction from examples. J. ACM, 24 (1):161-175, January 1977. ISSN 0004-5411.
doi: 10.1145/3211992.22200. URL https:
//doi. 11t...1111//......... 322002
Yu Sun, Xiaolong Wang. Zhuang Liu, John Miller, Alexei Efros, and Moritz Hardt. Test-time training with self-supervision for generalization under distribution shifts. In Hal Daumé III and
Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, vol- 119 of Proceedings of Machine Learning Research, pp. 2229-9248. PMLR, 13-18 Jul 2020.
URL https: httptt::/ppocccecinnntt http://ppcceeeei.e..ee mlr
wwt..mm.........eeeeemmmemm www.prrssvvv... sun2 html
Anna Vettoruzzo, Lorenzo Braccaioli, Joaquin Vanschoren, and Marlena Nowaczyk. Unsupervised meta-learning via in-context learning. In The Thirteenth International Conference on Learning Representations, 2025. URL
Guan Wang, Jin Li, Yuhao Sun, Xing Chen, Changling Liu, Yue Wu, Meng Lu, Sen Song, and Yasin Abbasi Yadkori. Hierarchical reasoning model. arXiv preprint 2.X.v:25066./7:44 2025.
Yudong Xu, Wenhao Li, Pashootan Vaezipoor, Scott Sanner, and Elias Boutros Khalil. LLMs and the abstraction and reasoning corpus: Successes, failures, and the importance of object-based representations. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. URL ptttps://openreviie
net/ Ed-
Huaxiu Yao, Linjun Zhang, and Chelsea Finn. Meta-leamiing with fewer tasks through task in- terpolation. In International Conference on Learning Representations, 2022. URL https:
peppenreview.net//e/e
d=a XXWF7
A ABDUCTION TRANSFORMER ENCODER
Given transformer outputs V E Rx the encoder linearly projects these representations to pre- activtion parameters μ E R=xP, log(o²) E R2',, a E R". Following Henderson & Fehr (2023), ) apply exponential activation function to our log- variancee parameters and ReLU (Nair & Hin- ton 2010) to our concentration parameters. This gives a posterior mixture distribution containing component (1i, oi, 0;) per transformer output, along with the prior component.
13

---

Published as a conference paper at ICLR 2026
B DENOISING ATTENTION
In standard attention, a query vector accesses set of vectors through a weighted sum parameterized by its attention vector. In denoising attention (Henderson & Fehr 2023), this interpretation is ex- tended by generalizing sets of vectors probability distributions over vectors, and treating attention as function of these probability distributions.
Scaled dot-product attention. We first consider standard cross-attention where some input query is mapped so single result vector. Here, the input w' E IR is projected to the query through w E Rxod We obtain keys and values by projecting the set of vectors Z E I through W WV E IRx respectively. We can regroup the standard expression of the attention function to operate in the space of Z, that is:
Attention Z: Z:W9, W wK) Attn(/"9(WK)) Z)WV = Attn(u, Z)WV
(6)
where u = ( u E 1xp, This new operation Attn(u, 2) can be characterized both as sum over vectors E Z, in terms of an integral over distribution with support at 2,. This equivalence restated below:
Attn(u, Z) = softmax (222ZT) Z = =DAttn(u; F2)
(7)
exp
2vd
||:/
Fz
6z:
(8)
i=1
1== exp
2 [|:113))
f(o) g(u;u, vdl) v do
DAttn(u;F)
(9)
Jo [()) (u; vdl) du
Here, 62, is impulse distribution (Dirac delta function) at zs. f(-) is the probability density func- tion for distribution F, and g(u; v, vdT) is Gaussian distribution with diagonal covariance.
Interpretation as denoising. The operation DAttn(u; Fz) which we call denoising attention can be thought of as the mean of the posterior distribution (over e) induced by making observation of some true vector o corrupted by noise, where e is drawn from a prior distribution F' specified generalization of regular attention, where instead of restricting sets
by Z. Denoising attention of vectors, it is defined for any distribution F over a vector space. When the input distribution F is discrete. it can be implemented naturally by including bias term in the cross-attention operation.
C TRAINING DETAILS: PSEUDO CODE AND KL-DIVERGENCE
Pseudo code. shows training procedure with gradient search enabled.
14

---

Published as a conference paper at ICLR 2026
Algorithm 1: Abduction Transformer training procedure with gradient search
Input: Encoder parameters Φ, decoder parameters 0
for / <- 1 to num traini ng steps do
Draw problem specification with input-output pairs (r4, 34). test query Equery and ground truth test output
x Sampl ling from DP
for i < 1 to do
L Sample He ~ ( 24.
/+ Gradient search
H --- H--Σ H4
n i=
for k t-1 to num www.grraieee..eerr..scctt
H 1- V log Ky H
==
// We ignore the second-order gradient w.r.t. 0
Ho -
(( 0) Ei= K F4, 9:) p(H)) log (query Ho) Update and via gradient descent on L(o, 0)
10
KL-divergence for Dirichlet processes. To compute the KL-divergence between DPs, we use an approximation due tofHenderson Fehr (023). Here, superscripts and p denote that the particular DP parameter belongs to the posterior and prior DP, respectively. Note that the number of input vectors (the additional term accounts for the prior component). and i refer to the gamma function and digamma function, respectively.
1, where is
DRI(g(H) X) |l p(H)) x T.D + LG
(10)
Ln log [(ag) - logT [(ag) + (ag ag) (-v(aa) + t
ai
Ku
+ kg [logT'
ax No
log
03 I
(11)
No
a+1
(m 12))
s0 )2
(0 )2
LG Kg
log
(12)
(o
(a
D 1-D ARC
D.I 1-D ARC TRANSPORMATIONS
We show the transformations contained in the training set and test set in Figures 4 and respectively.
Each pixel sequence is plotted as horizontal array. We plot 3 randomly generated input/output pairs for each transformation. The labels ('F.g.' the y-axis mark each input/output pair.
15

---

Published as a conference paper at ICLR 2026
1a1
1e:1
WXE Kagad tkds
XTEX Co 1ock
16-74
Colk
66.3
(6
coantin aqeeati
scara2 soaretl
16
ezerdi crtsorb
rett SIEELINE
NTDVE IngeE bodia andechr
66
Figure 4: 1-D ARC Training set transformations
coandtl
Co
Figure 5: 1-D ARC Test set transformations
D.2 ARC-LIKE TASK MODEL DETAILS
Model hyperparamters are shown in Table Training parameters used across all architectures are shown in Table
Transformer layers. All of the models we test utilize transformer layers that are structured in the usual way:
16

---

Published as a conference paper at ICLR 2026
Fatm MHA (Li)) + Fout FFN(LN (JFaan))))))))
Jatm
Our FFN is standard MLP with single hidden layer, using SiLU (Elfwing al. 2018 acti- vation. We use learned positional encodings and one-hot encode integer inputs corresponding to pixel/symbol values. Both positional and token embeddings accessed via dense embedding matrix.
Generation method. For all the models test, we generate predictions autoregressively. For Ab- duction Transformer, LPN and the encoder-decoder baseline, transformer decoder autoregressively generates its predictions by self-attending the test input and cross- atssednnnn to the encoded con- text. For the decoder-only baseline, the test input is given as prefix, and the prediction is generated by direct autoregression conditioned on the prefix.
Table 3: Details for ARC models.
Abduction Transformer
L.NN
Fncoder-decnder Baseline
Decoder-only Baseline
Num. Parameters
1,255,393
1,273,152
1,255,393
1,389,120
Encoder Layers
4
Decoder Layers
5
12
Emb. Dim.
96
96
96
96
MLP Dim.
384
384
384
384
No. Heads
6
Dirichlet KL Coef.
0.1
N/A
N/A
N/A
Gaussian KL Coef.
0.001
0.001
N/A
N/A
Gradient Search I.R
0.1
N/A
N/A
Table 4: Training hyperparameters. The decoder-only baseline differs only in the batch size used for training (shown in parentheses). All models converged in training loss and validation accuracy before reaching max. training steps.
Hyperparameter
Value
Training Steps
30,000
Batch Size
1024 (32)
Optimizer
Adam W
Gradient Clipping Norm
1.0
Learning Rate
0.001
D.3 SCALING GRADIENT SEARCH STEPS
Fig. shows the effect of scaling the number of gradient search steps take during test-time search. Both Abduction Transformer and LPN benefit from an increased number of search steps on both tasks.
17

---

Published as a conference paper at ICLR 2026
Abduction Transformer
Single Vector LPN
1-D ARC Task
SRAVEN (Train on 1%)
50%
40%
Accuracy (%)
30%
10%
0%
10
10
Gradient Steps
Gradient Steps
Figure 6: Accuracy against number of gradient search steps test-time. Left: Performance of Abduction Transformer and LPN on the 1-D ARC OOD composition task. Right: Performance of the same models on SRAVEN OOD composition task where we train on 1% of rule combinations.
D.4 ABDUCTION TRANSFORMER LATENT SPACE
Visualizing latent space. We hypothesize that latent space amenable to compositional general- ization requires the representations of various hypotheses be well organized in its geometry ac- cording to their semantics. To understand Abduction Transformer's latent space, we sample 15,360 randomly generated ARC problems from our training and test distributions and plot their repre- sentations Our mixture distribution representations their expectation, then projected into 2D using t-SNE; the results are presented Fig. 7 Our plots demonstrate that Abduction Transformer's latent space is remarkably well separated across different primitive transformations seen during training. Furthermore, the representations unseen compo- sitions map onto sensible locations near their constituent transformations.
first flattened into single vectors by taking
A
B
Figure 7: Visualization of latent space A (Primitive Transformations) t-SNE plot showing the structure of Abduction Transformer's latent space. Fach color corresponds to a particular primi- tive 1-D ARC transformation. Our mixture distribution representations are flattened by taking their expectation. B, C (Unseen Compositions) Latent representations unseen I-D ARC problem spec- ifications projected to the same space. B shows the representations for colorshi ft in blue, endpoints in orange and the unseen composition colorshife endpoints in magenta. C shows translate in black, denoise in purple, and the unseen composition translate denoise se in green.
'We encode problem specifications using an Abduction Transformer instance trained with step of gradient search on the training set. The reported representations are taken from the initial encoding before test-time gradient search is applied.
18

---

Published as a conference paper at ICLR 2026
E SRAVEN
E. SRAVEN PROGRESSION RULES
List of SRAVEN progression rules, taken from from/Schug et al. (0255. F refers to the size of the feature vocabulary.
1. Constant: Each row consists of random but fixed integer from [L....,FF:
2. Progression (+1); The first element of each row is sampled uniformly at random and incremented by modulo F for each successive column.
3. Progression (+2): The first element of each row is sampled uniformly at random and incremented by modulo F for each successive column.
4. Progression (-1); The first entry is sampled randomly, and each following entry is decre- mented by 1 modulo F.
5. Progression (-2): The first element of each row is sampled uniformly random and decremented by 2 modulo F for each successive column.
6. Addition: Two elements are sampled uniformly at random for each row and added modulo F to obtain the last column.
7. Subtraction: Two elements are sampled uniformly at random for each row and subtracted modulo F to obtain the last column.
8. Distribute three: independently sampled random permutations for each row.
Three elements are sampled uniformly at random and presented in three
E.2 SRAVEN MODEL DETAILS
Model hyperparamters shown in Table Training parameters identical those shown in Table Details regarding the transformer layers and generation method used in all models are identical those found in DD2
Table 5: Details for SRAVEN models.
Abduction Transformer
LPN
decoder-decoder Buseline
Decoder-only Baseline
Num. Parameters
1,093,44 441
1,148,640
1,093,441
1,131,840
Encoder Layers
4
Decoder Layers
Emb. Dim.
96
96
96
96
MLP Dim.
384
384
No. Heads
Dirichlet KL. Coef.
Gaussian Coef.
0.001
0.001
N/A
Grudient Search L.R
F VERIFYING NON-COMPOSITIONAL ABILITIES
Table 6: Performance non-compositional 1-D ARC problems.
Model
Solve Rate (%)
Gradient Search Steps
Train
Eval
Abduction Transformer (Ours)
98.39
100
LPN
97.90
100
19

---

Published as a conference paper at ICLR 2026
Table 7: Performance on non-compositional SRAVEN problems.
Model
Solve Rate
Gradient Search Steps
Train
Eval
Abduction Transformer (Ours)
99.95
100
LPN
99.90
100
G INTERPRETATION GRAMMAR AND META-LEARNING INSTANCE SAMPLING
An example interpretation grammar. taken from Lake & Baroni (2023) is shown in Fig. g Inter-
pretation grammars (which each correspond to individual meta-icarning instances) are randomly generated from simple meta-grammar.
Rewrite rules for primitives (the first rules in Fig. 8 are generated by randomly sampling input and output symbol pairs without replacement. Rewrite rules for functions generated by first sampling the LHS, followed by the RHS. The LHS is generated by sampling (without replacement) an input symbol random, then sampling whether the function is unary/binary, then sampling either primitve or non-primitive variables as its arguments (u and æ respectively in Fig. ( The RHS is generated by sampling random string of length S consisting of any of the function arguments defined in the LHS.
dax]
[wif]]
[lug]
[zup]
[u fepl
[n [u][uu]
[u blicket u2]
[u1] [2']
[u]]
[m] kiki x2]
[x2] [x1]
Figure 8: Example interpretation grammar. Double brackets [-] denote the interpretation func- tion. Variables x, apply to arbitrary non-empty strings, while u, apply only to dax, wif. lug, and zup.
H EXPERIMENTS ON GPT MODELS
We use prompts shown in Fig. 9 and Fig. 10 for our I-D ARC and SRAVEN experiments on GPT- 5 Thinking and GPT-4.1. We allow a maximum of 10, 000 reasoning tokens for GPT-5 Thinking (which corresponds to 'effort' set to 'low' for our tasks), and maximum of 10, 000 output tokens for G.T-4.1. We evaluate on 400 randomly selected problems from the test set for 1-D ARC, and 200 for SRAVEN.
20

---

Published as a conference paper at ICLR 2026
1-D ARC Prompt
You are sclving a 1D pattern recognition puzzle. Bach puzzle consists of sequences of pixel values from 0-9, where 0 represents empty/background.
You will be shown several onput-output examples that demonstrate a transformation rule. Your task is to identify the pattern and apply it to new test input.
EXAMPLES:
Example 1:
Input: <sequence>
Output: <sequence>
TEST:
Input: <sequence>
Output:
You may think through the problem in detail, but make sure to end your response with the final answer in this exact format:
FINAL ANSWER: lyour sequence herel
The sequence should be space- separatee integers only
Figure 9: Prompt used for 1-D ARC problems.
21

---

Published as a conference paper at ICLR 2026
SRAVEN Prompt
You are sclving a SRAVEN (Symbolic RAven' s) puzzic. This is a visual reasening task adapted to sequences.
Each puzzle consists grid of visual panels, where each panel is represented by a sequence of feature values (integers 0-7). Each
row in the grid follows consistent rule or pattern across its panels.
You will be shown several rows as examples, where each row contains: Input: The first few panels of the row (showing the pattern)
Output: The final panel that completes the pattern
Your task is identify the underlying rule and apply it predict the missing final panel the test row.
EXAMPLES:
EXAMPLES:
Example 1:
Input: <sequence>
Output: <sequence>
TEST:
Input: <sequence>
Output:
Analyze the pattern across the example rows and apply the same rule to complete the test row.
You may think through the problem step by step, but make sure to end your response with the final answer in this exact format:
FINAI, ANSWER: [your panel here]
The panel should be a comma- separatted 1ist of integers, e.g., [1, 2, 3, 4]
Figure 10: Prompt used for SRAVEN problems.
USE OF LARGE LANGUAGE MODELS
No LLMs were used in the writing or proofreading of the main text. We use LLM assisted tools for writing code used in models and experiments. We use LLM assistance for formatting LaTex figures.
22


Reviews:
[{"review_id": "IdgSTrW6Vn", "reviewer": "Reviewer_3Yws", "summary": "The paper introduces the Abduction Transformer (AT), a variational, test‑time‑adaptive architecture for systematic compositional generalization. AT encodes each few‑shot episode into a nonparametric latent space—a Dirichlet‑process (DP) mixture over a set of vectors—regularized by an information‑theoretic KL to a DP prior. A transformer decoder then cross‑attends to this latent mixture via denoising attention to generate outputs. Crucially, at test time the model refines the latent hypothesis $H$ by gradient descent on the few‑shot examples before answering the query, turning inference into iterative hypothesis testing. Across three tasks including 1‑D ARC‑like program induction, symbolic Raven’s (SRAVEN), and a linguistic systematicity benchmark, AT surpasses standard transformer baselines and a prior test‑time latent‑search model, and is competitive with zero‑shot GPT models under the authors’ protocol.", "questions": "* In table 2, why does the gradient-based refinement of $H$ harms the solve rate when trained on 90% of the possible compositions?\n* In table 1 and 2, how does the performance of GPT scale with the reasoning token budget?\n* How is the computational overhead compared to baseline methods such as Single Vector LPN?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "* This paper proposes a clear mechanism for compositionality through the procedure of (1) encoding few‑shot pairs, (2) inferring a DP posterior over latent hypotheses, (3) sampling a mixture $H$, (4) decoding, (5) refining $H$ with gradients on the few‑shot loss, and (6) answering the query. This operationalizes “search over hypotheses” rather than memorization.\n* t‑SNE visualizations in fig. 7 show clear clustering of primitives and sensible placement of unseen compositions near their components, consistent with the intended semantics of the latent hypothesis space.\n* Instead of a single latent vector, AT projects encoder outputs into DP parameters $(\\mu,\\sigma^2,\\alpha)$ and samples a mixture with a variable effective number of components, which is an upgrade from fixed-size latents and aligns with the transformer’s token‑proportional outputs.\n* Strong OOD results on diverse tasks are achieved with small models. AT achieves 25.1% perfect solve on strictly OOD 1‑D ARC compositions and 46.1% on SRAVEN with only 1% of rule combinations seen during training. The results match or surpass that from GPT-5 (Thinking) with only a 1.2M parameter budget.", "weaknesses": "* It is not very clear why denoising-attention is used in the AT decoder, and how much of the performance gain should be attributed to this design.\n* The comparison to GPT-5 can be a bit unfair, as it uses a \"low effort\" setting, while a large reasoning token budget is quite important for getting good performance in such reasoning tasks. Also, being able to update the latent hypothesis at test time could be a huge advantage.", "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": "iFTrrSpFCD", "reviewer": "Reviewer_xg9b", "summary": "The authors propose the Abduction Transformer to achieve better compositional generalisation on out-of-distribution (OOD) meta-learning tasks. The key idea is to represent each few-shot task with a nonparametric latent hypothesis, modelled as a Dirichlet-process (DP) mixture over latent vectors produced by the encoder. This allows the model to adaptively allocate a variable number of latent components to capture the compositional structure of the task. The authors evaluate their approach on three tasks: 1-D ARC-like program induction, symbolic Raven's Progressive Matrices (SRAVEN), and linguistic systematicity and show competitive performance with fewer parameters (1.2M), outperforming standard transformer baselines and comparable to GPT-5 Thinking on certain tasks.", "questions": "- Under what assumptions on the decoder does the test-time objective exhibit a landscape such that gradient search converges to semantically correct hypotheses rather than spurious modes?\n- It's interesting to see the model maintain high accuracy even on less than 7 samples from the linguistics task. How can you disentangle prior-based guesses from genuine compositional recovery?\n- Maybe I am missing something but can you explain why the number of gradient steps while training is 1 but is 100 for evals?\n- Why is it appropriate to average samples from different distributions at the encoder's output? Why is averaging done here rather than using a product of experts or other combination method?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper has a clear probabilistic framing, casting few-shot meta-reasoning as posterior inference over a task-specific hypothesis, trained with a variational free-energy bound and refined by test-time gradient search.\n- The DP prior over the latent space acts as a strong architectural prior that naturally aligns with the compositional learning premise.\n- The authors show the results of extensive ablation studies on modeling choices such as KL-regularisation over the DP posterior and test-time gradient search, indicating both are crucial.\n- The latent-space visualization suggests unseen compositions embed near their constituent primitives, which is qualitatively consistent with the paper’s goal of compositional reuse.", "weaknesses": "- From the paper, it seems that there are three types of composition: in data (function composition) as in 1-D ARC, in hypothesis space (production) as with the linguistic systematicity, and in features as independent submodules (parallel composition) as in SRAVEN. The issue is that the paper conflates these fundamentally different kinds of composition under a single formalism. Although all are expressed with the same symbolic notation $H_1(H_2(x))$, they correspond to distinct cognitive and computational challenges. Without explicitly distinguishing these regimes, it becomes unclear whether the reported improvements arise from genuine compositional reasoning or simply from learning parallel or modular mappings.\n- The paper relies on approximate KL and truncating DP samples ($\\kappa_0 = n + 1$), but does not quantify the approximation or truncation error in training or test-time search. Without rates of approximation or stability results, the “information-theoretic regularisation” claim remains empirical.\n- Why does gradient descent on latent representations lead to better hypotheses? It lacks convergence analysis and could benefit from more insights into the geometry and identifiability of the latent optimization: when and why gradient updates align with true compositional structure rather than overfitting to spurious patterns.\n- The method could benefit from even a small-scale analysis on a real compositional regime such as code editing, or some visual entities from ARC-AGI. Would it be possible to show such results?\n- The baseline comparisons for methods that perform test-time adaptation are limited to just one: single vector LPN and there are no comparisons to other compositional generalisation methods (for some examples which might not be up to date, neural module networks, program synthesis approaches, neural interpreters etc). Even if some methods need adaptation or don't directly apply, discussing why would strengthen the paper. \n- The notation is confusing since the paper uses $H$ both as a _function_ (the hypothesis mapping in compositional generalisation) and as a _latent variable_ (a mixture distribution in the nonparametric encoder). This can be made clearer by using, for eg, $h$ as the latent variable and $H_h(\\cdot)$ as the functional mapping. \n- The writing of the central claim can be made a bit clearer in my opinion. The 1-D ARC and SRAVEN setups are synthetic, controlled, and compelling for diagnosis, but the core gains are still modest, and comparisons to “GPT-5 Thinking” are zero-shot and prompt-based versus a trained model with test-time search making not like-for-like in compute or adaptation. \n- Appendix B claims denoising attention \"subsumes\" regular attention for _discrete distributions_ (categorical variables with additive noise), not for the _continuous Gaussian_ case used in the model. Without specifying assumptions like zero noise variance or small-variance limits, it is not clear the conditions under which this holds. \n\n\nI am happy to update the score of the paper if my questions and concerns can be addressed.", "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": "HirMIQba2B", "reviewer": "Reviewer_1Rbo", "summary": "This work addresses program synthesis tasks, with a focus on out-of-distribution (OOD) and compositional generalization. They introduce a new architecture named the Abduction Transformer, whose goal is to infer a latent cause (hence, abduction) of the observed specification for a given task, and use it to predict the solution of the considered task. This two-stage process implicitly defines a probabilistic graphical model corresponding to a deep variational Bayesian model, of the kind used in variational autoencoders (VAEs).\n\nThe main contributions of the paper include:\n- a VAE formulation for neural program synthesis\n- the representation of the latent space as a nonparametric Dirichlet Process (DP)\n- gradient-based search in the latent space to maximize the likelihood of the specification\n- training with 1 gradient step, using the ELBO loss\n\nThe authors experiment with the Adbuction Transformer on three different benchmarks: a 1D ARC dataset, Raven's Progressive Matrices, and a linguistic systematicity task. They highlight performance improvements due to latent search and some OOD generalization at test-time.", "questions": "1. What made you go for a DP to model the latent space? What justifies moving from the standard implementation of Kingma et al., 2013?\n\n2. What are the novel aspects introduced in this work that you would like to emphasize?", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 1, "strengths": "The paper is well written, the method is promising, and the experimental setups are somewhat well explained. I particularly enjoyed Figure 1, which details the method very well (although the DP equations could be skipped in this figure for higher clarity). In my opinion, the largest novelty in this work is the modeling of the latent space as a DP.\n\nI believe the paper tackles a crucial aspect of machine learning, i.e., OOD generalization in the context of program synthesis, and demonstrates the proposed architecture quite well with a set of 3 different benchmarks. The related work section successfully connects the Abduction Transformer to recent works such as Test-Time Fine-Tuning [Hübotter et al., 2025] or the Latent Program Network [Macfarlane & Bonnet, 2025]. Overall, the work tackles an important problem in machine learning, and the proposed architecture seems promising to overcome some of the limitations of current SOTA systems.", "weaknesses": "There are several weak points I would like to address in this review.\n\n## Novelty\nThe authors carefully relate this work to the framework of VAEs [Kingma et al., 2013] and the baseline Latent Program Network (LPN) [Macfarlane & Bonnet, 2025]. After carefully analyzing the [LPN paper](https://arxiv.org/abs/2411.08706) and comparing it to this work, it is not clear to me what novelty this work brings that is not covered in the LPN paper. This is the most important weakness to me; it seems that the LPN paper covers most of what is \"introduced\" in this work, in its current version.\n- VAE network for program synthesis: in LPN.\n- Training loss: identical to LPN, i.e., max likelihood and KL to prior.\n- Training dynamic: 1-step of gradient descent in latent space, no backprop through latent update.\n- Test-time search: gradient descent in latent space, maximizing likelihood of the specification.\n\nA notable difference from the LPN baseline is the DP process used to parameterize the latent space (the baseline is a standard normal Gaussian distribution). If this is an important novel contribution from this paper, I would like the authors to demonstrate it (see ablations below).\n\n## Experiments\nThe three experiment setups presented in this paper remain on the toy level. Since ARC-AGI is mentioned, it would improve the contributions of this paper to demonstrate the generalization abilities of the Abduction Transformer on the ARC-AGI v1 or v2 benchmark. Alternatively, the [ConceptARC benchmark](https://arxiv.org/abs/2305.07141) could be a more accessible version to try it on.\n\n## Baselines\nIn addition to the LPN baseline, it would improve the paper to include other program synthesis baselines, especially if ConceptARC or ARC-AGI is considered to be added to the experiments.\n\n## Ablations\nThe parameterization of the latent space as a DP seems interesting and would need to be ablated to be justified. The number of latent space gradient steps during training does not seem to be ablated or justified either.", "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": "OJp01TSLOR", "reviewer": "Reviewer_tPRa", "summary": "The authors build on recent work (LPN) that investigates learning compressed latent spaces that can be searched at test time to reduce the amortization gap from amortized inference. In this work, the paper proposes the Abduction Transformer, which uses multiple mixture distributions (via Dirichlet Processes), instead of a single Gaussian distribution, to learn the approximate posterior distribution. The method incorporates information-theoretic regularization and test-time gradient-based search in a nonparametric latent space. It is evaluated on OOD compositional meta-learning tasks such as ARC-like program induction, Raven’s progressive matrices, and linguistic systematicity tasks.", "questions": "1. Is it possible to control the number of distributions in the method? If so, how does performance of your method scale with number of distributions starting from 1 and increasing beyond what you investigated, to greater understand what is the key component generating compositionality. Any additional / alternative ablations that could give more insight into this would strengthen the paper also.\n\n2. How does performance compare on non-compositional datasets or datasets that do not explicitly target the requirement of compositionality?\n\n3. Can you give examples of problems solved and failed by LPN and the Abduction Transformer to understand more clearly the failure and success modes of each?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Targets a critical weakness in LPN: which does not necessarily handle compositionality in the latent space due to a single latent vector.\n- Approach to tackle compositionality is simple and scalable\n- Explores performance on a range of experiments, including ARC-like program induction, Raven’s progressive matrices, and linguistic systematicity tasks.\n- Uses appropriate baselines, like single-vector LPN, standard transformers, and GPT-5 Thinking.", "weaknesses": "- Lack of ablations on the method. Since the Abduction Transformer performance is supposedly derived from the larger number of distributions compared to LPN it is important to perform an ablation that scales this axis to see when compositionality emerges. Specifically, a graph with number of distributions on x-axis and test-time gradient search and non-gradient search performance plotted in the graph.\n\n- There is no evaluation on non-compositional data.", "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": "fD19KDgbW8", "reviewer": "Reviewer_VaZe", "summary": "This paper addresses the problem of compositional generalization and proposes the 'Abduction Transformer' as a solution. The model's core mechanism involves inferring Dirichlet Process Posterior parameters to effectively handle few-shot problems. The proposed algorithm consists of a 2-stage process: 1) An Encoder generates an initial hypothesis based on all few-shot examples, and 2) A separate Decoder architecture receives this hypothesis and performs gradient search (fine-tuning) to refine it into a hypothesis that best explains the examples. The authors validate their methodology on benchmarks such as ARC-LIKE REASONING and RAVEN’S PROGRESSIVE MATRICES.", "questions": "See weakness", "limitations": "", "rating": 2, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "This few-shot learning framework is generalizable, so can be applied to any benchmark.", "weaknesses": "This paper suffers from critical flaws in its experimental validation, primarily concerning benchmark and baseline selection.\n\n1. Omission of Core Baselines (MAML): The authors acknowledge that their methodology is a \"\"method of fine-tuning to minimize loss for a given few-shot input,\"\" and is fundamentally a meta-learning (learn-to-learn) approach. Therefore, the validity of this methodology must be directly compared against at least classic meta-learning algorithms like MAML. They should have verified whether their 'gradient-based hypothesis refinement' mechanism is superior to MAML's 'gradient-based model adaptation' on standard few-shot benchmarks (e.g., tieredImageNet). If their methodology indeed has such generalization capabilities, it must at least outperform classic meta-learning methods. Their method is, in fact, generally applicable. However, the paper intentionally avoids this crucial comparison, presenting only 'strawman' baselines like GPT and standard Transformers, which were not designed for meta-learning. This fails to prove any practical comparative advantage.\n\n2. Use of Arbitrary Benchmarks: The authors used a niche, self-created benchmark called \"\"1-D ARC-LIKE REASONING\"\", which they only state is \"\"inspired by\"\" the standard 2D ARC benchmark. A search for this benchmark reveals it is extremely niche. Furthermore, even the standard 2D ARC benchmark is not frequently used. On the contrary, on the 2D ARC benchmark—which should be more difficult—some existing LLMs have achieved nearly 70% accuracy. How does this discrepancy arise? Moreover, seeing that modern LLMs like GROK achieve nearly 0% on that benchmark suggests the benchmark has high variance between LLMs and is a task that LLMs are inherently bad at. Why does high performance here validate this model? For example, VLMs perform poorly at depth estimation. A very simple depth estimation model will easily outperform a VLM. What value does that simple model have? This merely proves the VLM's inability on that task; it does not guarantee the new model's performance. This appears to be an attempt to avoid direct comparison on recognized benchmarks, severely damaging the credibility and generalizability of the experimental results.\n\n3. Unrealistic Inference Cost (Test-Time Training): The most significant problem is that the method performs 100 steps of gradient search (fine-tuning) to solve a single test query. This is not fast few-shot inference. This is, in effect, Test-Time Training. This approach of training a model for 100 steps on a single problem renders any comparison against MAML (which takes only a few steps) or GPT (which requires no extra training) fundamentally meaningless. This is an unrealistic inference cost and not a fair comparison.\"", "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)