In-Context Compositional Learning vis Sparse Coding Transformer
Authors: Wei Chen
Jingxi Yu
Zichen Miao
Qiang Qiu
Venue: NeurIPS 2025 poster ()
Keywords: sparse coding
compositional learning
representational learning
Primary area: general_machine_learning

Abstract: Recent advances in AI, driven by Transformer architectures, have achieved remarkable success in language, vision, and multimodal reasoning, and there is growing demand for them to address in-context compositional learning tasks. In these tasks, models solve the target problems by inferring compositional rules from context examples, which are composed of basic components structured by underlying rules. However, some of these tasks remain challenging for Transformers, which are not inherently designed to handle compositional tasks and offer limited structural inductive bias.
Inspired by sparse coding, we propose a reformulation of the attention to enhance its capability for compositional tasks. In sparse coding,  data are represented as sparse combinations of basic elements, with the resulting coefficients capturing the underlying compositional structure of the input. 
Specifically, we reinterpret the standard attention block as projecting inputs into outputs through projections onto two sets of learned dictionary atoms: an *encoding dictionary* and a *decoding dictionary*. The encoding dictionary decomposes the input into a set of coefficients, which represent the compositional structure of the input. To enhance structured representations, we impose sparsity on these coefficients. The sparse coefficients are then used to linearly combine the decoding dictionary atoms to generate the output. Furthermore, to assist compositional generalization tasks, we propose estimating the coefficients of the target problem as a linear combination of the coefficients obtained from the context examples.
We demonstrate the effectiveness of our approach on the S-RAVEN and RAVEN datasets. For certain compositional generalization tasks, our method maintains performance even when standard Transformers fail, owing to its ability to learn and apply compositional rules.

In-Context Compositional Learning via Sparse Coding Transformer
Wei Chen, Jingxi Yu, Zichen Miao, Qiang Qiu
Purdue University, IN, USA {chen2732, yu667, miaoz, qqiu)@purdue, edu
Abstract
Transformer architectures have achieved remarkable success across language, vi- sion, and multimodal tasks, and there is growing demand for them address in-context compositional learning tasks. In these tasks, models solve the target problems by inferring compositional rules from context examples, which are com- posed of basic components structured by underlying rules. However, some of these tasks remain challenging for Transformers, which are not inherently designed handle compositional tasks and offer limited structural inductive bias. In this work, inspired by the principle of sparse coding, we propose a reformulation of the attention enhance its capability for compositional tasks. In sparse coding, data are represented as sparse combinations of dictionary atoms with coefficients that capture their compositional rules. Specifically, reinterpret the attention block as mapping of inputs into outputs through projections onto two sets of learned dictionary atoms: encoding dictionary and a decoding dictionary. The encoding dictionary decomposes the input into a set of coefficients, which represent the com- positional structure of the input. To enhance structured representations, we impose sparsity on these coefficients. The sparse coefficients are then used to linearly combine the decoding dictionary atoms generate the output. Furthermore, assist compositional generalization tasks, we propose estimating the coefficients of the target problem the context examples. We demonstrate the effectiveness of our approach on the S-RAVEN and RAVEN datasets. For certain compositional generalization tasks, our method maintains performance even when standard Transformers fail, owing to its ability to learn and apply compositional rules.
linear combination of the coefficients obtained from
Introduction
Recent advancements in artificial intelligence (AI) have led to significant breakthroughs in various domains 66 [TI [7 26). Models such as large-scale Transformers have demonstrated remarkable capabilities in natural language understanding, image classification, and multimodal reasoning. However, despite these successes, solving in-context compositional learning tasks remains major challenge [0] As illustrated in Figure [] such tasks involve data composed of basic components arranged by underlying compositional rules, requiring models to infer and transfer these structural patterns from context examples while achieving good representation and generalization.
Transformers primarily rely on dense attention mechanisms (26 without an explicit framework for representing compositional rules. As a result, they struggle to capture structured relationships and lack an effective mechanism to transfer these inferred rules across examples. The absence of structural inductive bias limits their ability to generalize tasks that demand compositional understanding.
In this paper, we extend the attention mechanism by explicitly encoding compositional rules, drawing inspiration from the principles of sparse coding. In sparse coding [14], signals are expressed as sparse combinations of basic elements, with the resulting coefficients capturing the compositional structure
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Data
Dictionary
Coefficients
Solving Steps:
Step I: Encode Data (Conlext Target)
2
5
8
?
Step 2: Esttmatte Turget Coefficients
is Target
0
Result
B
Dichioray
MB
Step 3: Decode Coefficients
Compositional Rule
Encoding Decoding
(a)
(b)
Figure 1: Illustration of the in-context compositional learning task. The input data includes both the context tasks and the target task. The goal is to solve the target task by inferring and applying the compositional rule observed the context tasks. (a) Applying the principles of sparse coding to represent the data. Given dictionary, the input data can be sparsely represented using set of coefficients that encode underlying compositional rules. Encoding/devoding data: An example of one task is composed of four elements from the dictionary, with indices "6, O, O, 4." After one-hot embedding. we obtain 4 10 matrix, where each nonzero entry corresponds specific element in the dictionary. By stacking all 9 examples, we obtain a 36 x 10 matrix representing the coefficients. Compositional rules: Each row of the input data follows an underlying pattern. If the first two shapes are constructed as (A, 0, 0, B) and (B, 0, 0, C),, where A, B, and C correspond to unique elements in the dictionary, 0 means an empty shape, then the third shape should be (C, 0, 0, A). (b) Representing effective way to estimate the coefficients of the
the compositional rules as coefficients provides target task from those the context tasks. Once inferred, these coefficients can be decoded into final output using the dictionary. Details of this task are described in Section
of the signal. Specifically, shown in Figure reinterpret the attention mechanism as mapping of inputs into outputs through projections onto two sets of learned dictionary atoms: an encoding dictionary and decoding dictionary. The encoding dictionary decomposes the input into set of coefficients, which represent the compositional structure of the input. The coefficients are then used linearly combine the decoding dictionary atoms generate the output.
In the attention mechanism (26). the attention map is generated by computing the inner product between inputs transformed the query and key matrices. In contrast, approach reinterprets this process as projecting the input onto learned dictionary. i.e., encoding dictionary, parameterized by the query and key matrices obtain the coefficients. To enhance structured representations, introduce sparsity into the coefficients, allowing them to explicitly represent the compositional rules inherent in the input. These sparse coefficients then used to combine another dictionary, i.e.,
decoding dictionary, parameterized by the value matrix, to generate the final output.
By projecting the input of both the context and target tasks onto shared encoding dictionary to obtain their respective coefficients, we can effectively infer the compositional rules of the target tasks. Inspired by the lifting scheme [22] we estimate coefficients of the target task through simple linear combination of the context task coefficients.
We first assess the effectiveness of our method on toy example with simple compositional rule, demonstrating that our approach successfully learns and generalizes the rule, whereas the standard Transformer fails in this case. The results are shown in Figure 3 We then evaluate our method on the
in-context compositional learning dataset, such as S-RAVEN 20 and RAVEN 29). Our approach
consistently outperforms standard Transformer baselines. These results indicate that integrating the attention mechanism with sparse coding enhances the ability of models to learn and apply compositional rules.

---

We summarize our contributions as follows:
We reformulate the attention mechanism, inspired by sparse coding, as a mapping of inputs outputs via projections onto two learned dictionaries: encoding dictionary and decoding dictionary.
We explicitly represent inputs as sparse combinations of the encoding dictionary to encode compositional rules.
We enable effective transfer of compositional rules across tasks by estimating target coeffi- cients via a simple linear combination of context coefficients.
We demonstrate the effectiveness of our approach on in-context compositional learning tasks, maintaining good performance even cases where standard Transformers fail.
Method
In this section, we first outline the problem setting in-context compositional learning and then introduce our framework, inspired by sparse coding, which reformulates the Transformer architecture to better capture compositional structure.
2.1 Preliminary
Problem formulation. We define the in-context compositional learning task as learning function purely from demonstrations provided within context window. Inspired by the RAVEN dataset [29]. we consider setting where the model is given L structured example (the context) and predicts the Th one (the target). We illustrate this task in Figure
Assume each example xy X' is governed
latent compositional rule R. Let the context set be:
C = {x1 T2 ,,, IL-1) x411
(1)
The model must produce & E X' such that á = f(C), where f is a learned model conditioned on the context C. The goal is minimize the expected error distribution of tasks:
min Ec,,, [e()
where l is task-specific loss function. To emphasize in-context compositional learning, the tasks in the distribution are constructed such that: (1) Each rule R. is composed from a finite set of primitive operations P. (2) Test-time tasks involve novel combinations of primitives not seen during training, ie., Ruest & span (Riran). This setting evaluates the model's ' ability to infer latent rules purely from examples and apply them unseen inputs in a compositional manner, mimicking human inductive reasoning in Raven's Progressive Matrices.
Sparse coding. Sparse coding represents signals using linear combinations of an overcomplete dictionary D E gg (where m > d is the number of atoms), and representing the signal as:
X SD,
where X Nxd is the input signal, S IR N is the sparse coefficient vector, which has only few nonzero elements. To achieve sparsity, the common usage the ooftthhresholdig function:
prox(S) = sign (S) O max S - E,
where @ is Hadamard product. It encourages sparsity by shrinking small values S toward zero. Sparse coding is widely used in signal processing, machine learning, and neuroscience, providing efficient and interpretable representations of data.
2.2 Revisiting Transformer Blocks
Multi-head attention (MHA). The attention layer (26) transforms the input sequence X Nxd to the output sequence O E Nxd where N denotes the sequence length, d is the dimension of

---

W xt
XW,
(X) (XX
(
Encoding Dictionury
(X)
Decoding Dictionary
(X)
XW,
XW
XW
X
(X)
(A(
a
Sparsify
oX (X))
(WW, X))
a(X + (X))
a
g(c)
Attention Map
Value
Coefficients
Dictionary
atherwise
Share Compositional
800)
Σ 4,
Z= AXW,
Z (x))
a(a)
05 15 2.. 424
(c)
Figure 2: (a) The attention block produces the output as linear combination of the value matrix, weighted by attention map. (b) Our framework reformulates the attention mechanism: Outputs are constructed as sparse combinations of learned dictionary atoms, i.e., decoding dictionary vx(X), ), and their coefficients a represent compositional rules. (c) Details of our method: The coefficients a are obtained by decomposing the input features over the encoding dictionary (XX), and then achieving sparse representations with nonlincar function (.). Since the coefficients of the target task only provide partial information about its compositional rule due to limited observations, propose to estimate the coefficients of the target task a as a simple linear combination of the context task coefficients, i.e., a' g(ax. Further details are provided in Section 2.3
input and output features. The attention layer projects the input using the corresponding projection matrices W g: We, W. E Rdxdd and calculates the attention map.
A = ATTN(X) (XW, -XXX),
where a--) = softmax(.) The attention map A E NxN captures the token-wise relationship by doing inner-product in space transformed by Wq. W..
Multi-head attention extends this by allowing multiple attention mechanisms work in parallel, with each head independently learning attention patterns. For H attention heads, each attention head
calculates attention maps A (h) o(X ((h)
gk XT), where W w (hoT is corresponding dk
projection matrices w
(h) w a #. 1....,.. H. The multi-head attention represents as,
H
M Σ
A
(h)
(W
Σ
X w (h)
XT)
(W 00
N0
k=1
where (hh)
W!
(A)T
w!!)
w!!)
E # w Rdxd
While the MHA offers a form of learned localization via query-key similarity, it suffers from two fundamental limitations in compositional tasks:
The use of the sof tmax function produces dense attention weights, resulting in indiscrimi- nate global mixing of information. This lack of sparsity hinders the model from representing the compositional structure inherent in contextual tasks.
There is no explicit mechanism for reusing local compositional rules. It struggles to disentangle meaningful subcomponents, limiting its capacity to generalize via transferring compositional rules.

---

2.3 Reformulate Transformer Using Sparse Coding
We propose to explicitly reinterpret the attention in the Transformer as a form of learned, sparse coding problem. We factorize MHA (6) as follows:
H
(X) = Σ
W
(N)
80
h=1
(7)
o(X Encoding dietionary
Decoding dictionary
where (5h)
(X) and (x(h(((X generate set of dictionary atoms conditioned the input X, ()()())
and db (.) are the basis functions parameterized by ( and W Our method is illustrated in g/
Figure
Learned dictionary atoms. Our method reformulates the attention mechanism as composition over learned dictionary atoms enable structured representations. Specifically, we introduce two sets of input-dependent dictionaries: (XX) and (XX). both parameterized by Icarnable functions of the input X.
The encoding dictionary ((X) is used to extract coefficients by computing the product (X))). which represents how the input X decomposed with respect the learned dictionary atoms. These coefficients encode the combination rule underlying the input structure.
The decoding dictionary (x) serves as a reconstruction dictionary that synthesizes the final output from the coefficients.
Both o(X) and (X) are dynamic and data-dependent, allowing the model to adaptively learn dictionary atoms that best represent the compositional patterns in each input instance.
Sparse coefficients. The coefficients X ((X) encode the combination rule underlying the input structure. To enhance the model's 's capability to capture compositional structure, we apply sparsity- promoting nonlincarities (.),, such as seft-thresholding, defined as prox(x) = sign(x) O max([æ] E, O) introduce sparsity in coefficients a, where is the threshold for setting values to zero, i.e.
== a (X ((X)).
(8)
Different from the attention map A, which applies omat operation, sparse coefficients preserve the most informative components while suppressing redundant interactions. The resulting representation is more structured and better aligned with the underlying compositional rules.
Update coefficients of the target task. By encoding the underlying compositional rule as sparse coefficients o, we aim to transfer this rule from context tasks to the target task. The coefficients the target task encode only partial information about its compositional rule due limited observations of itself. We can transfer the compositional rule to the target task by coefficient transfer.
To address this, propose estimating the target coefficients based on those of the context tasks. Inspired by the lifting scheme [22], we devise a procedure that predicts the target task coefficients through linear combination of the context task coefficients. Specifically, sparse coefficients portions derived from
N consist of contributions from both context and target tasks, with
the context tasks and a single portion az E R X xN corresponding to the target task. We can update the coefficients of target tasks az by,
L-1
aL t- aL
X
where A; is the learnable parameter for combining the coefficients of context tasks. The coefficients of context tasks remain unchanged. We represent this operation with function g(-).
g(a)= = a;
context tasks,
(10)
aL + El=1 Ar;; target task.
5

---

Baseline
Ground
Baseline
PSNR: 21.7
PSNR: 81.5
PSNR: 15.8
PSNR: 54.5
Attention Map
Attention Map
?
Question Pancl
Question Panel
B
Compasitional Rule
Baseline
Ours
Compasitional Rule
Baseline
Ours
(a)
(b)
Figure 3: The effectiveness of sparse coefficients (attention map). Models are trained on setting (a) and tested on both setting (a) and novel setting (b), which has a different compositional rule. The baseline method, Transformer with standard MHA. produces blurry outputs due dense coefficients, which lead to mixed and entangled results. In contrast, our sparse coefficients prevent this blurring and effectively transfer the construction rule from the context tasks to the target task. Further details are in Section
This method is parameter-efficient. At each layer, there L learnable parameters A correspond- ing to the number of context tasks, which remains relatively small compared to the overall parameter count of the Transformer blocks.
Variation of basis functions. This formulation allows us to explore various designs for the basis functions (6-) and v²-)) to modulate the expressiveness of the model. A detailed discussion is provided in Appendix
3 Discussion
We construct synthetic dataset designed to evaluate in-context compositional learning. Each input consists of panels. The panels are grouped such that panels 1-3, 4-6, and 7-9 share the same underlying compositional rule, illustrated in Figure The first two rows represent the context tasks, while the last row is the target task, which the model predicts based on the pattern observed in the first examples.
Compositional rules. Each panel an 8 8 binary image composed of four smaller basic shapes, arranged according to predefined rule. For example, in Figure 3 (a), every three panels are composed 3 basic shapes. Denoting these shapes as A, B, C. and use represent an empty position, the panels are arranged from left to right and top to bottom as follows: (A, 0, B), (B. C)., and (C,
A). Similarly, the compositional rule of Figure (b) (0, A, B, 0).
(0, B. C, 0). and (0, C. A, 0). The basic shapes are chosen from a set of 16 elements, allowing for about P(16, 9) = (16-9) ~ 4 x 10" distinct panel configurations. Details of the experimental setting
16
are described in Appendix
Learning configures. A single-layer Transformer block, containing only attention layer, is trained to predict the target panel given the 8 context panels. The model is trained with mean squared error (MSE) loss. The target panel is masked in the input, and the model is optimized to reconstruct it from the context examples. During training, the model exposed to data generated under one compositional rule and evaluated test data generated under different rule assess compositional generalization.
3.1 Effectiveness of Sparse Coefficients
We compare our approach with the baseline, where our method introduces sparsity in the coefficients. As shown in 33 the baseline model with dense attention fails to predict the target panel on the test set and produces only blurry predictions the training data. In contrast, through sparse attention and coefficient transfer, our method effectively infers and applies compositional rules to accurately predict the target panel both training and test data, illustrated in Figure (a) and (b).

---

50
Percentage
40
Training
30
Transforme
51.6 1.3
55.74 1.5
58.14 14
$9.84 1.4
63.32 1,9
65
HYLA 200
55+ 2.1
68.6+
73.24 0.6
72.5 + 6.6
27.1
79.3
1.8
20
Baseline
73.4
76.3
726
Durs, w/ g(o)
Ours
PSNR Threshald
Figure 4: (Table) Accuracy comparison between our method and baseline methods on the Symbolic RAVEN (S-RAVEN) dataset. Our method consistently achieves higher accuracy than baselines. (Plot) Results on the RAVEN dataset. shows the percentage of test samples with PSNR values exceeding a given threshold. At lower PSNR levels, the baseline method performs similarly ours. However, PSNR values above 40, the baseline achieves nearly coverage, whereas our method retains over 30% of the samples.
3.2 Effectiveness of Coefficient Transfer
By representing input X as [X,.... .x,',, where X,, Vi = 1,... , and X E 1 xd corresponding to context tasks and the target task, we have output according to 7
Z
(X, ((X)) (X)
[ay (X)
(1I)
Zr.
oX 6 (X)) (X)
ar (X)
We set X O, where
I2 M. is matrix with all zeros, since no observation for the target task.
Baseline methods. standard Transformer with o()) = softmax('), produces coefficients az = o(Lo((X) sof tmax
N 1, where 1 R ने. xd is a matrix with all ones. It leads to the
output of the target task as
Z
= 1 X 1XJW N
(12)
where 1 1X is average of the input. Estimating the output Zy by simply averaging the inputs N results in blurry output, as illustrated in Figure 3
Our method. Different from standard Transformer, method enforces sparsity in coefficients by applying (-) = prox(-) to obtain a = (Xx(((X)) = prox (0) = O, which produces
(13) This indicates that no estimation of the target output made when there no observation of the input. However, with the coefficient estimation 19 ag f- az + + \çay, we avoid a zero estimation of the target coefficients by linearly combining the coefficients of the context tasks, and produce nonzero output,
Z, 0z, 1 (X) 0.
L-11
Z₁ = az (X) + Σ \oa (X).
(14)
==1
Without coefficient estimation, neither standard Transformer nor our method yields informative outputs for Zz. However, by learning X; and leveraging the accurate reconstruction of context examples by Z,, Vi 1. meaningful outputs that reuse compositional rules from the context tasks. In practice, we observe the sharp and recurrent attention patterns produced by our method, as shown in Figure We provide details of our analysis in Appendix 9
1. Zz az (XX) + - X ((X) is capable to generate
Experiments
4.1 Symbolic RAVEN
The S-RAVEN dataset 20), detailed in the original paper, specifically designed to evaluate composi- tional reasoning, In S- RAVEN, each task is built from finite set of rule combinations systematically
7

---

Ground Truth
Baseline
Ground Truth
Bascline
Ours, w/o ((a)
Ours
Ours, w/o g(a)
Ours
Ground Truth
Bascline
Ground Truth
Bascline
Ours, w/o g(a)
Ours
Ours, w/o ((x)
Ours
(c)
(d)
Figure Example results of RAVEN. The model predicts 9th panel based on the first 8 panels. compare method with and without g(a), the coefficient estimation for the target task, alongside the baseline method. The baseline often yields blurry images with incorrect layouts, whereas our method preserves structure and improves compositional accuracy. However, all models occasionally fail on the most challenging cases, e.g., (d).
applied across panels, with each panel represented as a tuple of integers that symbolically encodes its features. Details of the experimental setting are described Appendix 7
Experimental settings. We train models using the standard decoder-only Transformer architecture and evaluate performance under varying numbers of training tasks and attention layers. Our method builds directly on the S-RAVEN implementation, introducing sparsity into the attention maps and applying lifting scheme to enhance compositional rule transfer within the attention mechanism.
Metries. To evaluate whether a model trained on a subset of rule combinations can generalize to unseen combinations, we partition all possible rule combinations into separate training and test sets, where 25% of the combinations are held out for testing. Model performance is assessed by measuring the accuracy of correctly predicted examples from the test set.
The results, summarized in Table are obtained by running the experiment three times. Our method consistently outperforms baseline approaches, including standard Transformer 26 and HYLA (20), achieving significantly higher accuracy even with fewer layers.
Results.
4.2 RAVEN dataset
The RAVEN dataset 29 was originally designed for visual reasoning, requiring models select the correct answer from eight candidates based on the underlying structure of context panels. In contrast, our work focuses evaluating the compositional capabilities of models by tasking them with generalizing the answer directly, rather than selecting from predefined options-a more challenging objective that demands better understanding and application of the compositional rule.
Experimental settings. For our experiments, we adapt the rule framework from RAVEN and focus on the simplified case where examples are arranged in the target answer based on the composition of the eight context images. We modify the standard Transformer architecture to serve as a baseline and compare its performance against our approach, which incorporates sparsity and estimation of the coefficients into the attention mechanism.
grid. The model generates
Metries. assess model performance, we adapt the Peak Signal-to-Noise Ratio (PSNR) metric to quantify the difference between the generated images and the ground truth. We report the percentage of test samples exceeding PSNR thresholds 20, 30, and 40, where higher percentage indicates better reconstruction quality and overall model performance.

---

Results. As shown in Figure 4 the results demonstrate that our method consistently achieves higher accuracy than the standard Transformer baseline. While the standard Transformer yields nearly 0% of test samples with PSNR above 40, our method maintains around 40%. Additionally, we observe further performance gains when the target coefficient estimation is applied.
Example predictions are visualized in Figure where the baseline model frequently produces blurry images with incorrect arrangements, while our method preserves clear structural information and generates more accurate compositions. Nevertheless, in particularly challenging cases, all models occasionally fail to produce satisfactory outputs.
5 Related Work
In-context compositional learning. Recent research on compositional reasoning with transformers has explored several key directions. Some studies focus on understanding and measuring com- positional generalization abilities, often identifying gaps between LLM performance on known components and novel compositions, and how these gaps evolve with model scale in-context learning 7 16 21]. Other works delve into the underlying mechanisms and offer explanations for how LLMs achieve fail at compositional reasoning, for example, by proposing that attention acts as a hypernetwork or by analyzing emergent algorithmic behaviors [TS] 20 (4]. Another line of inquiry compares the effectiveness of general pre-trat are-training against specialized architectures, investi- gating whether broad pre-training itself can endow models with strong compositional capabilities, sometimes rivaling or exceeding those of systems explicitly designed for such tasks 2 g. However, conventional Transformers often struggle with in-context compositional tasks due to insufficient structural inductive bias. We address this limitation by introducing a sparse coding attention, explicitly designed to capture and transfer structural rules from context examples.
Sparsity in attention. Sparsity has proven to be a powerful principle, and extensive research has investigated its application in Transformers, primarily to reduce the computational complexity (23
of the attention mechanism. Sparse attention mechanisms aim to reduce the number of token pairs being attended to. This includes methods employing fixed, pre-defined sparsity patterns, such as
local windowed attention combined with varying forms of global or random attention [1 3 28
Learnable or adaptive sparsity patterns have been explored, where the attention pattern is dynamically determined, for instance, through Some approaches seek to approximate full attention using kernel methods or low-rank projections, which implicitly reduce computational load without explicit sparse connections 4 27 In contrast to prior work, our approach introduces sparsity in attention mainly to enhance the representation of compositional rules. By replacing softmax with soft-thresholding, we promote the learning of
hashing [12] or learned routing strategies [18]
structured, localized attention patterns that better capture and encode compositional relationships.
6 Conclusion
In this work, we proposed a reformulation of the Transformer architecture to address the challenge of in-context compositional learning. By drawing inspiration from sparse coding, we introduced a frame- work that represents compositional rules as sparse coefficients over learned dictionaries, enhancing the transferability of structure across tasks. By enforcing sparsity in the coefficients and estimating target coefficients from those of the context tasks, our method further enhances rule transfer and localization within the attention mechanism. Experimental results on in-context compositional learning datasets, such as S S-RAVEN and RAVEN benchmark, demonstrate that our approach significantly outperforms standard Transformers, particularly tasks requiring compositional reasoning and generalization to unseen rule combinations. These findings highlight the potential of combining principles from sparse coding and attention to advance structured reasoning in neural models.
Limitations. While our approach shows promising results in training Transformers on relatively small-scale tasks, its application large pre-trained models remains unexplored. Although inte- grating the linear combination of attention maps into pre-trained models could potentially enhance compositional learning, we leave this for future work.

---

References
[1] Joshua Ainslie, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. Etc: Encoding long and structured inputs in transformers. In Empirical Methods in Natural Language Processing. 2020.
[2] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learn- ing algorithm is in-context learning? investigations with linear models. The International Conference on Learning Representations, 2023.
[3] Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv:2004.05150, 2020.
[4] Krzysziof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. The International Conference on Learning Representations, 2021.
[5] Bhishma Dedhia, Michael Chang, Jake Snell, Tom Griffiths, and Niraj Jha. Im-promptu: in- context composition from image prompts. Advances in Neural Information Processing Systems, 2023.
[6] Alexey Dosovitskiy. An image is worth 16x16 words: Transformers for image recognition at scale. The International Conference on Learning Representations, 2021.
[7] Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Sean Welleck, Peter West, Chandra Bhagavatula, Ronan Le Bras, et al. Faith and fate: Limits of transformers on compositionality. Advances in Neural Information Processing Systems, 2023.
[8] Daniel Furrer, Mare van Zee, Nathan Scales, and Nathanael Schärli. Compositional gen- eralization in semantic parsing: Pre-training vs. specialized architectures. arXiv preprint arXiv:2007.08970, 2020.
[9] Arian Hosseini, Ankit Vani, Dzmitry Bahdanau, Alessandro Sordoni, and Aaron Courville. On the compositional generalization gap of in-context learning. In Proceedings of the Fifth BlackboxNLP Workshop on Analyzing and Interpreting Neural Networks for NLP, 2022.
[10] Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. ICLR, 2022.
[11] Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wen Lo, et al. Segment anything. In Proceedings of the IEEE/CVF International Conference on Computer Vision, 2023.
[12] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. The International Conference on Learning Representations, 2020.
[13] Shih- Yang Liu, Chien-Ye Yi Wang, Hongxu Yin, Pavlo Molchanov, Yu-Chiang Frank Wang, Kwang- Ting Cheng, and Min-Hung Chen. Dora: Weight-decompossed low-rank adaptation. In ICML., 2024.
[14] Bruno Olshausen and David Field. Emergence of simple-cell receptive field properties by learning sparse code for natural images. Nature, 1996.
[15] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. Consortium for Reliability and Reproducibility, 2022.
[16] Ofir Press, Muru Zhang, Sewon Min, Ludwig Schmidt, Noah A Smith, and Mike Lewis. Measuring and narrowing the compositionality gap in language models. Empirical Methods in Natural Language Processing, 2023.
[17] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models, In The IEEE/ CVF Computer Vision and Pattern Recognition Conference, 2022.
10

---

[18] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content based sparse attention with routing transformers. Transactions of the Association for Computational Linguisties, 2021.
[19] Abulhair Saparov and He He. Language models are greedy reasoners: A systematic formal analysis of chain-of-thougee thoughh... The International Conference on Learning Representations, 2023.
[20] Simon Schug, Seijin Kobayashi, Yassir Akram, João Sacramento, and Razvan Pascanu. Atten- tion hypernetwork. The International Conference Learning Representations, 2025.
[21] Aarohi Srivastava, Abhinav Rastogi. Abhishek Rao, Abu Awal Md Shoeb, Abubakar Abid, Adam Fisch, Adam R Brown, Adam Santoro, Aditya Gupta, Adrià Garriga-Alonso, et al. Beyond the imitation game: Quantifying and extrapolating the capabilities of language models. The International Conference on Learning Representations, 2025.
[22] Wim Sweldens. The lifting scheme: A construction of second generation wavelets. SIAM journal on mathematical analysis, 1998.
[23] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, 2022.
[24] Eric Todd, Millicent L Li, Arnab Sen Sharma, Aaron Mueller, Byron C Wallace, and David Bau. Function vectors in large language models. The International Conference on Learning Representations, 2024.
[25] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13977 2023.
[26] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 2017.
[27] Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.0477 04766, 2020.
[28) Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 2020.
[29] Chi Zhang, Feng Gao, Baoxiong Jia, Yixin Zhu, and Song Chun Zhu. Raven: dataset for relational and analogical visual reasoning. In The IEFE / CVF Computer Vision and Pattern Recognition Conference, 2019.

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We have provided an accurate discussion of our contribution in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations at the end of the paper.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
12

---

Justification:
We have included assumptions and proof.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We have provided sufficient experimental details the main paper. Additional information is also included in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
13

---

Answer: [Yes]
Justification:
We have provided the code for our main experimental results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https Cooeeeoodeesssseeeeee //nips.com htte:/ubbl///liiiee////eeeueassl ssiniioiscyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: We have provided sufficient experimental details the main paper. Additional information is also included in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification:
We conduct our experiments for multiple runs and include error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
14

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We have provided this information in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics http:p/ps///.capp..c.o/oo//////a
Answer: [Yes]
Justification: Our paper conforms with NeurIPS code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our paper only proposes study to learn compositional generalization tasks. It doesn't have a clear impact on the real-world applications.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
15

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: There is no such risk.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We have cited the models, datasets that used in the paper.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
16

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: No new assets released.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: It is not a crowdsourcing experiment.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
17

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: the core method development in this research does not involve LLMs
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy fattps://neurd/ /ttppp/p...///oooeeeere////////0 for what should or should not be described.
18

---

Experimental Details
The experimental setting of synthetic data. We conduct experiments on synthetic dataset composed of 16 distinct basic elements, which are shown in Figure 6 (a), where each panel is constructed by selecting and combining two of these elements. The examples of training and test data are displayed in Figure (b-c) The model architecture consists of single-layer Transformer using only self-attention, omitting the feedforward layer. The input sequence length is fixed N 32, with feature dimension of 16 and single attention head (H = 1). Training is performed over 200 epochs using a batch size of 128. We optimize the model using Adam optimizer with a learning rate of 0.001 and employ mean squared error (MSE) loss as the training objective.
The experimental setting of S-RAVEN. We evaluate on the S-RAVEN benchmark 20 where each task is defined by features, sampled from pool of possible rules. For each task, we generate three input-output sequences of length three, using random inputs for each rule to form the context. Our model architecture follows HYLA [20], varying the number of layers between 4 and 8. The input has a feature dimension of 128 and 16 attention heads (H 16). For baseline comparisons, including HYLA and standard Transformer, we adopt the original configurations as specified in the HYLA paper. All models use Root Mean Square (RMS) normalization for attention activations. To promote structured representations, our method applies soft thresholding to the attention weights, encouraging sparsity. Training is conducted for one epoch using batch size of 128. the Adam optimizer with learning rate of 0.001 and a weight decay of 0.1, and the cross-entrope loss as the objective.
restricted version of the RAVEN dataset [29], focusing solely on the 2-by-2 grid layout. To ensure deterministic target generation, remove stochastic variations in rotation and color, so that the target panel is uniquely determined by the eight context panels. Each image is resized to 40 x 40 pixels. The model is a standard Transformer with 4 layers, sequence length of N 36, feature dimension of 512, and 16 attention heads (H 16). Training is performed over 2000 epochs with a batch size of 256, using the Adam optimizer with learning rate of 0.0001. The model is trained minimize squared error (MSE) loss.
The experimental setting of RAVEN. We conduct experiments
Experimental Results
Sparsity and threshold To investigate the impact of the threshold attention sparsity, we conduct experiments on the RAVEN dataset. Specifically, we measure the average sparsity of the attention maps across all layers, where sparsity is defined as the proportion of zero-valued entries after thresholding. As the threshold increases, more small-magnitude values are suppressed, leading to higher sparsity levels. Our results confirm this trend: larger thresholds consistently yield sparser attention maps, demonstrating the controllable nature of sparsity in our model through the threshold parameter. The results are shown in Table [
(a) basic elements
(b) An example training data
(c) An example
Figure 6: Examples of the synthetic dataset.

---

Table 1: The effect of threshold on the sparsity of the attention map.
Threshold (g)
0.003
0.01
0.03
0.1
0.3
Sparsity 18.53
57.82
90.45
97.82
99.38
Table 2: Variation of basis functions.
Condigs of o(X) and vi()
W x. w
ReLU(W dt X). w W X. Relu(WW)X))
ReLu/ x). BeLu (eLl(
Accuracy
71.7
72.3
72,9
73.6
Variation of basis functions. With the above formulation, explore different designs for the basis functions of-) and (-) to adjust the expressiveness of models. In the baseline configuration, the basis functions constructed through linear projections of the input, parameterized by w-(h) ql W simple variation is introduce nonlinearity into the basis construction by applying an activation function, such as ReLU, after the linear projections. For instance, different basis function can be redefined as o XX ( X) 15 (X) ReLU(W Do X). Incorporating nonlinearity into the basis functions can increase the representational capacity, enabling the model to capture more complex localized patterns beyond those achievable with purely linear projections.
We conduct experiments on the S- S-RAVEN dataset using 4-layer Transformer architecture, training the model on a dataset of 20 million samples. We compare the different designs of ((X) and (XX) by adding the ReLU. The results are shown in Table 2
Application to language modeling tasks. While our primary focus address specific limitations of attention mechanisms in compositional tasks, we have conducted experiments on language models to demonstrate our method's effectiivnnss on standard benchmarks.
We integrate the proposed sparse-coding inspired attention into the Llama-7B [25] model. We then fine-tuned these modified models on several widely-used commonsense reasoning bench- marks and compared the results against both the original base models and those fine-tuned using
LoRA/DoRA 10 13..
In our implementation. we target the model's Attention blocks. We treat the original attention weights (denoted vi( (X) and ((X)) fixed and introduce our core components: new, learnable parameters for sparsity (E) and the coefficient transfer mechanism (A,). We initialize these new parameters to zero, ensuring that our module has no impact on the model's 's output before training. By fine-tuning only these new parameters, we can cleanly measure the influence of our method.
Our findings show that models incorporating over the base Llama model, which is shown in Table 3 Although these results do not yet surpass those from LoRA and DoRA fine-tuning, it's important consider that our approach
model achieve notable performance improvement
significantly fewer trainable parameters (over a hundred vs. over 50 million) and has not undergone extensive hyperparameter optimization. The performance gains over the base models suggest that large language models benefit from our mechanism on reasoning tasks, providing compelling evidence of its value. We believe that with further refinement, on language modeling tasks. We see the exploration of its application to other benchmarks, such as translation and summarization, promising direction for future work.
approach has the potential to achieve better performance
Evaluate the models on Im-promptu benchmark. We evaluate our method on the Im tn-promptu benchmark [5], including 3D Shapes, BitMoji Faces, and CLEVr Objects datasets. For this com- parison, we adopt the Object-Centric Learner from the original paper our baseline and integrate our approach by modifying its attention layer. As detailed the Table our method consistently achieves a lower MSE, demonstrating improvement over the bascline.

---

Table Results on language modeling tasks.
Model
Params
BoolQ
PIQA
HellaSwag
WinoGrande
ARC-c
OBQA
Avg.
Ilama-7B
56.5
79.8
76.1
70.1
63.2
77.0
70.5
+ Ours
128
57.3
80.7
80.6
71.1
64.2
77.6
71.9
LoRA
55.9M
67.5
80.8
83.4
80.4
62.6
79.1
75.6
LoRA + Ours
55.9M
69.5
81.8
81.6
80.8
65.1
79.0
76.3
DoRA
56.6M
69.7
83.4
87.2
81.0
66.2
79.2
77.8
DoRA + Ours
56.6M
70.0
83.6
87.3
81.2
67.4
78.9
78.1
Table 4: Results on Im-promptu benchmark.
MSE
3D Shapes
BitMoji Faces
CLEVr Objects
OCL
4.36 4.31
4.77
37.54
Ours
4.42
36.23
9 Additional Analysis
By representing an input as [...... ".,"", where X,, 1. 1, --1 and X E RR xd are corresponding to context tasks and the target task, we have,
Z1
(x ((x)) (xx)
an (X)
(15)
[z
oxx 6(X)) (X))
aL (X))
We set X O, where 0 E IR y. xd is matrix with all zeros, since no observation for the target task.
Our method. Different from standard Transformer, method enforces sparsity in coefficients by applying (.) = (rox(-) to obtain αg = ox(xx(( prox(0) = 0, which produces
Zy aL (X) 0.
(16)
This indicates that no estimation of the target output is made when there is no observation of the input. However, with the coefficient estimation (9) Oz <- Σ A,Oa, we avoid a zero estimation of the target coefficients by linearly combining the coefficients of the context tasks, and produce nonzero output,
Z = aL (X) Σ A (X).
(17)
1=1
Without coefficient estimation, neither standard Transformer nor our method yields informative outputs for Z, However, by learning X; and leveraging the accurate reconstruction of context examples by Z,, Vi 1,.... -L-1, Zr ar (X)) + 2-11 - λ; (XX) capable to generate the target outputs that reuse compositional rules from the context tasks.
9.1 Compositional Reconstruction of the Target Output
We have dictionary of basis elements, (X) { Each output Z; for 1, 1,..., is expressed as a linear combination of elements (X() using coefficient vectors a; RRW i.e.
Z, のの = a.
(18)
1=1
where 3 [wh, ... T.
Assumption 9.1. The dictionary (X) is sufficient to represent the target output Zr.
Assumption 9.2. Each of the I. outputs Zı,..., .., Zr- is correctly constructed using coefficient vectors . at-1
3

---

Assumption 9.3, Across , every dictionary element w is used at least once, i.e., vj. there exists such that a; (5) + 0.
Proposition 9.4. There exists set of weights A1...,,A AL-1 such that:
L-1
0
X
(19)
and aL reconstructs nμ using only elements in (X).
R.N denote the set of known coefficient vectors. Let V span(A) C TRN be the subspace spanned by them. Since from Assumption 9.2 each ds reconstructs Z, correctly and the union of their support covers all dictionary elements, the span V includes directions along all dictionary elements used for constructing Zμ.
Proof. Let A { {a,,..., 01--1
From Assumption 9.1
we know there exists some 0j, E RR such that:
Z a 4..
(20)
Because supp(art) S - supp( i.e., the dictionary elements needed for 2,, have already been used in A. and all such directions are already present in V, it follows that:
ai E
(21)
Therefore, there exist scalars ......, dL-1 such that:
L-1
aj =
A
Thus, by setting az ELs L--1 \ obtain the desired coefficient vector such that:
2, = =f=
Given that the dictionary is sufficient, and the L outputs collectively utilize all necessary dictionary elements, the coefficient vector for the L-th output can be expressed as a linear combination of previous coefficient vectors. This demonstrates the ability to transfer compositional rules from context examples to new tasks via linear combination coefficients.
10 Computational Resource
We conducted development and experiments on Linux workstation equipped with single NVIDIA A5000 GPU (24GB memory). A single run of the synthetic task typically takes 3-5 minutes, while single S-RAVEN experiment run takes between 60 and 200 minutes. For RAVEN experiments, full run requires approximately 200 minutes.


Reviews:
[{"review_id": "murHNvVM18", "reviewer": "Reviewer_1bo9", "summary": "This paper proposes an improved Transformer architecture for in-context compositional learning, which enhances the compositional generalization ability of Transformers. Specifically, a novel attention mechanism inspired by sparse coding is proposed, where inputs are projected onto an encoding dictionary to produce sparse coefficients that represent compositional rules. These coefficients combine elements from a decoding dictionary to generate the final output. To enable rule transfer, the sparse coefficients of the target task are estimated via a linear combination of the context coefficients. The proposed architecture outperforms standard Transformers on S-RAVEN and RAVEN, demonstrating stronger compositional generalization ability.", "questions": "In Figure 3, the baseline model generates relatively blurry images. Would the results change if the MSE loss is replaced with the cross-entropy loss, which is commonly used when training Transformers?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths\n\n1. This paper identifies a key weakness that Transformers lack the inductive bias to solve compositional generalization tasks. The motivation is practical and attempts to solve an important challenge in current AI research.\n\n2. The attention mechanism based on sparse coding, along with the learnable encoding and decoding dictionaries, provides an intuitive and partially interpretable way to capture discrete compositional rules.\n\n3. The overall architecture is relatively simple, suggesting that the architecture is potentially scalable and extensible.\n\nWeaknesses\n\nMy main concerns lie in the evaluation.\n\n1. A key of the architecture is its ability to handle novel rule compositions at test time. However, the RAVEN dataset does not seem to support this setting. As far as I know, in the standard RAVEN configuration, the training and test splits differ only in the values of image attributes, not in the combinations of rules and attributes. This means the rule-attribute compositions in the test set are already seen during training.\n\n2. The PGM dataset [1] may be more appropriate. It includes held-out splits where rule combinations are disjoint between training and testing. I suggest the authors run experiments under the held-out splits of PGM to better validate the methods.\n\n3. The paper lacks discussion and comparison with models that address matrix reasoning tasks. Many existing models, such as neuro-symbolic methods like PrAE [2] and NVSA [3], are capable of in-context compositional learning and compositional generalization in matrix reasoning tasks like RAVEN.\n\n4. [4] discusses compositional generalization in a broader context and introduces a benchmark to evaluate the compositional generalization ability. I recommend the authors compare with the method proposed in [4] and evaluate the models on their benchmark.\n\n[1] Measuring abstract reasoning in neural networks. ICML 2018.\n\n[2] Abstract spatial-temporal reasoning via probabilistic abduction and execution. CVPR 2021. \n\n[3] A neuro-vector-symbolic architecture for solving raven’s progressive matrices. Nature Machine Intelligence, 2023.\n\n[4] Im-promptu: in-context composition from image prompts. NeurIPS 2023.", "was_revised": true, "final_justification": "The response has addressed my concerns. The experiments on more datasets or benchmarks improve the applicability of the proposed method, and the authors also indicated that these experiments will be included in the final version. So I raise my score to borderline accept.", "rebuttal": "Thank you for the constructive comments. Below are our responses to each of your comments. We will incorporate this feedback into our revised work.\n\n**Q1. The standard RAVEN dataset is inadequate for testing novel rule compositions.**\n\nA: Different from the standard RAVEN dataset, we used two distinct datasets in our paper to evaluate different aspects of compositional learning:\n\n1. S-RAVEN: As stated in the paper, this dataset is explicitly designed to evaluate compositional reasoning. Our experimental setup for S-RAVEN involved holding out 25% of all possible rule combinations for the test set, ensuring that the model was evaluated on novel compositions not seen during training.\n\n2. RAVEN: For the visual RAVEN dataset, our primary goal was to evaluate the model's ability to generate the correct answer *directly* from the context images, rather than selecting from a list of candidates. This generative task is inherently more challenging than the standard discriminative format and serves as a strong test of the model's ability to understand and apply the underlying compositional rule, even if the rule structure was seen during training.\n\nWe believe this two-pronged approach allowed us to test our model's ability to generalize to both unseen rule structures and to perform a challenging generative task.\n\n**Q2 & Q3. Recommend using the held-out splits from the PGM dataset for more robust validation. Provide discussion of relevant neuro-symbolic models for matrix reasoning, such as PrAE and NVSA.**\n\n\nA: To evaluate our approach, we integrate our coefficient transfer mechanism, g($\\alpha$), into the MRNet model [1]. We then compare the performance of our enhanced model against the original MRNet as well as other baselines, including PrAE and NVSA, on both the RAVEN and PGM datasets. The logic and details of our experimental design are provided below.\n\n\n| PGM | Accuracy |\n|--------|----------|\n| MRNet   | 68.34    |\n| NVSA   | 68.30    |\n| Ours   |    **68.92**      |\n\nOur approach demonstrates better performance across both benchmarks. Specifically, it outperforms MRNet and NVSA on the PGM dataset and shows consistent improvement over PrAE and MRNet on the RAVEN dataset.\n\n|   RAVEN   |  Avg  | Center  | 2x2 grid  | 3x3 grid  |  L-R  |  U-D  | O-IC  | O-IG |\n|:----------:|:-----:|:-------:|:---------:|:---------:|:-----:|:-----:|:-----:|:----:|\n| ResNet+DRT | 59.6  |  58.1   |   46.5    |   50.4    | 65.8  | 67.1  | 69.1  | 60.1 |\n|    PrAE    |  60.3 |   70.6  |    83.5   |    30.5   |  88.5 |  89.2 |  37.3 | 22.7 |\n|    NVSA    |  87.7 |   99.7  |    93.5   |    57.1   |  99.8 |  99.7 |  98.6 | 65.4 |\n|    MRNet   |  74.7 |   96.2  |    49.1   |    45.9   |  93.7 |  94.2 |  92.5 | 51.3 |\n|    Ours    |  77.8 |   96.7  |   53.1    |    51.6   |  95.5 |  95.8 |  94.3 | 57.7 |\n\n\n**Experiment Details:** Our model's primary goal is to test the in-context compositional learning capability of a *standard Transformer*. We intentionally kept the architecture simple to isolate this ability, rather than introducing complex, task-specific designs to maximize performance on a specific benchmark.\n\nThis focus differs from previous state-of-the-art works on the RAVEN and PGM datasets, which are typically designed for classification. These models often rely on complex reasoning blocks built upon CNN backbones like ResNet. A direct comparison is challenging:\nSimply swapping the CNN backbone in these models with a ViT causes a significant performance drop (e.g., from 57.1% to 43.5% on RAVEN). This is because ViT modules are data-hungry and struggle to generalize on these structured reasoning tasks without extensive pre-training or data.\nOur method is not designed to be a plug-in replacement for the specialized reasoning blocks in those models. While our approach does improve the performance of a vanilla ViT baseline (e.g., from 43.5% to 47.6%), it doesn't close the gap with highly engineered, CNN-based architectures.\n\nTo demonstrate the versatility of our coefficient transfer mechanism beyond generative tasks, we adapt it for classification by integrating it into the MRNet baseline. We select MRNet because its multi-head attention model is highly compatible with our approach. The integration is straightforward: we apply our coefficient transfer mechanism, g($\\alpha$), directly to the panel-level features extracted by MRNet's panel-processing blocks.\n\n\n[1] Benny, Y., Pekar, N. &Wolf, L. Scale-localized abstract reasoning. CVPR (2021).\n\n\n\n\n**Q4. Evaluate the models on Im-promptu benchmark.**\n\nA: We evaluate our method on the Im-promptu benchmark, including 3D Shapes, BitMoji Faces, and CLEVr Objects datasets. For this comparison, we adopt the Object-Centric Learner from the original paper as our baseline and integrate our approach by modifying its attention layer. As detailed in the table below, our method consistently achieves a lower MSE, demonstrating an improvement over the baseline.\n\n\n| MSE      | 3D Shapes | BitMoji Faces | CLEVr Objects |\n|----------|-----------|---------------|---------------|\n| OCL | 4.36      | 4.77          | 37.54         |\n| Ours     | **4.31**      | **4.42**          | **36.23**         |\n\n\n\n**Q5. Will cross-entropy loss help blurry images of the baseline model?**\n\nA: Yes, using cross-entropy loss results in visually sharper outputs. However, it does not improve the performance on our primary evaluation metric, PSNR. Since PSNR is our quantitative standard for image reconstruction quality, MSE is the more appropriate choice. Cross-entropy loss is most straightforwardly applied to binary or one-hot categorical data. The RAVEN dataset, however, contains non-binary images with grayscale values. Adapting cross-entropy for this type of continuous data is non-trivial and can introduce training instability. Given these considerations, we choose MSE loss as it provides a more stable and suitable objective for the specific data and evaluation metrics used in our paper."}, {"review_id": "Ylo9ZTie3q", "reviewer": "Reviewer_Dgmp", "summary": "The paper builds on sparse coding to propose a new attention mechanism by which transformers perform better at compositional tasks. Precisely the authors propose to reformulate the attention mechanism as an encoding and a decoding dictionary. The encoding dictionary transforms an input into a sparse set of compositional endoding atoms, while the decoding dictionary uses the previous coefficients to recombine decoding atoms, generating an output. At test time, the coefficients are updated by taking into account  both the coefficients of the target task and of context tasks.\n\nThe paper asesses the method by evaluating it first on a toy example and on the RAVEN dataset, both of them being examples of in-context compositional tasks. Overall, their proposal outperforms the standard Transformer architecture.", "questions": "Why is prox(x) used instead of training with L1  regularization on the coefficient values? \nHow do the results change as a function of atoms in both dictionaries? And as a function of dimensionality?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 2, "strengths_and_weaknesses": "## Strengths\n\n- The paper is well-written, the motivation, goal and results are clearly conveyed\n- The method is simple and provides with an interesting reformulation of self-attention\n- The proposal shows a remarkable improvement with respect to the baseline Transformer architecture\n\n## Weaknesses\n\n- Not clear why the toy example is relevant, it seems to me like it is redundant with RAVEN.\n- The work is very similar to \"Attention as a Hypernetwork\" paper. The paper is cited, but I would like to see a more direct explanation on how both proposals differ.\n- No application to language, where a compositional inductive bias could also be beneficial, while showing whether the method scales to bigger network sizes.\n- I would like to see more results on atom number and hidden dimensionality, and how these affect performance.", "was_revised": true, "final_justification": "The paper, while similar to HYLA, proposes an interesting interpretation of self-attention including sparsity as an inductive bias.", "rebuttal": "Thank you for the insightful feedback. We will incorporate this feedback into our revised work.\n\n**Q1. Clarification of the toy example.**\n\nA: The toy example in Figure 3 serves a distinct and crucial purpose that complements the main RAVEN experiments. While the RAVEN datasets demonstrate the model's performance on complex, established benchmarks, the toy example is a controlled, synthetic environment designed for better illustration and detailed analysis:\n\n1. It provides an interpretable visualization of the core problem. More importantly, it allows us to test for compositional generalization in a direct way. The model is trained on one compositional rule (Figure 3a) and tested on an unseen rule (Figure 3b). This demonstrates the model's ability to infer and transfer abstract rules, which is the central claim of our paper.\n\n2. The simplicity of the toy example enables a step-by-step mathematical analysis of why the standard Transformer fails and our method succeeds. As detailed in Section 3.2, we show how the baseline's softmax operation on a masked input leads to an indiscriminate averaging of features, resulting in a *blurry output*. In contrast, our method's use of a sparsity-promoting function ($prox(x)$) and the coefficient transfer mechanism (g($\\alpha$)) avoids this issue and produces a sharp, correct output. This level of granular analysis would be difficult to illustrate with the complexity of the full RAVEN dataset.\n\n\n**Q2. Compare with \"Attention as a Hypernetwork (HYLA)\".**\n\nA: While both approaches reinterpret the attention mechanism, our proposal differs in two fundamental ways that stem from our core inspiration: sparse coding.\n\n1. The central contribution of our work is the introduction of explicit sparsity into the attention coefficients. We replace the standard softmax function with a sparsity-promoting nonlinearity like soft-thresholding ($prox$). This is a key architectural change designed to produce structured, disentangled representations that capture compositional rules. HYLA does not enforce this type of sparsity. The empirical results in our paper (Table in Figure 4) show our method consistently outperforms HYLA, suggesting that this explicit sparsity is a beneficial inductive bias for these tasks.\n\n2. We introduce an explicit mechanism, g($\\alpha$), to transfer compositional rules from context to target. This step estimates the target coefficients as a learnable linear combination of the context coefficients. This mechanism is directly inspired by the lifting scheme and provides a clear procedure for rule transfer. This explicit transfer step is a unique component of our framework.\n\n**Q3. Application to language model.**\n\nA: We tested this approach on language benchmarks, including commonsense reasoning (PIQA, HellaSwag, WinoGrande) and in-context question-answering (BoolQ, ARC-c, OBQA) tasks. The details of our experimental design are provided below.\n\n| Model | Params | BoolQ | PIQA | HellaSwag | WinoGrande | ARC-c | OBQA | Avg. |\n|:------------:|:------:|:--------:|:------:|:-------:|:----------:|:----------:|:----------:|:----------:|\n| Llama-7B | - | 56.5 | 79.8 | 76.1 | 70.1 | 63.2 | 77.0 | 70.5 |\n| Llama-7B (Ours) | 128 | 57.3 | 80.7 | 80.6 | 71.1 | 64.2 | 77.6 | 71.9 |\n| LLama-7B-LoRA | 55.9M | 67.5 | 80.8 | 83.4 | 80.4 | 62.6 | 79.1 | 75.6 |\n| LLama-7B-DoRA | 56.6M | 69.7 | 83.4 | 87.2 | 81.0 | 66.2 | 79.2 | 77.8 |\n| Llama-2-7B | - | 57.4 | 78.8 | 77.2 | 69.2 | 63.9 | 78.6 | 70.9 |\n| Llama-2-7B (Ours) | 128 | 58.8 | 79.3 | 80.4 | 73.2 | 64.1 | 80.0 | 72.6 |\n| LLama-2-7B-LoRA | 55.9M | 69.8 | 79.9 | 83.6 | 82.6 | 64.7 | 81.0 | 76.9 |\n| LLama-2-7B-DoRA | 56.6M | 71.8 | 83.7 | 89.1 | 82.6 | 68.2 | 82.4 | 79.6 |\n\nOur method delivers a notable performance improvement over the original Llama [1] and Llama-2 [2] models. While not yet surpassing fune-tuning methods such as LoRA [3] or DoRA [4], our approach is significantly more parameter-efficient and has been achieved without extensive hyperparameter tuning. We believe that with further refinement, our approach has the potential to achieve better performance on language modeling tasks.\n\n**Experiment Details:** The scope of this work was to propose and validate our sparse coding-inspired attention mechanism in a domain where the need for compositional reasoning is clear and standard Transformers struggle. The visual reasoning tasks in S-RAVEN and RAVEN provided a controlled environment to establish the effectiveness of our approach.\n\nTo evaluate our method on large pre-trained models, we adapt it into a lightweight component for fine-tuning, using LLaMA as our foundation model. Our approach integrates into the existing architecture by adding a new module while keeping the vast majority of the pre-trained weights frozen.\n\nSpecifically, we target the model's Attention blocks. We treat the original attention weights (denoted as $\\phi(X)$ and $\\psi(X)$) as fixed and introduce our core components: new, learnable parameters for sparsity ($\\xi$) (equation (4) in our paper) and coefficient recombination ($\\lambda_i$) (equation (9) in our paper). We initialize these new parameters to zero, ensuring that our module has no impact on the model's output before training. By fine-tuning only these new parameters, we can measure the influence of our method.\n\n[1] Touvron, Hugo, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière et al. \"Llama: Open and efficient foundation language models.\" arXiv (2023).\n\n[2] Touvron, Hugo, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov et al. \"Llama 2: Open foundation and fine-tuned chat models.\" arXiv (2023).\n\n[3] Hu, Edward J., Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. \"Lora: Low-rank adaptation of large language models.\" ICLR (2022).\n\n[4] Liu, S.Y., Wang, C.Y., Yin, H., Molchanov, P., Wang, Y.C.F., Cheng, K.T. and Chen, M.H., Dora: Weight-decomposed low-rank adaptation. ICML (2024).\n\n\n**Q4. Effect of atom number and hidden dimensionality on performance.**\n\nA: We conducted a study on the S-RAVEN dataset to analyze the impact of atom count and hidden dimensionality. Our results indicate a direct correlation between these parameters and model accuracy; performance consistently improves as the number of atoms increases from 16 to 128, and likewise, as the hidden dimension expands from 16 to 128. As detailed in the table, this improved accuracy reflects an increase in model capability, which comes at the cost of more parameters per layer.\n\n| Num of atoms | 16    | 32    | 64   | 128   |\n|--------------|-------|-------|------|-------|\n| Accuracy     | 50.12 | 60.67 | 63.1 | 65.68 |\n|Params|35.8K|38.9K|45.1K|57.3K|\n\n| Hidden Dim | 16   | 32   | 64  | 128  |\n|------------|-------|-------|------|-------|\n| Accuracy   | 45.03 | 55.32 | 63.1 | 65.12 |\n|Params|11.3K|22.6K|45.1K|90.2K|\n\n**Q5. Why is prox(x) used instead of training with L1 regularization on the coefficient values?**\n\nA: The proximal operator $prox(x)$ is the mathematical tool we used to solve an optimization problem that includes L1 regularization. We add the L1 penalty ($\\lambda ||x||_1$) to our objective function to encourage sparsity, but since this term is non-differentiable everywhere, we can't use standard gradient descent. Instead, we use the proximal operator, where the $prox(x)$ step is the specific operation that correctly applies the L1 penalty by shrinking coefficients and forcing small ones to zero."}, {"review_id": "GkeH3qAyJc", "reviewer": "Reviewer_Qnor", "summary": "The paper adopts a sparse-codes stance on ICL compositional generalisation. The authors reformulate the Transformer attention module so that the latent rule is expressed through sparse weights over a learned dictionary of \"atoms\". They show the proposed method is more parameter-efficient and slightly outperforms a hyper-network-based baseline. Experiments are conducted solely on the symbolic RAVEN dataset.", "questions": "Please see the Weaknesses", "limitations": "yes", "rating": 4, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "### Strengths\n- The paper is very well written.  \n- The paper is accompanied by code.  \n- The proposed method appears to be more parameter- and data-efficient than HYLA baseline.  \n\n### Weaknesses\n- **Lack of a theoretical argument.**  \n  The paper would benefit from a formal treatment of the problem, discussing when this type of generalisation is achievable and how the proposed method advances toward it.  \n\n- **Limited experimental evaluation.**  \n  I acknowledge the authors statement in the checklist that \"our paper […] doesn't have a clear impact on real-world applications\". Nevertheless, the evaluation would benefit from experiments on more complex data beyond RAVEN. It remains unclear whether the learned dictionary can scale or generalise to a mixture of tasks (e.g., training on both RAVEN and Multi-dSprites) or to an arbitrary number of ICL demonstrations.   \n\n- **Missing connection to object-centric work?**   \n  There is a large body of object-centric learning literature, with (I believe) a direct link between the concept of \"atoms\" and \"slots\". Some work explicitly frames compositional generalisation via \"local sparsity patterns in Jacobians\" with proposed attention regualrization [1, 2]. Because RAVEN dataset can be viewed as an example of higher-order interactions across context examples, the paper would benefit from distinguishing itself from OCL work and discussing existing approaches to task learning in that area.   \n\n[1] Brady et al., *Interaction Asymmetry: A General Principle for Learning Composable Abstractions*  \n[2] Brady et al., *Provably Learning Object-Centric Representations*", "was_revised": true, "final_justification": "After reading other reviews and the rebuttal from the authors, who presented more evidence of the effectiveness of their method across diverse datasets and models, I decided to maintain my score still leaning towards acceptance.", "rebuttal": "Thank you for the constructive comments. We will incorporate this feedback into our revised work.\n\n**Q1. Provide theoretical justification for the generalization of the method.**\n\nA: We provide a formal analysis in Section 3.2 that grounds our method's success. We demonstrate mathematically that for unseen targets, standard softmax attention produces a blurry average, whereas our coefficient transfer mechanism, g($\\alpha$), can precisely construct the solution by reusing rules from the context. This advantage is further supported by a proof in the appendix, which guarantees a correct solution when the target information is available in the context. We agree that expanding this specific analysis into a more general theory is an important direction for future work.\n\n**Q2. Application to more complex data.**\n\nA: We validated our method's applicability by evaluating its compositional generalization on the Im-promptu benchmark and its performance in language modeling applications.\n\n**Im-promptu benchmark:** Including 3D Shapes, BitMoji Faces, and CLEVr Objects datasets [1]. For this comparison, we adopted the Object-Centric Learner from the original paper as our baseline and integrated our approach by modifying its attention layer. As detailed in the table below, our method consistently achieves a lower MSE, demonstrating an improvement over the baseline.\n\n\n| MSE      | 3D Shapes | BitMoji Faces | CLEVr Objects |\n|----------|-----------|---------------|---------------|\n| OCL | 4.36      | 4.77          | 37.54         |\n| Ours     | **4.31**      | **4.42**          | **36.23**         |\n\n**Language Model**: We tested this approach on language benchmarks, including commonsense reasoning (PIQA, HellaSwag, WinoGrande) and in-context question-answering (BoolQ, ARC-c, OBQA) tasks.\n\n| Model | Params | BoolQ | PIQA | HellaSwag | WinoGrande | ARC-c | OBQA | Avg. |\n|:------------:|:------:|:--------:|:------:|:-------:|:----------:|:----------:|:----------:|:----------:|\n| Llama-7B | - | 56.5 | 79.8 | 76.1 | 70.1 | 63.2 | 77.0 | 70.5 |\n| Llama-7B (Ours) | 128 | 57.3 | 80.7 | 80.6 | 71.1 | 64.2 | 77.6 | 71.9 |\n| LLama-7B-LoRA | 55.9M | 67.5 | 80.8 | 83.4 | 80.4 | 62.6 | 79.1 | 75.6 |\n| LLama-7B-DoRA | 56.6M | 69.7 | 83.4 | 87.2 | 81.0 | 66.2 | 79.2 | 77.8 |\n| Llama-2-7B | - | 57.4 | 78.8 | 77.2 | 69.2 | 63.9 | 78.6 | 70.9 |\n| Llama-2-7B (Ours) | 128 | 58.8 | 79.3 | 80.4 | 73.2 | 64.1 | 80.0 | 72.6 |\n| LLama-2-7B-LoRA | 55.9M | 69.8 | 79.9 | 83.6 | 82.6 | 64.7 | 81.0 | 76.9 |\n| LLama-2-7B-DoRA | 56.6M | 71.8 | 83.7 | 89.1 | 82.6 | 68.2 | 82.4 | 79.6 |\n\nOur method delivers a notable performance improvement over the original Llama [2] and Llama-2 [3] models. While not yet surpassing fune-tuning methods such as LoRA [4] or DoRA [5], our approach is significantly more parameter-efficient and has been achieved without extensive hyperparameter tuning. We believe that with further refinement, our approach has the potential to achieve better performance on language modeling tasks.\n\n[1] Dedhia, Bhishma, Michael Chang, Jake Snell, Tom Griffiths, and Niraj Jha. \"Im-promptu: in-context composition from image prompts.\" NeurIPS (2023).\n\n[2] Touvron, Hugo, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière et al. \"Llama: Open and efficient foundation language models.\" arXiv (2023).\n\n[3] Touvron, Hugo, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov et al. \"Llama 2: Open foundation and fine-tuned chat models.\" arXiv (2023).\n\n[4] Hu, Edward J., Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. \"Lora: Low-rank adaptation of large language models.\" ICLR (2022).\n\n[5] Liu, S.Y., Wang, C.Y., Yin, H., Molchanov, P., Wang, Y.C.F., Cheng, K.T. and Chen, M.H., Dora: Weight-decomposed low-rank adaptation. ICML (2024).\n\n\n\n**Q3. Discussion of object-centric work.**\n\n\nA:  Both the \"dictionary atoms\" in our framework and the \"slots\" in OCL literature serve as learned primitives for composing complex inputs.\n\nOur approach is distinguished from most OCL work by its motivation and mechanism: Our method is explicitly derived from the principles of sparse coding. The goal was to reformulate the entire attention block within a Transformer as a sparse decomposition and reconstruction process. OCL models, while also focused on compositionality, often employ different architectural paradigms built from the ground up, such as iterative refinement or specialized slot-based attention mechanisms.\n\nOur core idea is the replacement of softmax with a sparsity-promoting function ($prox$) and the introduction of an explicit, learnable coefficient transfer scheme (g($\\alpha$)). This makes our method a direct modification of the standard Transformer block, aimed at improving its inherent compositional reasoning.\n\nWe will incorporate a discussion of OCL in the related work section of a revised version."}, {"review_id": "l7fnNNXsZC", "reviewer": "Reviewer_i17N", "summary": "The paper proposes a drop-in replacement of the attention mechanism motivated by sparse-coding. The proposed mechanism replaces the softmax with soft-thresholding. The authors rename the attention as coefficients. Given a context of L tokens, the coefficient of token at position L+1 may be defined as a weighted average of the coefficients of the prior token. The authors demonstrate their mechanism in toy tasks that demonstrate compositional generation.", "questions": "The key concerns are listed above.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The authors study the problem of compositional via sparse-coding. The idea of using sparsity to induce structure and improve generalization is well-motivated in the literature.\n\nHowever, this article have several issues as it currently stands.\n\nThe authors assert strong claim which are not backed by citations or evidences. For example, the authors assert that the softmax function produces dense attention weights [..., which] hinders the model from representing compositional structure inherent in contextual tasks.\" However, there are several example of transformers that use soft-attention and can generalize outside of its training distributions. I can use any LLMs and make them exhibit systematic generalization.\n\nThe paper lacks clarity. For example,\n- the authors re-use the symbol $\\sigma$ for both the softmax and the \"soft-thresholding\" function,\n- the authors defines $l$ as a task-specific loss function. However, $l$ is not indexed by the task index. Is it really task-specific? This is ambiguous.\n\nThe paper is limited to compositional tasks and does not demonstrate the proposed mechanism on standard benchmarks. To be convincing, the authors should compare the method on standard benchmarks (e.g. translation tasks, summarization, etc.) and they should compare against baseline numbers published in the literature. Without such comparison, it is hard to assess the potential impact of the proposed mechanism.", "was_revised": true, "final_justification": "Experiments on fine-tuning addressed my major concern.", "rebuttal": "Thank you for the constructive comments. We will incorporate this feedback into our revised work.\n\n**Q1. The softmax attention hinders compositional generalization is an unsubstantiated argument.**\n\n\nA: We acknowledge that many powerful models using standard softmax attention can achieve systematic generalization. Our central argument is that explicitly introducing sparsity into the attention mechanism provides distinct advantages for representing compositional structure. While dense attention can learn these relationships implicitly, our work investigates whether adding an explicit inductive bias for sparsity can lead to a more robust or interpretable model of compositional rules. The claim is supported by both a formal analysis and empirical evidence within this specific problem setting:\n\n1. In Section 3.2, we provide a mathematical breakdown of what happens when a standard Transformer must predict a target based on context examples with no prior observation of the target itself (i.e., a masked input). In this scenario, the softmax operation on a zero input produces a uniform, dense attention map. This forces the model to create the output for the target task by simply averaging the features from the context tasks, leading to the \"blurry output\" seen in our experiments. Our proposed method, which uses a sparsity-promoting function, explicitly avoids this failure mode.\n\n2. The experiments in Figure 3 directly support our claim. The baseline model, using standard softmax attention, fails to generalize to a novel compositional rule and produces blurry, averaged outputs. In contrast, our model, under the exact same conditions, successfully generalizes and produces a sharp, correct output. This comparison highlights a clear limitation of dense attention for this type of task.\n\n3. While LLMs are powerful, the challenge of compositional generalization is not entirely solved. As noted in our related work section, several studies have identified gaps between LLM performance on known components and novel compositions.\n\nWe will revise the manuscript to clarify that our work aims to demonstrate the benefits of adding sparsity, rather than arguing that standard softmax attention is inherently flawed.\n\n\n\n\n**Q2. Clarification: (a) The authors reuse the symbol for both the softmax and the \"soft-thresholding\" function. (b) The definition of a task-specific loss function is ambiguous.**\n\nA: (a) Our intention was to use $\\sigma$ as a general placeholder for the non-linear activation function that processes the pre-computed coefficients. In the standard Transformer, this function is softmax. In our proposed model, this function is a sparsity-promoting one, such as $prox(x)$ (soft-thresholding). We will clarify that in our revised paper. \n\n(b) The term \"task-specific\" was intended to convey that the form of the loss function $l$ is chosen to be appropriate for the overall problem domain being addressed (e.g., image reconstruction for the RAVEN experiments), which is not indexed by the task index. Within a given set of experiments (e.g., all RAVEN tasks), the same loss function—MSE in our case —is used for all individual tasks.\nTherefore, the loss $l(f(C), x_L)$ is dependent on the specific task instance because its arguments (the model's prediction and the ground truth) are determined by that task, even though the mathematical form of the loss function itself remains fixed across all tasks drawn from that distribution. We will revise the wording to \"a loss function appropriate for the task\" to make this clearer.\n\n**Q3. Application to language modeling tasks.**\n\n\nA: While our paper's primary focus was to address specific limitations of attention mechanisms in compositional tasks, we have conducted new experiments on language models to demonstrate our method's effectiveness on standard benchmarks.\n\nWe integrated our proposed sparse-coding inspired attention into the Llama-7B [1] and Llama-2-7B [2] models. We then fine-tuned these modified models on several widely-used commonsense reasoning benchmarks and compared the results against both the original base models and those fine-tuned using LoRA/DoRA [3, 4].\n\nOur findings show that models incorporating our model achieve a notable performance improvement over the base Llama and Llama-2 models. Although these results do not yet surpass those from LoRA and DoRA fine-tuning, it's important to consider that our approach uses significantly fewer trainable parameters (over a hundred vs. over 50 million) and has not undergone extensive hyperparameter optimization. The performance gains over the base models suggest that large language models benefit from our mechanism on reasoning tasks, providing compelling evidence of its value. We believe that with further refinement, our approach has the potential to achieve better performance on language modeling tasks. We see the exploration of its application to other benchmarks, such as translation and summarization, as a promising direction for future work.\n\n| Model | Params | BoolQ | PIQA | HellaSwag | WinoGrande | ARC-c | OBQA | Avg. |\n|:------------:|:------:|:--------:|:------:|:-------:|:----------:|:----------:|:----------:|:----------:|\n| Llama-7B | - | 56.5 | 79.8 | 76.1 | 70.1 | 63.2 | 77.0 | 70.5 |\n| Llama-7B (Ours) | 128 | 57.3 | 80.7 | 80.6 | 71.1 | 64.2 | 77.6 | 71.9 |\n| LLama-7B-LoRA | 55.9M | 67.5 | 80.8 | 83.4 | 80.4 | 62.6 | 79.1 | 75.6 |\n| LLama-7B-DoRA | 56.6M | 69.7 | 83.4 | 87.2 | 81.0 | 66.2 | 79.2 | 77.8 |\n| Llama-2-7B | - | 57.4 | 78.8 | 77.2 | 69.2 | 63.9 | 78.6 | 70.9 |\n| Llama-2-7B (Ours) | 128 | 58.8 | 79.3 | 80.4 | 73.2 | 64.1 | 80.0 | 72.6 |\n| LLama-2-7B-LoRA | 55.9M | 69.8 | 79.9 | 83.6 | 82.6 | 64.7 | 81.0 | 76.9 |\n| LLama-2-7B-DoRA | 56.6M | 71.8 | 83.7 | 89.1 | 82.6 | 68.2 | 82.4 | 79.6 |\n\n\n**Experiment Details:** \nIn our implementation, we target the model's Attention blocks. We treat the original attention weights (denoted as $\\phi(X)$ and $\\psi(X)$) as fixed and introduce our core components: new, learnable parameters for sparsity ($\\xi$) (equation (4) in our paper) and the coefficient transfer mechanism ($\\lambda_i$) (equation (9) in our paper). We initialize these new parameters to zero, ensuring that our module has no impact on the model's output before training. By fine-tuning only these new parameters, we can cleanly measure the influence of our method.\n\n\n[1] Touvron, Hugo, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière et al. \"Llama: Open and efficient foundation language models.\" arXiv (2023).\n\n[2] Touvron, Hugo, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov et al. \"Llama 2: Open foundation and fine-tuned chat models.\" arXiv (2023).\n\n[3] Hu, Edward J., Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. \"Lora: Low-rank adaptation of large language models.\" ICLR (2022).\n\n[4] Liu, S.Y., Wang, C.Y., Yin, H., Molchanov, P., Wang, Y.C.F., Cheng, K.T. and Chen, M.H., Dora: Weight-decomposed low-rank adaptation. ICML (2024)."}]

Author rebuttal:

Meta-review: This paper proposes a sparse coding-inspired reformulation of Transformer attention to improve in-context compositional learning, introducing explicit sparsity and a coefficient transfer mechanism. Reviewers appreciated the clear motivation, simplicity of the approach, and promising improvements across symbolic reasoning tasks, vision benchmarks, and even language models, highlighting its potential as a lightweight and interpretable inductive bias. Concerns were raised regarding limited theoretical grounding, initial evaluation scope, and connections to related work, but the rebuttal provided clarifications, additional experiments on diverse datasets and language tasks, and comparisons with relevant baselines, which strengthened confidence in the contribution. While further work on theoretical foundations, scaling, and broader applications would enhance the impact, the consensus is that the paper makes a meaningful and timely contribution. I therefore recommend acceptance.
Decision: Accept (poster)