In-Context Compositional Learning vis Sparse Coding Transformer

Thank you for the constructive comments. We will incorporate this feedback into our revised work.

Q1. The softmax attention hinders compositional generalization is an unsubstantiated argument.

A: 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:

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

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

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

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

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.

A: (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.

(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. Therefore, 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.

Q3. Application to language modeling tasks.

A: 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.

We 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].

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

Experiment Details: In 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.

[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).

[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).

[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).

[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).

Thank you for the constructive comments. We will incorporate this feedback into our revised work.

Q1. Provide theoretical justification for the generalization of the method.

A: 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.

Q2. Application to more complex data.

A: We validated our method's applicability by evaluating its compositional generalization on the Im-promptu benchmark and its performance in language modeling applications.

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.

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.

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

[1] Dedhia, Bhishma, Michael Chang, Jake Snell, Tom Griffiths, and Niraj Jha. "Im-promptu: in-context composition from image prompts." NeurIPS (2023).

[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).

[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).

[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).

[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).

Q3. Discussion of object-centric work.

A: Both the "dictionary atoms" in our framework and the "slots" in OCL literature serve as learned primitives for composing complex inputs.

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

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

We will incorporate a discussion of OCL in the related work section of a revised version.

Thank you for the insightful feedback. We will incorporate this feedback into our revised work.

Q1. Clarification of the toy example.

A: 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:

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

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

Q2. Compare with "Attention as a Hypernetwork (HYLA)".

A: While both approaches reinterpret the attention mechanism, our proposal differs in two fundamental ways that stem from our core inspiration: sparse coding.

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

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

Q3. Application to language model.

A: 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.

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

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.

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

Specifically, 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.

[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).

[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).

[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).

[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).

Q4. Effect of atom number and hidden dimensionality on performance.

A: 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.

Q5. Why is prox(x) used instead of training with L1 regularization on the coefficient values?

A: 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.

Thank you for the constructive comments. Below are our responses to each of your comments. We will incorporate this feedback into our revised work.

Q1. The standard RAVEN dataset is inadequate for testing novel rule compositions.

A: Different from the standard RAVEN dataset, we used two distinct datasets in our paper to evaluate different aspects of compositional learning:

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

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

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

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.

A: 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.

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

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.

This 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: Simply 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. Our 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.

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

[1] Benny, Y., Pekar, N. &Wolf, L. Scale-localized abstract reasoning. CVPR (2021).

Q4. Evaluate the models on Im-promptu benchmark.

A: 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.

Q5. Will cross-entropy loss help blurry images of the baseline model?

A: 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.