Title: In-Context Compositional Learning vis Sparse Coding Transformer
Venue: NeurIPS 2025
Accepted: yes
Keywords: sparse coding

compositional learning

representational 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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nThe 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.\n\nQUESTIONS:\nThe key concerns are listed above.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-18"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n### 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*\n\nQUESTIONS:\nPlease see the Weaknesses\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 1, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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.\n\nQUESTIONS:\nWhy 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\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.\n\nQUESTIONS:\nIn 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-01"}