Discrete Dictionary-based Decomposition Layer for Structured Representation Learning

Authors: taewon  park
hyun-chul  kim
minho  lee
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 14:53:08

Abstract: Neuro-symbolic neural networks have been extensively studied to integrate symbolic operations with neural networks, thereby improving systematic generalization. Specifically, Tensor Product Representation (TPR) framework enables neural networks to perform differentiable symbolic operations by encoding the symbolic structure of data within vector spaces. However, TPR-based neural networks often struggle to decompose unseen data into structured TPR representations, undermining their symbolic operations. To address this decomposition problem, we propose a Discrete Dictionary-based Decomposition (D3) layer designed to enhance the decomposition capabilities of TPR-based models. D3 employs discrete, learnable key-value dictionaries trained to capture symbolic features essential for decomposition operations. It leverages the prior knowledge acquired during training to generate structured TPR representations by mapping input data to pre-learned symbolic features within these dictionaries. D3 is a straightforward drop-in layer that can be seamlessly integrated into any TPR-based model without modifications. Our experimental results demonstrate that D3 significantly improves the systematic generalization of various TPR-based models while requiring fewer additional parameters. Notably, D3 outperforms baseline models on the synthetic task that demands the systematic decomposition of unseen combinatorial data.

Main paper: 
 introduction: 
 Compositional generalization, aiming at understanding unseen data by combining known concepts, is essential for neural networks to handle complex tasks [2, 13, 12, 16, 8, 6] . Tensor Product Representation (TPR) framework [33] facilitates this by embedding the symbolic structure of data within vector spaces, providing neural networks with compositional capabilities. Within this framework, individual objects are decomposed at the representation level into distinct symbolic components called role-filler pairs 2 . The TPR framework encodes each object by taking a tensor product of its role vector and filler vector, represented as T = f iller ⊗ role, and then superimposes them to represent multiple objects within a single representation. During decoding, the TPR framework retrieves specific fillers-essential for solving tasks-from the superimposed representation through matrix multiplication using an unbinding operator correlated to a particular role, f iller = T • unbind. This retrieved filler is then utilized in downstream tasks. Based on this property, TPR-based neural networks have demonstrated significant generalization and applicability in fields such as associative reasoning [28, 30] , mathematical problem-solving [29] , and natural language processing [9, 32, 21, 34] .Figure 1 : Overview of D3. D3 generates structured TPR representations by mapping input data to the nearest pre-learned symbolic features stored within discrete, learnable dictionaries. Each dictionary is linked explicitly to specific TPR components, such as roles, filler, and unbinding operators. Notably, D3 uses a shared dictionary configuration between the roles and unbinding operators. This figure illustrates, for example, that role 1 and unbind 1 share one dictionary, while role 2 and unbind 2 share another. T denotes a superimposed representation that represents multiple objects.Despite their successes, the TPR-based approaches pose a significant challenge known as a decomposition problem [33, 23] , which refers to the difficulty of decomposing input data into TPR components, such as roles, fillers, and unbinding operators. Without accurate decomposition, TPR-based models fail to represent the symbolic structure of data, causing a decline in the performance of the TPR operations. Recently, inspired by an object-centric learning method [18] , Park et al. [23] proposes an attention-based iterative decomposition (AID) module to address this issue. AID uses competitive attention to iteratively refine structured representations, thereby enhancing the systematic generalization of TPR-based models. However, it still struggles to generalize all possible combinations of known symbols in simple synthetic tasks. This failure is likely attributable to its insufficient mechanism for explicitly mapping input data to known symbolic features observed during training. Therefore, the decomposition module may need an additional mechanism to store observed symbolic features during training and utilize it to effectively decompose unseen combinatorial data of known symbols.In another line of work, discrete representation learning has been explored to improve the efficiency, interpretability, and generalization capabilities of neural networks [39, 14, 17, 37, 7] . This approach involves mapping continuous input data into discrete representations by finding the nearest features in a predefined codebook. The features within the codebook are learnable parameters, specifically trained to capture the latent features of data during training phase [39] . Some researchers have applied discrete representation techniques to extract specific types of representations from unstructured data [11, 43, 44] . Other researchers have integrated discrete symbolic embeddings within the TPR framework to improve its interpretability [21, 9] . However, these methods are designed for specific applications, such as question-answering and summarization tasks, making them difficult to integrate into other TPR-based models.In this work, we propose a Discrete Dictionary-based Decomposition (D3) layer for structured representation learning within the TPR framework. D3 employs the discrete representations techniques to utilize prior knowledge acquired during training for decomposition operations. Inspired by prior discrete key-value architectures [14, 38] , D3 consists of multiple dictionaries, each comprising discrete, learnable key-value pairs. Unlike prior work, each dictionary of D3 is linked explicitly to individual TPR components, such as role, filler, and unbinding operator. This design allows each dictionary to capture and store the discrete features of its corresponding TPR components during training. D3 acts as a drop-in layer that maps input data into pre-learned discrete features for the decomposition of TPR components through a three-step process, as illustrated in Fig. 1 . First, it generates multiple queries from the input data, with each query utilized for different TPR components. Next, it identifies the nearest codebook keys within each dictionary based on these queries. Finally, D3 generates structured TPR representations by aggregating the codebook values corresponding to these keys. Moreover, D3 can be seamlessly integrated into any TPR-based model by replacing the TPR component generation layer without requiring further modifications.Our main contributions are as follows.• We propose a novel D3 layer to tackle the decomposition problem inherent in the TPR-based approaches. D3 leverages discrete, learnable dictionaries to enhance the decomposition capabilities of TPR-based models. By mapping input data to pre-learned discrete features stored within the dictionaries, D3 effectively generates structured TPR representations.• We conduct extensive experiments across various systematic generalization tasks, including synthetic associative recall and text/visual question-answering tasks. Our experimental results show that D3 significantly enhances the generalization performance of TPR-based models, demonstrating its effectiveness on systematic generalization tasks.• Our analyses show that D3 generates well-bound structured representations that are satisfactory for the requirements of the TPR framework, utilizing the discrete, learnable dictionaries.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: This paper includes the paper's contributions and scope in the abstract and introduction, as follows. This paper tackles the decomposition problem inherent in the TPR-based approaches. To address this, this paper proposes a discrete dictionary-based decomposition (D3) layer designed to enhance the decomposition capabilities of the TPRbased models.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in 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. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
related work: 
 Decomposition Problem. Compositional generalization in neural networks, which allows for generalizing beyond training data, has been extensively studied [2, 13, 12, 16, 8, 6, 41] . One important capability for achieving this is a segregation, as discussed in Greff et al. [6] , which enables the formation of meaningful representations from structured and unstructured data [3, 18] . TPR-based neural networks also rely on this capability to generate structured representations for TPR components such as roles, fillers, and unbinding operators. In the TPR framework, these structured representations must satisfy specific conditions to ensure accurate encoding and decoding. First, roles need to be linearly independent to avoid filler overlap. Second, the unbinding operator must correlate with the corresponding roles to accurately retrieve associated fillers. Recent work [23] has shown that existing TPR-based models often fail to generate structured representations that meet these conditions, undermining their symbolic operations. To address this, an attention-based decomposition module [23] has been introduced, but it still shows limited performance on synthetic tasks involving the decomposition of unseen combinatorial data. In this work, we address the decomposition problem within the TPR framework using a discrete dictionary-based method, advancing the research further.Discrete Representation Learning. Discrete neural representation learning has introduced a codebook of discrete, learnable representations into neural networks [39] . During training, each discrete representation captures underlying latent features by mapping continuous input data to the nearest features within the codebook, which are then used for downstream tasks. Recent work on object-centric learning has utilized discrete representations to extract specific types of features from unstructured data, leveraging latent features learned during training [11, 43] . Some researchers have proposed a separate key-value codebook for learning discrete representations, demonstrating its effectiveness in systematic generalization [17] and robustness against distributional shifts [38] .Inspired by these findings, we develop a separate key-value-based discrete dictionary method to enhance the decomposition capabilities of TPR-based models. Other researchers have introduced a discrete symbolic embedding layer to improve the interpretability of TPR-based models, showing the feasibility of discrete representations in the TPR framework [21, 9] . However, their methods focus on encoding processes and specific tasks such as question-answering [21] and abstractive summarization [9] . In contrast, our work addresses the decomposition problem in TPR-based approaches, and our D3 method is a drop-in solution that can be easily adapted to any TPR-based model.Memory Network. Research on memory networks has focused on enhancing neural network capacity by integrating external memory [36, 4, 5, 24, 27, 41] . Memory-augmented neural networks store variable lengths of sequential data in this external memory and retrieve necessary information using various addressing methods [36, 5] . These writing and reading mechanisms share many similarities with our D3 approach. However, while memory networks store input features sequentially in their memory states as a continuous stream, D3 updates symbolic feature information through gradient descent into codebook parameters within dictionaries. This distinctive characteristic allows D3 to leverage the learned discrete features to decompose unseen data after training. In another work, Lample et al. [14] introduces a learnable key-value memory layer to improve the efficiency of the Transformer by replacing the feed-forward layer. Unlike their memory layer, D3 employs key-value pairs in dictionaries explicitly linked to individual TPR components, making it well-suited for the TPR framework.Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: This paper discusses the limitations of the work performed by the authors in Section 5, as follows. The model introduced in this paper requires additional computational overhead and configuration search when the proposed model is integrated into the existing baseline models.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 in 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 on 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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
method: 
 In this section, we explain how the D3 module generates structured representations of the TPR components using discrete, learnable dictionaries. We then introduce configurations of D3 and how it can be applied to our baseline models.D3 is a discrete dictionary-based drop-in layer designed to enhance the decomposition capabilities of TPR-based approaches. At every time step, D3 decomposes input data into TPR components, such as roles, fillers, and unbinding operators, by mapping input data to pre-learned symbolic features within dictionaries. These dictionaries consist of discrete, learnable codebook key-value pairs, denoted as {D j } N component j=1 as shown in Eq. 1. Each dictionary D j is explicitly linked to a j-th TPR component, allowing it to learn the symbolic features required for generating the specific TPR component. This design also enables the generation of structured representations for different TPR components individually and in parallel.D j := {(k j i , v j i ) | k j i ∈ R Dquery , v j i ∈ R Dcode } Ncode i=1where j = 1, ..., N component (1) where D j denotes the discrete, learnable dictionary for the j-th TPR component, k denotes a learnable codebook key, and v denotes a learnable codebook value. In the next paragraph, we describe how D3 generates TPR components using these dictionaries in three steps.Step 1: Query Generation. At each time step t, D3 takes input data, denoted as input t ∈ R Dinput , and generates the query, denoted as queries t ∈ R N component×Dinput , for each j-th TPR component using a query network, f j query : input t → query j t ∈ R Dquery . The query network can be any neural network; in this study, we use a feed-forward network with a single layer. Additionally, we apply a layer normalization [1] and a dropout of p dropout [35] to query j t .Step 2: Sparse Key Access. D3 searches for the nearest keys from each dictionary, D j , based on the generated query j t . We measure the similarity using the inner product between query j t and {k j i } Ncode i=1 . Then, D3 selects top-k codebook keys in order of largest similarity, as follows.EQUATIONwhere T k denotes the top-k operator that finds the indices of k largest values, and I j denotes the indices of the k most similar keys within D j . We found that applying L2 normalization to keys before the inner product mitigates the codebook collapse problem.Step 3: Aggregation of Code Values. D3 computes the normalized score for selected codebook keys, denoted as w j t , and aggregates codebook values corresponding to selected codebook keys with w j t , as follows.EQUATIONThen, D3 maps query j t to a dimension of D code and adds this projected vector to code j t . The summed vectors are mapped to a dimension of D component to generate structured representations of TPR components, as follows.EQUATIONwhere layer residual and layer final denote a feed-forward network with a single layer. Those components t are then utilized for TPR operations to solve the downstream tasks.In this section, we describe the configurations of D3 when applied to TPR-based models.Shared Dictionary between Role and Unbinding Operator. As discussed in Section 2, roles and unbinding operators should have correlated features for accurate TPR operations. Considering this characteristic of the TPR framework, we share the dictionaries of roles and unbinding operators. This shared dictionary also reduces the number of learnable parameters.D3 Applied to Filler. While the TPR framework requires specific conditions for roles and unbinding operators for accurate TPR operations, there are no such requirements for fillers. Therefore, we explore two configurations in this study: applying D3 to generate fillers (w/ F) and not applying D3 to generate fillers (w/o F). In the w/o F configuration, we follow the baseline models to generate the filler representations.In this section, we introduce our baseline models and explain how D3 is applied to them, considering the configurations of D3. We use three TPR-based models as our baselines: FWM [30] , TPR-RNN [28] , and Linear Transformer [10] . Notably, integrating D3 into these baseline models requires only substituting their TPR component generation layer with D3 without further modifications.Fast Weight Memory. Fast Weight Memory (FWM) [30] is a TPR-based memory network designed for understanding long sequential contexts. It proposes a single word-level TPR operation related to the perceptron learning rule [25] . It has shown significant associative reasoning capability in reinforcement learning and natural language processing tasks. FWM requires two types of roles (role 1 and role 2 ) and one filler for encoding, as well as two types of unbinding operators (unbind 1 and unbind 2 ) for decoding. When D3 is integrated into FWM, it employs three dictionaries for the shared dictionary configuration: one for the role 1 and unbind 1 , another for the role 2 and unbind 2 , and the other for filler, as shown in Fig. 1 .TPR-RNN. TPR-RNN [28] is a sentence-level TPR-based memory network designed for basic text question-answering tasks [42] . It incorporates various encoding operations such as writing, moving, and backlink to process sequential data at the sentence level. These operations necessitate different encoding components with varying dimensions, making direct connections to the decoding components challenging. As a result, we do not apply the shared dictionary configuration to TPR-RNN; instead, we use a shared query network without layer normalization. Furthermore, due to the differing dimensions of the TPR components in TPR-RNN, we employ distinct layer final layers for each TPR component.Linear Transformer. Linear Transformer [10] linearizes the attention mechanism to improve the computational efficiency of the Transformer [40] . Recently, Schlag et al. [31] demonstrated the equivalence between TPR and the linear attention mechanism, indicating that the key, value, and query in linear attention correspond to the role, filler, and unbinding operator, respectively. Building on this work, we apply D3 to generate the query, key, and value in the Linear Transformer. Unlike TPR-RNN and FWM, the Linear Transformer utilizes multi-head operations. Therefore, we use distinct dictionaries for each head, with the key and query of each head sharing the same dictionary.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [NA] Justification: This paper does not include theoretical results.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 crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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 or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
experiment: 
 In this section, we evaluate the effectiveness of D3 across various tasks, including a synthetic task, text/visual question-answering tasks, and a language modeling task. To assess the decomposition capabilities, we follow the experimental settings of the AID [23] , a prior work addressing the decomposition problem in the TPR framework, and closely compare our D3 model to baseline models and AID.Systematic Associative Recall (SAR) task. This task evaluates systematic generalization in memorizing and recalling combinatorial data [23] . It consists of a discovery phase and an inference phase. During the discovery phase, the model receives the combinatorial sequential items, each combining two symbols, x ∈ X and y ∈ Y where X= X 1 ∪ X 2 ∪ X 3 and Y = Y 1 ∪ Y 2 .The model is then required to predict an associated y when a specific x is presented. The SAR task uses different combination settings between training and evaluation to target systematic generalization specifically. During training, the model learns the following combination settings: (1) X 1 and Y 1 , (2) X 2 and Y 2 , and (3) X 3 and Y . At the evaluation, on the other hand, the model should generalize unseen combination settings, specifically X 1 and Y 2 . Additionally, the task includes a hyper-parameter p = |X3| |X2|+|X3| where |X i | denotes the cardinality of set X i . By adjusting p, this task tests the systematic generalization of models under varying levels of exposure to different symbol combinations during training. In our study, we focus solely on the most challenging setting of the SAR task (p = 0.0), where the subset X 3 is excluded. In the SAR task, the TPR framework regards x as the role and the unbinding operator, and y as the filler. Therefore, TPR-based models should systematically decompose the combinatorial data into structured representations by mapping x to the role and y to the filler during the discovery phase, and mapping x to the unbinding operator during the inference phase to solve this task.Systematic bAbI (sys-bAbI) task. This task is a variant of the bAbI task [42] designed to evaluate systematic generalization in text understanding and reasoning [23] . It consists of 20 distinct sub-tasks, each comprising stories, relevant queries, and corresponding answers. The sys-bAbI task requires the models to remember the stories and predict corresponding answers to the queries. Unlike the original bAbI task, the sys-bAbI task evaluates the models with two aspects: (a) in-distribution (w/o sys diff ) and (b) with the systematic difference (w/ sys diff ) where each sub-task includes unseen words during training. Therefore, the models should learn task-independent text understanding to solve the sys-bAbI task.Sort-of-CLEVR task. This task [26] evaluates compositional generalization in visual relational reasoning. It consists of scene images, queries, and corresponding answers. This task requires the models to understand the properties of individual objects (Unary) or the relationships between multiple objects (Binary or Ternary) within visual scene images, and predict the correct answers to the queries [20] . Therefore, the model should capture relationships within each object and between objects to solve this task.WikiText-103 task. This task [19] is a language modeling dataset consisting of lengthy corpora from Wikipedia. Although the WikiText-103 task does not directly measure the systematic generalization of the models, it is used to evaluate the effectiveness and applicability of D3 on a large-scale task beyond relatively simple tasks.In this section, we present the experimental results of the SAR task, sys-bAbI task, sort-of-CLEVR task, and WikiText-103 task. In our experiments, we set D query as D code /2. First, we evaluate FWM with D3 on the SAR task, which requires understanding the composition of two types of symbols, x and y. TPRbased models are expected to solve this task perfectly by mapping each symbol to a specific TPR component during decomposition. However, as shown in Fig. 2 , FWM and AID fail to generalize unseen combinations of known symbols. In contrast, our D3 module significantly outperforms other baseline models, achieving nearly 100% accuracy. This result demonstrates that D3 effectively decomposes unseen combinatorial data into TPR components using discrete dictionaries.Next, we test TPR-RNN and FWM with D3 on the sys-bAbI task. This task involves compositional information in each story sentence, such as the relation between objects and their locations. It makes a sentence-level model more suitable for capturing the structural information of data than a word-level model. However, as shown in Table 1 , TPR-RNN shows a larger performance gap between the w/o sys diff and w/ sys diff cases than FWM. Notably, D3 enhances the systematic generalization of both TPR-RNN and FWM with fewer additional parameters, significantly reducing the performance gap for TPR-RNN. These results highlight the efficacy of D3 in text understanding tasks.We also evaluate the Linear Transformer with D3 on the sort-of-CLEVR task and WikiText-103 task. Following the AID [23] , we use a 4-layered Linear Transformer with shared parameters for the sortof-CLEVR task and apply D3 to a 16-layered Linear Transformer at intervals of 4 out of the 16 layers for the WikiText-103 task. As shown in Tables 2 and 3 , D3 improves the performance of the Linear Transformer, with these improvements increasing as the capacity of the dictionaries grows. These results demonstrate the effectiveness of D3 on visual relational reasoning and language modeling tasks, as well as its applicability to the Linear Transformer. In addition, D3 shows comparable performance to the attention-based decomposition method, even with fewer parameters.In this section, we conduct a qualitative analysis of the structured TPR representations generated by D3 and an ablation study of D3. For these analyses, we experiment with D3 (w/o F) on the SAR task.TPR framework requires its structured representations to satisfy the following conditions for accurate TPR operations: (i) linearly independence between distinct roles, and (ii) high correlation between role and unbinding operator for the same symbol x. We analyze the orthogonality of generated representations to investigate whether they satisfy these TPR conditions. Specifically, we consider the case of varying x while keeping y fixed for simplicity. Fig. 3(c ) shows the cosine similarity between the roles during the discovery phase, and Fig. 4(c) shows the cosine similarity between the roles during the discovery phase and the unbinding operator during the inference phase. Both results demonstrate that the generated representations by D3 satisfy the TPR conditions, resulting in an accuracy of nearly 100%. We also conduct the same analysis for intermediate features, particularly query and code. Figs. 3 and 4 show that each intermediate representation complements the others to satisfy the TPR condition, indicating the effectiveness of D3. Furthermore, we analyze the similarity patterns of codebook keys and codebook values. Fig. 5 shows that the codebook features learn orthogonal patterns despite being learned without constraints. This result implies that the learnable parameters of dictionaries implicitly capture TPR conditions to ensure accurate TPR operations.We investigate the effect of hyper-parameters of D3, specifically N code , D code , and top-k, on performance on the SAR task. Fig. 6 (a) shows the effect of D code . We observe that the value of D code significantly affects the performance of D3. Notably, D3 fails to solve the SAR task when D code is set to 8, indicating a need for adequate capacity of D code . Fig. 6 (b) shows the effect of varying top-k while holding N code constant, indicating that D3 achieves optimal performance when top-k is set to 2. This result demonstrates the efficacy of the sparse mechanism employed by D3. Fig. 6(c ) examines the effect of varying N code while holding top-k constant, showing that D3 generally performs better with larger values of N code .Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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)?Answer: [Yes] Justification: This paper discloses all the information needed to reproduce the experimental results. This paper explains the mechanism of the proposed model and how it is applied to existing baseline models in Section 3 and presents the experiment details and hyperparameter settings in Appendices A and B. 
discussion and limitations: 
 Motivation. From the perspective of systematic generalization, the decomposition operations in the TPR framework can be viewed as mapping unseen data to TPR components observed during training. Motivated by this, we design a decomposition module based on discrete representations, which maps input data to discrete, learned features facilitating systematic generalization in the decomposition operations of TPR. This design choice differentiates our contribution from AID's competitive attention-based decomposition module. Additionally, each dictionary in D3 is explicitly linked to a specific TPR component, ensuring that each dictionary is responsible solely for generating its corresponding component. The generated components are then utilized in predefined TPR operations of the TPR-based models. This design ensures that each dictionary is trained to specialize in a specific TPR component.Interpretability. The TPR framework decomposes data at the representation level into distinct symbols, such as role-filler pairs for encoding and unbinding operators for decoding. This characteristic enhances the interpretability of models because the relationships between roles and unbinding operators explain which parts of the input the model focuses on to predict the output. However, this interpretability is reliable only when the generated structured representations satisfy the TPR conditions. In this context, D3 enhances the interpretability of models by providing structured representations that more effectively satisfy the TPR conditions than baseline models like FWM and AID. Figs. 9 and 10 demonstrate that the representations generated by D3 better conform to the TPR conditions than those from other baseline models, supporting our claim that D3 contributes to increased interpretability.D3 Applied to Filler (w/o F and w/ F). In the TPR framework, roles and unbinding operators must meet specific conditions, such as linear independence among roles and high correlation between roles and unbinding operators, to ensure accurate TPR operations. However, there are no such requirements for fillers, which are features related to downstream tasks. This characteristic affects the performance of D3 depending on whether it is applied to generate the fillers (w/ F) or not (w/o F). In our experiments, the w/ F configuration performs well on the sys-bAbI and sort-of-CLEVR tasks with relatively few labels (~200). In contrast, the w/o F configuration excels on the SAR and WikiText-103 tasks, which have a larger number of labels (500~). These findings suggest that the w/o F configuration may be more effective for large-scale practical tasks. Nevertheless, beyond these experimental results, we do not fully understand the conditions under which each configuration performs better. Consequently, one limitation of D3 is the additional burden of determining the suitable configuration for various tasks when applying it to other domains.Sparse Key Selection. D3 integrates seamlessly with existing TPR-based models, significantly enhancing their generalization performance across various tasks. However, this integration introduces additional computational overhead to the baseline models. Specifically, the sparse key selection mechanism of D3 has a computational complexity of O(N code × (D query + logk)) for each TPR component. Therefore, this complexity can become a drawback as the capacity of the dictionaries increases. One potential solution to address this capacity issue is to incorporate product keys into the sparse key selection mechanism of D3, a technique studied in prior discrete key-value architectures [14] . We leave this enhancement for future work.Scalability. The scalability of D3 is inherently linked to TPR operations of baseline models since the number of dictionaries in the D3 layer aligns with the number of TPR components required for their operations. As TPR operations require increasing components to handle large datasets, our method also requires a proportional increase in dictionaries, resulting in significant computational and memory overhead. As explored in prior work, one potential solution to mitigate this issue is distributing shared dictionaries across multiple heads or layers [14] . However, this approach requires further investigation and experimentation, which we plan to research in future work. 
conclusion: 
 In this paper, we tackle the decomposition problem inherent in the TPR framework, which poses a significant challenge for TPR-based models. To address this, we introduce a discrete dictionary-based layer, D3, designed to enhance the decomposition capabilities of TPR-based models. D3 employs the discrete dictionaries to map input data to pre-learned symbolic features within each dictionary, thereby generating structured TPR representations. Our comprehensive experiments demonstrate that D3 significantly enhances the systematic generalization of the TPR-based models with fewer additional parameters. Furthermore, our qualitative analysis verifies that D3 effectively generates structured representations that are satisfactory for the requirements of the TPR framework.We use the open-source implementation of the baseline models, TPR-RNN 6 [28] and FWM7 [30] . Following the experimental settings of baseline models, we use different configurations for each model. We train the TPR-RNN with D3 using an embedding size of 179 and the Adam optimizer with a batch size of 128 and a learning rate of 1e -3 , β 1 of 0.9, and β 2 of 0.99 for 100 training epochs. For FWM with D3, we use an embedding size of 256 and the Adam optimizer with a batch size of 64 and a learning rate of 1e -3 , β 1 of 0.9, and β 2 of 0.98 for training iterations of 60K. Furthermore, following the AID, we use the reconstruction loss for the bAbI task, introduced in Park et al. [22] , in our experiments on the sys-bAbI task. Each experiment took approximately 7 hours per seed for the TPR-RNN with D3 and 8 hours per seed for the FWM with D3. 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA]Justification: The authors do not foresee a negative societal impact on the work presented in this paper beyond the general effects of ML advancements.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or 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. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: This paper does not have a high risk for misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 to 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 a best faith effort.12. Licenses 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: This paper cites the original paper that produced the code package or dataset, and includes URLs in Appendix A.Guidelines:• The answer NA means that the paper does not use 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 a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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. • If this information is not available online, the authors are encouraged to reach out to the asset's creators.  
 
 Appendix: The sort-of-CLEVR task [26] evaluates compositional generalization in visual relational reasoning. It consists of scene images, queries, and corresponding answers. This task requires the models to understand the properties of individual objects (Unary) or the relationships between multiple objects (Binary or Ternary) within visual scene images and predict the correct answers to the queries. Therefore, the model should capture relationships within each object and between objects to solve this task.Each scene image, with a size of 75×75 pixels, includes 6 distinct objects in 6 different colors (red, blue, green, orange, yellow, or gray) and 2 different shapes (square or circle). This scene image is encoded by a visual encoder. The encoded visual feature is then concatenated with the embedding vector of the query. These concatenated features are provided to the model. Following the experimental settings of the AID [23] , we use a single CNN layer with a kernel size of 15 and a stride of 15 for the visual encoder, and an embedding size of 128 for the word embedding method. Also, we use a 4-layered Transformer, where each layer shares its parameters with others, as our baseline model.To build the experimental environment for the sort of CLEVR task, we utilize the open-source implementation 8 from Mittal et al. [20] . We train the model using the Adam optimizer with a batch size of 64 and a learning rate of 1e -4 for 100 training epochs. Each experiment took approximately 2.5 hours per each seed.The WikiText-103 task [19] is a language modeling dataset consisting of lengthy corpora from Wikipedia. Although the WikiText-103 task does not directly measure the systematic generalization of the models, it is used to evaluate the effectiveness and applicability of D3 on a large-scale task beyond relatively simple tasks.The WikiText-103 task comprises 28,475 articles for training, 60 for validation, and 60 for testing. Following the experimental settings of Schlag et al. [31] , we partition the articles into segments of L words. During training, the gradient is back-propagated only within spans of L words. The performance of the model is evaluated using the measure of perplexity. During evaluation, the model processes an input sequence of L words by sliding a segment over the article with a stride size of 1. Perplexity is then computed based on the last position of each segment, except for the first segment, where every position is taken into account.To build the experimental environment for the WikiText-103 task, we utilize the open-source implementation 9 from [31] . Following the AID [23] , we apply D3 to a 16-layered Linear Transformer at intervals of 4 out of the 16 layers. We train the model using the Adam optimizer with a batch size of 96, an initial learning rate of 2.5e -4 , and a learning rate warmup step of 2,000 for 120 epochs. Each experiment took approximately ~3 days. indicates that although D3 improves the performance of the Linear Transformer in language modeling tasks, the choice of updating rules has a more substantial impact on performance for tasks involving the comprehension of lengthy corpora than the decomposition operation. In this section, we extend our ablation studies to investigate the effects of varying the number of keys in the codebook and the impact of removing either the residual connection or the codebook from the D3 layer.The Effect of Varying the Number of Codebook Keys. Fig. 7 shows that even with a significantly reduced number of keys, the model with D3 maintains high accuracy on the SAR task. This observation prompts the question of how consistent performance is achieved despite the reduction in codebook size. To explore this further, we examine the impact of removing the codebook or the residual connection within the D3 layer on the SAR and sys-bAbI tasks. Specifically, removing the codebook means that the components are generated solely by the shared feed-forward networks (layer residual and layer final ) while removing the residual connection implies that the components are derived solely from the codebook values. The ablation studies in Tables 12 and 13 demonstrate that incorporating the codebook mechanism leads to nearly 100% accuracy on the SAR task and significantly improves the systematic generalization of models on the sys-bAbI task. However, as shown in the ablation study on the residual connection (Fig. 8 ), the codebook alone does not achieve the same level of generalization and exhibits instability within the D3 layer.In conclusion, our experimental results indicate that the combination of the codebook and the shared residual networks within the D3 layer is crucial for enhancing systematic generalization performance and stability. By integrating these two components, our D3 layer significantly improves the systematic generalization capabilities of TPR-based models. • The answer NA means that the paper does not include experiments.• If the paper includes experiments, a 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 a 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 a 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 it 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 to 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 a 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 submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to 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., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 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 circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).The code of D3 is publicly available at https://github.com/taewonpark/D3The roles and fillers depend on the task at hand. For example, in a tree structure, the role corresponds to a position within the tree, while the filler represents the label associated with that position[34]. In associative memory, the role is analogous to an associative key, and the filler corresponds to the associative value[28,30]. 38th Conference on Neural Information Processing Systems (NeurIPS 2024).We used the following seed values: {0, 1111, 2222, 3333,4444, 5555, 6666, 7777, 8888, 9999} 4 https://github.com/taewonpark/AID/tree/main/SARtask 5 https://github.com/taewonpark/AID/tree/main/bAbItaskhttps://github.com/APodolskiy/TPR-RNN-Torchhttps://github.com/ischlag/Fast-Weight-Memory-publichttps://github.com/sarthmit/Compositional-Attention/tree/main/Sort-of-CLEVRhttps://github.com/IDSIA/lmtool-fwp 


Decision: The paper proposes D3, a layer of discrete, learnable dictionaries for the Tensor Product Representation (TPR) framework. With each dictionary connected to a specific TPR component (e.g., roles, fillers, unbinding operators), D3 enables an effective decomposition of input data into TPR components. Experiments with 3 different TPR architectures on 4 different tasks demonstrate the benefits of the proposed layer compared to baselines without it and alternative solutions to the decomposition problem in TPR architectures.

All 4 reviewers feel cautiously positive about the paper, commending its novel approach to the decomposition problem in TPR architectures, thorough evaluation, and insightful analyses. The main concerns expressed by the reviewers were related to (1) architectural / performance differences relative to prior work — primarily AID — with similar goals; (2) missing ablation studies / comparisons to state-of-the-art baselines; (3) the relationship between the dictionaries and the TPR components, in particular in the context of interpretability and enforcement of TPR conditions; and (4) scalability. Accessibility was noted as a potential issue for readers without familiarity with the TPR framework.

The authors responded well to these concerns. They convincingly positioned D3 relative to AID (both in terms of technical novelty and practical performance), provided the requested ablation studies and comparisons to state-of-the-art baselines, and discussed D3 design choices and their effects on TPR conditions / interpretability. It is expected that these discussions and results, as well as a gentle introduction to the TPR framework, will be integrated into the paper. Scalability was identified as an issue that requires further attention.

Given that the major concerns have been addressed, as well as the relevance of the problem and novelty of the proposed approach, the paper is recommended for acceptance as a poster.

Reviews: review: paper_summary: paper_summary: This paper address the difficulty of decomposing input data into Tensor Product Representation (TPR) components, namely, roles, fillers, and unbinding operators. The proposal called D3 includes the use of learnable dictionaries for these components, and the mapping of input data into intermediate features (code, query, component), in order to generate TPR components from input data. The authors demonstrate that D3 can be easily integrated into existing TPR-based models and improves their systematic generalization performance across various tasks, including synthetic associative recall, text/visual question-answering, and language modeling.
strength_weakness: strength_weakness: * The paper is an alternative to the recently introduced AID method (Ref #22), and the goal/settings are nearly identical. It is helpful to clearly differentiate the contributions of this work to AID.
* There are design choices in the D3: dictionaries, query, code, components and the ways they are connected. There should be clear motivations for each of these choice.
* While the paper provides some analysis of the learned representations, a deeper theoretical analysis of why D3 works well could enhance the contribution.
* The experiments are primarily focused on synthetic or relatively simple tasks. Testing on more complex, real-world tasks would strengthen the paper's claims about generalization. The comparison to baselines could be expanded to include a wider range of state-of-the-art methods in compositional generalization.
questions: questions: The questions naturally arise from the weaknesses above.
limitations: limitations: The authors discuss some limitations of their approach, including the need for task-specific configuration and added computational overhead. They could expand on potential limitations in terms of scalability to very large models/datasets.
paper_summary: paper_summary: Drawing inspiration from discrete representation learning with dictionaries, the author introduced a novel Tensor Product Representation (TPR) framework, a Discrete Dictionary-based Decomposition (D3) layer designed to retain the learned symbolic features during training and apply them effectively to address decomposition challenges in previously unseen data. Owing to its architectural properties, this method is defined as a drop-in layer that maps input to pre-learned symbolic features, facilitating smooth integration with existing TPR methods. Comprehensive experimental results of this novel approach showcase superior or comparable performance to that of other baseline methods.
strength_weakness: strength_weakness: - W1: The proposed method's technical novelty is a slight improvement over the previous work, AID. It appears that the proposed method's performance in more complex cases, such as Sort-of-CLEVR and WiKiText, is due to the introduction of additional configuration complexity and more learnable parameters. This may cause scalability issues in cases where more complex TRP operations are required since the number of dictionaries a model directly corresponds to the number of roles/unbind operators. 
- W2: Although the proposed method is a drop-in layer applicable to other existing TPR methods, it raises the question of how such an attention key-query-value mechanism can meet the TPR operation conditions. In particular, how can you ensure that the shared dictionary satisfies these properties during training to satisfy near orthogonality between roles or between unbinding operators?
questions: questions: Please check out the Weakness section first. I listed additional questions as follows:
- Q1: In Fig 2, the experimental results of the SAR task appear to show a huge difference in performance between the proposed method and AID. In the original AID paper, it achieved about 90% accuracy for the SAR task. The original experiment setup that AID ran on appears to be more difficult than yours, so why did AID perform significantly worse on your setup?
- Q2 (related to W2):  In Fig. 5(a), there are some strong similarities among specific keys, which indicates some redundancy of the learned keys. How can these redundant properties of keys not violate the TPR operation conditions?
limitations: limitations: Please check out the Weakness section.
paper_summary: paper_summary: The paper presents a discrete dictionary-based decomposition (D3) layer for tensor product representation (TPR). The purpose is to enhance the decomposition capabilities of TPR so that it can perform downstream tasks more effectively. D3 uses the discrete, trainable key-value dictionaries to map input data to “discrete” features within each dictionary. The dictionaries correspond to the TPR components, which are roles, fillers, and unbinding operators. More specifically, one dictionary is trained for the roles and unbinding operators and an optional one for the fillers. Experiments were conducted by using D3 in three TPR-based models. The paper shows that the use of D3 in these models outperforms the ones without using it.
strength_weakness: strength_weakness: -	The paper assumes that the readers are familiar with terminologies used in the paper. It would be better if the paper provided definitions of some key concepts frequently used in the paper (such as roles, fillers, etc.). I had to look into references to fully understand the concepts. 
-	D3 is evaluated with multiple TPR-based models on multiple tasks (such as text understanding and reasoning, visual relational reasoning and large language modeling). But it is unclear how these TPR-with-D3-based models compares with the SOTA models for these individual tasks. Without such a comparison, the significance of the work is unclear. 
-	The paper mentioned at the beginning that using discrete representation can help with the interpretability of the models. However, there is no discussion or evaluation of this aspect of the model as the result of using D3. 
-	I am not sure if the resulting “discrete” representation of the input data can be called “symbolic” features (which the paper uses a few times). The representations are still in the vector form, which is numeric. They may be considered “discrete”, but they are not symbols.
-	It is not clear how the dictionaries trained with D3 correspond to three components (roles, fillers, and unbinding operators) in TPR. Each component seems to be trained in the same way. There do not seem to be distinctions among them according to the descriptions in Section 3.1. In the experiments, only two or one dictionary is used, with one dictionary corresponding to roles and unbinding operators and an optional one for the fillers. If only the number of dictionaries matters, then the “w/o F” option could be considered “w/o roles/unbinding operators but with fillers”. More explanation on this would be beneficial.
-	There are some small errors in the presentation. For example, Subfigues (b) and (c) in Figure 6 should be switched to match their references in the text.
questions: questions: -	How is the TPR-with-D3-based models compared with the SOTA method on their individual tasks?
-	Does the use of the D3 layer increase the interpretability of the models? 
-	How do the dictionaries trained with D3 correspond to three components (roles, fillers, and unbinding operators) in TPR?
limitations: limitations: Yes.
paper_summary: paper_summary: The authors propose the Discrete Dictionary-based Decomposition method, a discrete representation learning for TPRs. It encodes the roles and unbinding queries (and potentially also the fillers) using a learned dictionary, which serves as a codebook. The roles and then unbindings share the codebook, encouraging correct, noise-free retrieval. The layer is a drop-in replacement for any TPR-based layer. The authors demonstrate strong performance on tasks requiring systematic generalization, matching, or outperforming AID without using an iterative attention mechanism.
strength_weakness: strength_weakness: Given that there is a residual connection around the codebook, it is unclear what component of the roles and unbindings are generated by the codebook and which part of them comes from the residual. It would be nice to see ablations on the model without residual. If the model is not able to learn without it, an alternative is to replace the output of the residual by their mean over the dataset and see the effect on the performance. 

Additionally, I would like to see analysis like Figures 3 and 4, but for the baselines without D3. Currently, in 282, the authors claim "indicating the effectiveness of D3", but it is unclear how orthogonal representations are the baselines' representations. 

I believe that this additional analysis could deepen the understanding of why and how D3 works and significantly strengthen the paper.
questions: questions: In line 115, the authors write: "This distinctive characteristic allows D3 to leverage the learned symbolic features to decompose unseen data after training.". Could the authors provide an additional explanation for why discretization aids the decomposition of unseen data, as opposed to denoising the representations?

On Fig 5 (b), the authors show an ablation over the number of keys in the codebook. Even with the smallest key size, the model works well. I would like to see a similar plot going down in the number of keys as long as the performance doesn't suffer significantly or the number of keys doesn't reach 1. This would help to understand better the effect of the codebook.

Calling WikiText 103 a large-scale language modeling dataset in 2024 is a stretch. Please just call it "language modeling" dataset.
limitations: limitations: The authors discuss the limitations of their method in the paper. As an additional limitation, as with all TPR papers, it is unclear how the improved generalization on toy tasks would transfer to real-world scenarios, like improving the generalization of LLMs. However, this should not be a reason to downvalue a contribution of the paper, given that it aims to improve an important shortcoming with a novel idea.

Comments: title: Response to the reviewers' rebuttal
Concerns about the usefulness codebook
comment: Thank you for detailed response!
We thank Reviewer 3Nfz for your time and your constructive comments during the rebuttal period.
Thank you for the clarification. I've adjusted my rating accordingly.
We sincerely thank the reviewer for thoughtful feedback and for acknowledging our efforts to improve the paper.

The primary distinction in our approach lies in the application of shared residual networks. In the FWM, *queries* are treated as structured representations. In contrast, in the D3 layer, we introduced shared residual networks that are applied across all *queries*. This means that, regardless of the specific type of TPR components, the same residual network is employed to map *queries* into the dimensional space of *codes*. While FWM directly utilizes the *queries* as structured representations, our shared residual approach further transforms these *queries* through a shared feed-forward network before performing with subsequent TPR operations.

We also appreciate the reviewer’s insightful suggestion. The suggested idea can potentially enhance the D3 layer by increasing the dependency on the codebook. We plan to explore this direction further in our future work.

Thank you once again for your valuable input and your continued interest in our research.
I'm thankful for the author's effort to improve the paper. That makes the paper more correct, but I'm afraid I'm not convinced enough about the method to increase my score. 

Could the authors summarize what is the main difference between the shared residual (without a codebook) and the previous work?

A suggestion if the authors want to keep working on improving the paper (not for the rebuttal, but for general curiosity): it might help to improve the codebook utilization and possibly the generalization properties if the authors would introduce a low-rank bottleneck on the residual, giving the network more incentive to rely on the codebooks.
We thank Reviewer QJJM for your time and your constructive comments during the rebuttal period.
We sincerely appreciate Reviewer FjyM’s thorough review and valuable suggestions.

In response to your feedback, **we will update our manuscript to include the additional experimental results you recommended and address your concerns** as follows:

- Add ablation study for the effect of the residual connection (Figure 13)
- Add additional orthogonal analysis on baseline models (Figures 14 and 15)
- Add ablation study for the effect of varying the number of keys (Figure 16)
- Add ablation study for the effect of the codebook (Tables A1 and A2)

Furthermore, we will expand our discussion in the manuscript to emphasize the effect of both codebook and shared residual networks, which are the primary components of our D3 layer.

The ablation study for the codebook on the SAR task (Table A1) indicates that the shared residual networks significantly enhance the generalization performance of the D3 layer. However, the results from the sys-bAbI task (Table A2) suggest that, while these networks improve performance, they alone struggle to generalize more structured data.

The ablation studies for the codebook (Tables A1 and A2) demonstrate that incorporating the codebook mechanism results in nearly 100% accuracy on the SAR task and significantly improves the systematic generalization of models on the sys-bAbI task. However, as shown by the ablation study on the residual connection (Fig. 13), the codebook alone does not achieve the same level of generalization and exhibits instability in the D3 layer.

In conclusion, our experimental results indicate that **the combination of the codebook and shared residual networks within the D3 layer is crucial for enhancing systematic generalization performance and stability**. By integrating these two components, **our D3 layer significantly enhances the systematic generalization capabilities of TPR-based models**.
Thank you for the authors' clarification and for addressing my concerns. 
I have decided to increase the score.
I appreciate the effort of the authors to clarify the effects of the codebook.

The results show that the codebook is used, but the shared residual network may have a more significant impact. I think more analysis should be needed to determine the effect of each of these components. Additionally, more emphasis should be placed on the shared residual network, which is not discussed in the current version of the paper.

Given this, I am maintaining my score.
Thank you for your detailed feedback and valuable suggestions.

To address your concerns, **we conducted additional ablation studies on the SAR and sys-bAbI tasks**. Specifically, we investigated **the D3 layer without incorporating the codebook** (referred to as "*w/o codebook*"), where *components* are generated solely using the shared feed-forward networks (layer$\_\text{residual}$ and layer$\_\text{final}$). In this configuration, the *component* is calculated as "layer$\_\text{residual}$(*query*)" instead of the "*component* = *code* + layer$\_\text{residual}$(*query*)", which is described in Eq. 4.

As shown in Table A1, even without the codebook, the D3 layer improves the generalization performance of the baseline model on the SAR task. This result indicates that **the shared feed-forward networks significantly contribute to performance enhancement on the SAR task**, which may explain why the model maintains robust performance even with fewer keys.

However, it is important to note that without the codebook, the D3 layer does not achieve near-perfect accuracy on the SAR task (as shown in Table A1) and fails to significantly enhance the systematic generalization of the baseline on the sys-bAbI task (as shown in Table A2). **These results demonstrate that the codebook plays a crucial role in enhancing the model's overall performance and generalization capabilities**, especially in tasks requiring systematic generalization.

Additionally, we experimented with $N_\text{code} = 1$ on the SAR task, where the codebook may act as a bias term. The results in Table A1 show that using a single codebook element results in degraded generalization performance compared to the "*w/o codebook*" configuration, indicating that multiple codebook elements are essential for achieving optimal results.

| **Table A1.** SAR task | $D_\text{code}$ | $N_\text{code}$ | top-$k$ | Accuracy |
| --- | :---: | :---: | :---: | :---: |
| FWM | - | - | - | 44.9$_{\pm31.5}$ |
| --- |
| D3 | 32 | 4 | 2 | 87.38$_{\pm11.10}$ |
| D3 | 32 | 64 | 8 | 99.27$_{\pm0.88}$ |
| D3 *w/o codebook* | 32 | - | - | 89.02$_{\pm4.56}$ |
| --- |
| D3 | 64 | 1 | 1 | 89.10$_{\pm7.99}$ |
| D3 | 64 | 4 | 2 | 94.47$_{\pm2.35}$ |
| D3 | 64 | 64 | 8 | 94.29$_{\pm8.06}$ |
| D3 *w/o codebook* | 64 | - | - | 91.65$_{\pm3.66}$ |

| **Table A2.** sys-bAbI task | *w/o sys diff* | *w/ sys diff* |
| --- | :---: | :---: |
| FWM | 0.79$_{\pm0.14}$ | 2.85$_{\pm1.61}$ |
| --- |
| D3 | 0.75$_{\pm0.17}$ | 1.96$_{\pm0.88}$ |
| D3 *w/o codebook* | 1.19$_{\pm0.41}$ | 3.55$_{\pm1.04}$ |

In summary, **combining of the codebook and the shared residual networks is crucial for achieving systematic generalization in TPR-based models**. We hope these additional experiments provide clarity regarding our approach.
We are thankful to the authors for their efforts in the rebuttal, and they answered most of my questions positively.

However, I find the results of "Q2: Ablation study for the effect of varying the number of keys *extremely concerning*. The performance of the network doesn't change with even as low as 2 keys in the codebook. This begs the question of whether the codebook, which is the main component of the model, is doing anything, or are there other differences in the training pipeline or the model that cause the observed difference in Fig 2 compared to the baselines? This makes me question the validity of the approach.

Do the authors have an explanation of how this is possible? Can you please run additional ablations with (1) the codebook removed, but otherwise use the same code and pipeline, (2) the number of codebook elements set to 1, which is equivalent to having a linear layer down-projecting to 1 elements and then projecting it back.

I would appreciate any evidence that shows that the codebook is crucial, like the one I described above, or one that show that codebooks with very small number of entries do not work on other tasks.