Discrete Dictionary-based Decomposition Layer for Structured Representation Learning
Authors: Taewon Park
Hyun-Chul Kim
Minho Lee
Venue: NeurIPS 2024 poster (2024)
Keywords: tensor product representation
neuro-symbolic neural network
systematic generalization
compositional generalization
decomposition problem
structured representation learning
discrete representation learning
Primary area: deep_learning_architectures

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.

Discrete Dictionary-based Decomposition Layer for Structured Representation Learning
Taewon Park Hyun-Chul Kim
Minho Lee
1 Kyungpook
National University, South Korea
2 ALI Co., Ltd., South Korea
pt#79980gmail.cm com, Syunchul__kimkmuu kr, mbolee@ggail..commmmmm com
Abstract
Neuro-symbolic neural networks have been extensively studied integrate sym- bolic operations with neural networks, thereby improving systematic generalization. Specifically, Tensor Product Representation (TPR) framework enables neural net- works 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, undermin- ing their symbolic operations. To address this decomposition problem, we propose a Discrete Dictionary-basee 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 decom- position operations. It leverages the prior knowledge acquired during training to generate structured TPR representations by mapping input data to pre-learned discrete features within these dictionaries. that can be seamlessly integrated into any TPR-based model without modifications. Our experimental results demonstrate that D3 significantly improves the system- atic 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.
a straightforward drop-in layer
Introduction
Compositional generalization, aiming at understanding unseen data by combining known concepts, is essential for neural networks handle complex tasks [2, 13, 12, 16, 8, 6J. Tensor Product Representa- tion (TPR) framework [33] facilitates this by embedding the symbolic structure 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?. The TPR framework encodes each object by taking tensor product of its role vector and filler vector, represented as 7 = filler t role, and then superimposes them to represent multiple objects within single representation. During decoding, the TPR framework retrieves specific fillers- eesseralsseettall for solving tasks-from from the superimposed representation through matrix multiplica- tion using an unbinding operator correlated a particular role, filler 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).
'The code of D3 is publicly available at https://github.. http://gghuub.commewwopararaa com/tacwonpark/D3
"The roles and fillers depend on the task at hand. For example, in 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 associative key, and the filler corresponds to the associative value [28, 30).
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Step Quary Generation
Step Sparse Key Access
Step Aggregating Code Values
Key
Volue
Query
TPR Encoding Operations
filler @ rale & rolez - T
Query I [iller]
Code (flleer
Components ( [iller
Input (Bicodlegg)
Discc wte Vay-ealue Dictionary
[ ([lller)
Query
Code
Components (roic) )
TPR Decoding Operations T Unbindy Unbind = filler
Key
Value
Quary S
Input (Derading
: Encoding Path Decoding Path
Discrete Key-Value Dictionary rrols / untinding Operator) )
Code (unlituing
Components unninniing
(anbisding
Figure I: 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 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, and unbind, share one dictionary, while rolez and unbind, share another. T denotes a superimposed representation that represents multiple objects.
Despite their successes, the TPR- approaches pose a significant challenge known as a decompo- sition problem [33, 23], which refers to the difficulty of decomposing imput 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 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 address this issue. AID uses competitive attention to iteratively refine structured representations, thereby enhancing the systematic general- ization 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 improve the efficiency, interpretability, and generalization capabilities of neural networks [39, 14, 17, 37, 7J. This approach involves mapping continuous input data into discrete representations by finding the nearest features in 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 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, Iearnable key-value pairs. Unlike prior work, each dictionary 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 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, 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 novel D3 layer 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 syn- thetic associative recall and text/visual werrstioonanswerinn tasks. Our experimental results show that D3 significantly enhances the generalization performance of TPR-based models, demonstrat- ing 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.
2 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 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. 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, 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 codebook of discrete, learnable representations into neural networks [39]. During training, each discrete representation captures underlying latent features by mapping continuous input data 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 separate key-value-bare enhance the decomposition capabilities of TPR-based models. Other researchers have introduced discrete symbolic embedding layer to improve the interpretability of TPR-base 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 cersstion-answrinng [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 casily adapted to any TPR-base model.
discrete dictionary method to
Memory Network. Research memory networks has focused on enhancing neural network capacity by integrating external memory [36, 4, 5, 24, 27, 41J. 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 continuous stream, D3 updates symbolic feature information through
in their memory states 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 individual TPR components, making well-suited for the TPR framework.
3

---

3 Method
In this section, explain how the 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.
3.1 Discrete Dictionary-based Decomposition module
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-Icarned symbolic features within dictionaries. These dictionaries consist of discrete, learnable codebook key-value pairs, denoted as (D) } ccomponent as shown in Eq. Each dictionary DD 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': [kk, v/) K/ E RR v E I.. (ePaa Made li=l
where
1, componeut
(1)
where Di denotes the discrete, learnable dictionary for the j-th TPR component, k denotes learnable codebook key, and v denotes learnable codebook value. the next paragraph, we describe how D3 generates TPR components using these dictionaries in three steps.
Step 1: Query Generation. At each time step f, D3 takes input data, denoted as input, E TR Dinpur and generates the query, denoted as queries, E N Neompommntt Duper for each j-chj-th TPR component using a query network, fileery input, + query! E IR Pquery The query network can be any neural network; in this study, we use feed-forward network with single layer. Additionally, we apply layer normalization [1] and a dropout of Paropout [35] to query?.
Step 2: Sparse Key Access. D3 searches for the nearest keys from each dictionary, D), based on the generated query/- We measure the similarity using the inner product between query/ and {k) j Then, D3 selects top-k codebook keys order of largest similarity, as follows.
I' = To (query) k)) where Ri = k./ /</
(2)
where Tk denotes the top-k operator that finds the indices of & largest values, and I' denotes the indices of the & most similar keys within D) 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/, and aggregates codebook values corresponding to selected codebook keys with w/, as follows.
code Exerw where wa Softmax(query! T R
Then, D3 maps query? to dimension of Deode and adds this projected vector to code/. The summed vectors are mapped to a dimension component to generate structured representations of TPR components, as follows.
component code + layer residmml (query!) E Pack
(4)
component, = layer (component!) E R Dosmproee
(5)
where layer and layer denote feed-forward network with single layer. components, are then utilized for TPR operations to solve the downstream tasks.
Those
3.2 Module Configurations
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 filerss. 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 configuration, we follow the baseline models to generate the filler representations.
3.3 Integration of D3 into Existing TPR-based Models
In this section, we introduce our baseline models and explain how D3 applied to them, considering the configurations of D3. We use three TPR-based models 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 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) and rolez) and one filler for encoding, as well as two types of unbinding operators (unbind) and unbindz) for decoding. When D3 is integrated into FWM, it employs three dictionaries for the shared dictionary configuration: one for the role; and umbind1, another for the rolez and unbindz, and the other for filler, shown in Fig.
TPR-RNN. TPR-RNN [28] is sentence-level TPR- -Raeddedd memory network designed for basic text question- question-answerree 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 compo- nents challenging. As result, we do not apply the shared dictionary configuration to TPR-RNNI 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 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, apply D3 to generate the query, key, and value in the Linear Transformer. Unlike 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.
4 Experiment
In this section, we evaluate the effectiveness of D3 across various tasks, including synthetic task, text/visual question-answeringg tasks, and a language modeling task. To assess the decomposition capabilities, we follow the experimental settings of the AID [23], prior work addressing the decomposition problem in the TPR framework, and closely compare our D3 model to baseline models and AID.
4.1 Task
Systematic Associative Recall (SAR) task. This task evaluates systematic generalization mem- orizing and recalling combinatorial data [23]. It consists of a discovery phase and inference phase. During the discovery phase, the model receives the combinatorial sequential items, each combining two symbols, z E X and y E Y where X Xi U X X3 and Y Yi U Yg. The
5

---

model is then required to predict an associated y when specific z' is presented. The SAR task uses different combination settings between training and evaluation target systematic general- ization specifically. During training, the model learns the following combination settings: (1) X1 and V1, (2) X2 and Y2, and (3) X3 and Y. At the evaluation, on the other hand, the model should generalize unseen combination settings, specifically X1 and Yg. Additionally, the task includes hyper-parameter p [ where [X,] denotes the cardinality of set X,, 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 X3 is excluded. In the SAR task, the TPR framework regards the role and the unbinding operator, and the filler. Therefore, TPR-ban -PR-based models should systematically decompose the combinatorial data into structured representations by mapping z to the role and y the filler during the discovery phase, and mapping z to the unbinding operator during the inference phase to solve this task.
[Xal
Systematic bAbl (sys-bAbl) task. This task is variant of the bAbl 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-bAbl task requires the models to remember the stories and predict corresponding answers to the queries. Unlike the original bAbl task, the sys-bAbl 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 aask-independent text understanding to solve the sys-bAbl 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 Terarry) 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-102 task. This task [19] is a language modeling dataset consisting of lengthy corpora from Wikipedia. Although the WikiTexxt-103 task does not directly measure the systematic general- ization of the models, it is used to evaluate the effectiveness and applicability of D3 on a large-scale task beyond relatively simple tasks.
4.2 Experimental Results
In this section, we present the experimental results of the SAR task. sys-bAbl task, oort-of-CLEVR task, and WikiText- 103 task. In our experiments, we set Dquery as Doode/2. 2.
4.2.1 -aR-based Memory Networks
First, we evaluate FWM with D3 on the SAR task, which requires understanding the compo- sition of two types of symbols, I' and y. TPR- based models are expected to solve this task
100
80
Accuracy
Model
perfectly by mapping each symbol to a specific TPR component during decomposition. How- ever, as shown in Fig. 2, FWM and AID fail to generalize unseen combinations of known sym- bols. In contrast, our D3 module significantly outperforms other baseline models, achieving nearly 100% accuracy. This result demonstrates that D3 effectively decomposes unseen combina- torial data into TPR components using discrete dictionaries.
FWM
40
AID
D3 (w/ F)
20
D3 (w/o F)
0
10000
20000
30000
Iteration
Figure 2: Test accuracy curve [%] on the SAR task
for 10 seeds, with shadowed area indicating SD.
Next, we test TPR-RNN and FWM with D3 on the syb- task. This task involves compositional information in each story sentence, such as the relation between objects and their locations. It makes

---

Table 1: The mean word enror rate [%] on the sys-bAbl task for 10 seeds, with indicating SD.
Model
w/o sys diff
sys diff (4)
Gap (4)
# params (4)
TPR-RNNN
0.79 0.16
8.74 3.7.4
7.95
0.14 M
0.69
5.61
4.92
0.32 M
D3
0.65 0.255
3.50 2.07
2.85
0.17 M
FWM
0.79
2.85
2.06
0.73 M
AID
0.45 0.16
1.21 0.66
0.76
1.23 M
D3 (w/o F)
0.79 0.300
2.58 1.12
1.79
0.75 M
D3 (w/ F)
0.75 0.0.7
1.96 0.88
1.21
0.75 M
Table 2: The mean accuracy [%] the sort-of-CLEVR task for 10 seeds, with = indicating SD,
Model
Doude
Unary (t)
Binary (t)
Ternary (1)
# params (1)
Linear Transformer
69.3
75.5
56.4
0.68 M
+ AID
98.9 1
78.6 =03
63.7
0.83 M
+ D3 (w/o F)
128
73.9 16.5
77.2
57.3 4.6
0.75 M
256
73.7 + 16.5
77.8 2.5
57.9 5.8
0.96 M
+ D3 (w/F)
128
98.9 0.0.2
79.5
63.1 1.9
0.80 M
256
99.0 + 0.3
82.1 2.4
68.8 1.2
1.13 M
Table 3: Perplexity the WikiText- 103 task.
Model
Dadla
Valid (4)
Test (4)
# params (1)
Linear Transformer
36.473
37.533
44.02 M
AID
36.159
37.151
44.16 M
D3 (w/o F)
32
36.061
37.220
44.12 M
64
35,975
37.009
44,36 M
D3 (w/ F)
32
36,630
37,620
44.22 M
64
36,220
37.128
44.62 M
sentence-level model more suitable for capturing the structural information of data than 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 understanding tasks.
4.2.2 Linear Transformer
We also evaluate the Linear Transformer with D3 on the sort-of-CLEVR task and WikiText- 103 task. Following the AID [23], of-CLEVR task and apply D3 to a 16-layered Linear Transformer at intervals of out of the 16 layers for the WikiTexx- 103 task. As shown in Tables 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 visual relational reasoning and language modeling tasks, as well as its applicability to the Linear Transformer. In addition, D3 shows comparable performance the attention-ban based decomposition method, even with fewer parameters.
4-layered Linear Transformer with shared parameters for the sort-
4.3
Analysis
In this section, we conduct qualitative analysis the structured TPR representations generated by D3 and ablation study of D3. For these analyses, experiment with D3 (w/o F) on the SAR task.
4.3.1
Qualitative Analysis
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
7

---

(a) query
code
component
Discovery Phase
Discovery Phase
Discovery Phase
Discovery Phase
Figure 3: The heatmap displays the cosine similarity hetween the generated representations during the discovery phase for the SAR task. We explore the similarity across diffferent types of representations: (a) queries of roles, (b) codes of roles, and (c) the roles themselves.
(a) query
(b) code
(c) component
Inference Phose
33
Discovery Phase
Discovery Phase
Discovery Phase
Figure 4: The heatmap displays the cosine similarity between the generated representations during the discovery phase (represented on the x-axis) and the inference phase (represented the y-axis) for the SAR task. We explore the similarity across different types of representations: (a) queries of roles and unbinding operators, (b) codes of roles and unbinding operators, and (c) the roles and unbinding operators themselves.
(a) Codebook Key
(b) Codebook Value
825
075
Figure 5: The heatmap visualizes the cosine similarity of the learned codebook features for the SAR task. There are two parts to each heatmap: (a) the similarity among codebook keys, denoted (k, Neode and (b) the similarity among codebook values, denoted as {v, Neode For better visualization, the heatmap values are reordered to reflect the cluster of similar codebook keys.
role and unbinding operator for the same symbol z. We analyze the orthogonality of generated representations to investigate whether they satisfy these TPR conditions. Specifically, consider the case of varying z 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 TPR conditions, resulting in an accuracy 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 satisfy the TPR condition, indicating the effectiveness of D3.

---

(a)
100
140
Accuracy
Accuracy
Accuracy
60
20
20
128
26
Code Size
Number Keys
Top-k
Figure 6: The mean accuracy on the SAR task for seeds in the ablation study, with error bar indicating SD. The default setting uses Dede of 64, Neode of 64, and top-k of 8, Each figure shows the experimental results for the following settings: (a) Varying Deade. (b) Varying Neode with top-k constant. (c) Varying top-k with Neoie constant.
Furthermore, analyze the similarity patterns of codebook keys and codebook values. Fig. 5 shows that the codebook features result implies that the learnable parameters of dictionaries implicitly capture TPR conditions ensure accurate TPR operations.
learn orthogonal patterns despite being learned without constraints. This
4.3.2 Ablation Study
We investigate the effect of hyper-parameters, of D3, specifically Neode. Deode and top-k, on perfor- mance on the SAR task. Fig. 6(a) shows the effect of Dende. We observe that the value of Dodee significantly affects the performance of D3. Notably, fails solve the SAR task when Deade is set to 8, indicating a need for adequate capacity of D Fig. 6(b) shows the effect of varying top-k while holding Noote constant, indicating that achieves optimal performance when top-k is set 2. This result demonstrates the efficacy of the sparse mechanism employed by D3. Fig. 6(c) examines the effect of varying Noude while holding top-k constant, showing that D3 generally performs better with larger values of Nexde-
5 Discussion and Limitations
Motivation. From the perspective systematic generalization, the decomposition operations in the TPR framework can be viewed mapping unseen data TPR components observed during training. Motivated by this, we design a decomposition module based on discrete representations, which maps input data 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 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- baseb models. This design ensures that each dictionary is trained to specialize specific TPR component.
Interpretability. The TPR framework decomposes data at the representation level into distinct symbols, such role-filler pairs for encoding and unbinding operators for decoding. This character- istic enhances the interpretability of models because the relationships between roles and umbinding 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 rep- resentations that more effectively satisfy the TPR conditions than baseline models like FWM and AID. Figs. 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 such

---

requirements for for fillers, which are features related to downstream tasks. This characteristic affects the performance of D3 depending on whether 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-bAbl and cort-of-CLEVV tasks with relatively few labels (-200). In contrast, the w/o F configuration excels on the SAR and Wiki 103 tasks, which have a larger number of labels (500-). These findings suggest that the w/o 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 (onntexxqueen component. Therefore, this complexity 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.
(Dquery + logk)) for each TPR become drawback as the capacity of the dictionaries
Scalability. The scalability of D3 inherently linked 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 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.
6 Conclusion
In this paper, we tackle the decomposition problem inherent in the TPR framework, which poses significant challenge for TPR-based models. To address this, we introduce discrete dictionary-based layer, D3, designed to enhance the decomposition capabilities of TPR-based models. D3 employs the discrete dictionaries 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 IPR-based models with fewer additional parameters. Furthermore, qualitative analysis verifies that D3 effectively generates structured representations that are satisfactory for the requirements of the TPR framework.
Acknowledgements
This work was supported by the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (No. A2C301A2C301 1169 No. A0226145 No. RS-2022-00166735 & No. RS-2023-00218887).
References
[1] Ba, J. R. Kiros, and G. E. Hinton. Layer normalization. arXiv preprint 1:X:v:6007006500 2016.
[2] J. A. Fodor and Z. W. Pylyshyn. Connectionism and cognitive architecture: A critical analysis. Cognition, 28(1 28(1-2):3-1 -2)33771, 1988.
[3] A. Goyal, A. Lamb, J. Hoffmann, S. Sodhami, S. Levine, Y. Bengio, and B. Schölkopf. Recurrent independent mechanisms. arXiv preprint arXiv: 10893, 2019.
[4] A. Graves, G. Wayne, and Danihelka. Neural turing machines. arXiv preprint arXiv: 1410.5401, 2014.
10

---

[5] A. Graves, G. Wayne, M. Reynolds, T. Harley, I. Danihelka, A. Colmenarejo, E. Grefenstette, T. Ramalho, J. Agapiou, et al. Hybrid computing using a neural network with dynamic external memory. Nature, 538(7666)471-477, 2016.
S. G.
[6] K. Greff, S. Van Steenkiste, and J. Schmidhuber. On the binding problem in artificial neural networks. arXiv preprint arXiv:2012.05208, 2020.
[7] K. Hsu, W. Dorrell, J. Whittington, J. Wu, and C. Finn. Disentanglement via latent quantization. Advances in Neural Information Processing Systems, 36, 2024.
[8] D. Hupkes, V. Dankers, M. Mul, and E. Bruni. Compositionality decomposed: How do neural networks generalise? Journal of Artificial Intelligence Research, 67:757-795, 2020.
[9] Y. Jiang, A. Celikyilmaz, P. Smolensky, P. Soulos, S. Rao, H. Palangi, R. Fernandez, C. Smith, M. Bansal, and J. Gao. Enriching transformers with structured tensor-product representations for abstractive summarization, arXiv preprint arXiv:2 2021.
[10] A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 1156-5165 PMLR, 2020.
[11] A. Kori, F. Locatello, D. S. Ribeiro, F. Toni, and B. Glocker. Grounded object-centric learning. In The Twelfth International Conference on Learning Representations, 2023.
[12] B. Lake and M. Baroni. Generalization without systematicity: On the compositional skills of sequence-to-seyy sequence recurrent networks. In International conference machine learning, pages 2873-2882. PMLR, 2018.
[13] B. M. Lake, T. D. Ullman, J. B. Tenenbaum, and S. J. Gershman. Building machines that learn and think like people. Behavioral and brain sciences, 40:e253, 2017.
[14] G. Lample, A. Sablayrolles, M. Ranzato, L. Denoyer, and H. Jégou. Large memory layers with product keys. Advances in Neural Information Processing Systems, 32, 2019.
[15] H. Le, T. Tran, and S. Venkatesh. Self-attentive associative memory, In International Conference on Machine Learning, pages 5682-569 PMLR, 2020.
[16] A. Liška, G. Kruszewski, and M. Baroni. Memorize or generalize? searching for a compositional mn in haystack. arXiv preprint arXiv:1802.066 00766666..7.,... 2018.
[17] D. Liu, A. M. Lamb, K. Kawaguchi, A. G. ALIAS PARTH GOYAL, C. Sun, M. C. Mozer, and Y. Bengio. Discrete-valued neural communication. Advances in Neural Information Processing Systems, 14:2109-2121, 2021.
[18] F. Locatello. D. Weissenborn, T. Unterthiner, A. Mahendran, G. Heigold, J. Uszkoreit, A. Doso- vitskiy, and T. Kipf. Object-centric learning with slot attention. Advances in Neural Information Processing Systems, 33:1 13:11525-11538 2020.
[19] Merity, C. Xiong, Bradbury, and R. Socher. Pointer sentinel mixture models, arXiv preprint arXiv: 1609.08833. 2016.
[20] S. Mittal, S. C. Raparthy, I. Rish, Y. Bengio, and G. Lajoie. Compositional attention: Disentan- gling search and retrieval. arXiv preprint arXiv:2110.09411 2021.
[21] H. Palangi, P. Smolensky, X. He, and L.. Deng. Question-answering with grammatically- interpretable representations. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018
[22] T. Park, I. Choi, and M. Lee. Distributed associative memory network with memory refreshing loss. Neural Networks, 144:33-48, 2021.
[23] Park, I. Choi, and M. Lee. Attention-based iterative decomposition for tensor product representation. In The Twelfth International Conference on Learning Representations, 2023.
11

---

[24] J. Rae, J. Hunt, I. Danihelka, T. Harley, A. W. Senior, G. Wayne, A. Graves, and T. Lillicrap. Scaling meeory-aaggeentne neural networks with sparse reads and writes. In Advances in Neural Information Processing Systems, pages 3621-3629, 2016.
[25] Rosenblatt. The pereeptron: a probabilistic model for information storage and organization in the brain. Psychological review, 65(6):386, 1958.
[26] A. Santoro, D. Raposo, D. G. Barrett, M. Malinowski, Pascanu, Battaglia, and T. Lillicrap. A simple neural network module for relational reasoning. Advances in neural information processing systems, 30, 2017.
[27] A. Santoro, R. Faulkner, D. Raposo. J. Rac, M. Chrzanowski, T. Weber, D. Wierstra, O. Vinyals. R. Pascanu, and T. Lillicrap. Relational recurrent neural networks. Advances in neural information processing systems, 31, 2018.
[28] I. Schlag and J. Schmidhuber. Learning to reason with third order tensor products. Advances in neural information processing systems, 31, 2018.
[29] I. Schlag, P. Smolensky, R. Fernandez, N. Jojic, J. Schmidhuber, and J. Gao. Enhancing the transformer with explicit relational encoding for math problem solving, arXiv preprint arXiv:1910.06611, 2019.
[30] I. Schlag, T. Munkhdalai, and J. Schmidhuber. Learning associative inference using fast weight memory. arXiv preprint 1xXiv:00110088// 2020.
[31] I. Schlag, K. Irie, and J. Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning, pages 9355-992 536699366. PMLR, 2021.
[32] Z. Shi, Q. Zhang, and A. Lipani. Stepgame: new benchmark for robust multi-hop spatial reasoning in texts. In Proceedings of the AAAI conference on artificial intelligence, volume 36, pages 11321 11329, 2022.
[33] P. Smolensky. Tensor product variable binding and the representation of symbolic structures in connectionist systems. Artificial intelligence, 46(1-2):159-211 1990.
[34] P. Soulos E. J. Hu, K. McCurdy, Y. Chen, R. Fernandez, P. Smolensky, and Gao. Differentiable tree operations promote compositional generalization. International Conference on Machine Learning, pages 22499-325222 PMLR, 2023.
[35] N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever, and R. Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15 (1):1929-1958, 2014.
[36] S. Sukhbaatar, J. Weston, R. Fergus, et al. End-to-end memory networks. In Advances in neural information processing systems, pages 2440-2448, 2015.
[37] A. Tamkin, M. Taufeeque, and N. D. Goodman. Codebook features: Sparse and discrete interpretability for neural networks. arXiv preprint 1722v 2023.
[38] Träuble, A. Goyal, N. Rahaman, M. C. Mozer, K. Kawaguchi, Y. Bengio, and B. Schölkopf. Discrete key-v bottleneck. In International Conference on Machine Learning. pages 34431-34455, PMLR. 2023.
[39] A. Van Den Oord, O. Vinyals, et al. Neural discrete representation learning, Advances in neural information processing systems, 30, 2017.
[40] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and Polosukhin. Attention is all you need. Advances neural information processing systems, 30, 2017.
[41] Webb, I. Sinha, and J. D. Cohen. Emergent symbols through binding in external memory. arXiv preprint arXiv:2012.146011 2020.
12

---

[42] J. Weston, A. Bordes, S. Chopra, A. M. Rush, B. Van Merriënboer, A. Joulin, and T. Mikolov. Towards ai-complete question answering: A of prerequisite toy tasks. arXiv preprint arXiv: 150..00022999 2015.
[43] Y.-F. Wu, M. Lee, and S. Ahn. Structured world modeling via semantic vector quantization. arXiv preprint arXiv:2402.0120. 2024.
[44] X. Zhuang, Q. Zhang, K. Ding, Y. Bian, X. Wang, J. Lv, H. Chen, and H. Chen. Learning invariant molecular representation in latent discrete space. Advances in Neural Information Processing Systems, 36, 2024.
13

---

Appendix
A Experiment Details
This section provides a detailed description of our experiments on the SAR task, sys-bAbl task, sort-of-CLEVR task, and WikiText-103 task. We followed the experimental settings outlined by AID [23] to assess the decomposition capabilities of D3. To ensure stability and reproducibility. we ran all experiments, except for the WikiText- 103 task, using 10 different random seeds'. For the WikiText- 103 task, we experimented with single seed of 1111. Each experiment was conducted on a single 48GB NVIDIA RTX A6000 GPU and an AMD EPYC 7513 32-Core Processor.
A.1 Systematic Associative Recall task
The SAR task [23] evaluates systematic generalization in memorizing and recalling combinatorial data. It consists of discovery phase and an inference phase. During the discovery phase, the model receives the combinatorial sequential items, each combining two symbols,
X and y E Y where X = X1 U X2 U X3 and Y = Yi U Y2 The model is then required to predict an associated when specific r presented. The SAR task uses different combination settings between training and evaluation to target systematic generalization specifically. During the training, the model learns the following combination settings: (1) X1 and Yi. (2) X2 and Y2, and (3) X3 and Y. At evaluation, however, the model should generalize unseen combination settings, specifically X₂ and Y2, In our study, unlike the AID paper [23], we only consider the most challenging setting of the SAR task by excluding the subset X3-
Each combinatorial item is constructed as follows. First, symbols and y are sampled from their respective sets X and Y, where X1| = [X2l = |Yı|= [Y2] 250. The sampled symbols are mapped into 50-dimensional space using word embedding method. These embedding vectors are then concatenated to construct the combinatorial item. For training. 100 randomly generated combinatorial items are sequentially provided to the model during the discovery phase. During the inference phase, the model receives only the symbols sequentially, with the embedding vector of set to zero This task also provides binary flags to indicate the start of each phase. At evaluation, all possible combinations that can be formed in X1 and Y2 are tested.
To build the experimental environment for the SAR task, we utilize the open-source implementation' from the AID [23]. We train the model using the Adam optimizer with a batch size of 64 and a 3 81 of 0.9, and 82 of 0.98 for training iterations of 30K, Each experiment took approximately 3 hours per each seed.
learning rate of le-3.
A.2 Systematic bAbl task
The sys-bAbl task [23] is a variant of the bAbl task [42] designed to evaluate systematic generalization in text understanding and reasoning. It consists of 20 distinct sub-tasks, each comprising stories, relevant queries, and corresponding answers. The bybl task requires the models to remember the stories and predict corresponding answers to the queries. Unlike the original bAbl task, the sys-bAbl 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-bAbl task.
The bAbl dataset includes various versions, such as en-10k and en-valid-10k. The sys-bAbl task uses the en-valid-100 version, which is already divided into training, validation, and test datasets. To create the experimental environment for the sys-bAbl task, we use the open-source implementation' provided by the AID.
*e used the following seed values: (0. 1111, 2222, 3333, 4444, 5555, 6666, 7777, 8888, 9999) ARtask
4 'hups://githubb hhtnps//gihuu..oo//awoooooo hww.://wwaaoooooa/////tttaaaaaan
14

---

We use the open-source implementation of the baseline models, TPR-RNN [28] and FWM [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 batch size of 128 and a learning rate of le 13. f1 of 0.9, and (3 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 le -3, 31 of 0.9, and 32 of 0.98 for training iterations of 60K. Furthermore, following the AID, use the reconstruction loss for the bAbI task, introduced in Park et al. [22], in our experiments on the sys-bAbl 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.
A.3 Sort-of-CLEVR task
The sort-of-CLEER 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 75x75 pixels, includes 6 distinct objects in 6 different colors (red, blue, green, orange, yellow, gray) and 2 different shapes (square or circle). This scene image encoded by 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 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" from Mittal et al. [20]. We train the model using the Adam optimizer with a batch size of 64 and a learning rate of le for 100 training epochs. Each experiment took approximately 2.5 hours per each seed.
A.4 WikiText-103 task
The WikiText-103 task [19] is a language modeling dataset consisting of lengthy corpora from Wikipedia. Although the WikiTeext 103 task does not directly measure the systematic generalization of the models, it is used to evaluate the effectiveness and applicability of D3 on 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 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 imple- mentation from [31]. Following the AID [23], apply D3 to 16-layered Linear Transformer at intervals of 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 and learning rate warmup step of 2,000 for 120 epochs. Each experiment took approximately ~3 days.
http:://ggthuh..mmmiiihagga/aaa/ Weight Weiirss-----tttt--ee-----e----e Memory-
15

---

B Hyper-parameter Settings
Table 4: Hyper-parameter settings of the D3.
SAR task
sys-bAbl task
Cort-of-CLEVe task
WikiText-103 task
Dodte
8, 16, 32, 64, 128
32, 64, 128, 256
128, 256
32, 64
Nooe
64
Duuay
Dodee /2
top-k
Diropout
0.1
Table 5: Hyper Syrra--prsmeetes of TPR-RNN.
sys-bAbl task
Dentiy (D componen)
90
Dreatoon (Domppoem))
20
N
ene cumponent
N Idee composent
Table 6: Hlyper-parameters of FWM.
SAR task
sys-bAbl task
DESTM
256
256
DrwM (D
32
32
Nieads
Nase
companent
3
3
N rdee companent
1+NN
1+N
Table 7:
of Linear Transformer.
Sort-of-CLEVR task
WikiText- 103 task
Ducads (comppmen))
64
16
Nreads
4
8
N Mene
companent
2 $ Naeads
2 * Nuaads
Adce companent
Niedk
Nacds
16

---

C
Additional Experiments
Table 8: The mean word error rate [%] on additional experiments of the sys-bAbl task for 10 seeds.
Model
Dedde
w/o sys diff (4)
w/ sys diff (1)
Gap (4)
params (1)
TPR-RNN
0.79 + 0.16
8.74 + 3.74
7.95
0.14 M
+ AID
0.69
+: 0.88
5.61
+ 1.78
4.92
0.32 M
+ D3
32
1.16 4: 0.25
3,44
1.78
2.28
0.13 M
64
0.65
+ 0.25
3.50
+ 2.07
0.17 M
128
0.68 = 0.14
3.94
2.20
3.26
0.26 M
FWM
0.79 + 0.14
2.85
2.06
0.73 M
+ AID
0.45 :: 0.16
1,21 0.66
0.76
1.23 M
+ D3 (w/o F)
64
0.79 + 0.30
2.58 1.12
1.79
0.75 M
128
0.93 + 0.20
3.82 + 1:21
2.89
0.82 M
256
1.04 +: 0:40
3,33 1.21
2.29
0.97 M
+ D3 (w/ F)
32
1.20 + 0.31
7.23 4.33
6.03
0.71 M
64
0.75 4: 0.17
1.96 0.88
1.21
0.75 M
128
0.89 +: 0.32
2.48 0.67
1.59
0.84 M
256
0.75
3.09
2.34
1,02 M
D Additional Comparisons
In this section, we expand our comparisons to include broader range state-of-the-an methods, as detailed below.
sys-bAbl task. We compare D3 state-of-the-art methods (DAM [22] and STM [15]) on the original bAbl task. Table 9 shows that existing memory networks struggle with the sys-bAbl task, highlighting the efficacy of D3 compared to these state-of-the-ar memory networks.
Table 9: The mean word eтror rate [%] on additional comparison of the sys-bAbl task for 10 seeds.
Model
w/o sys diff (L)
w/sys diff' (L)
Gap (L)
DAM
0.48 0.20
5.25 + 1.64
4.77
STM
0.49 +: 0.16
4.79 +: 1.53
3.70
TPR-RNN
0.79 + 0.16
8,74 4 3.4
7.95
+ AID
0.69 4: 0.08
5.61 +: 1.78
4.92
+ D3
+ 0.25
3.50 + 2.07
2.85
FWM
0.79 0:14
2.85 + 1.61
2.06
+ AID
0.45 0.16
1.21 :: 0.66
0.76
D3 (w/o
0.79 1: 0.30
2.58 : 1.12
1.79
D3 (w/ F)
0.75
1.96
1.21
Sort-of-CLEVR task. We compare D3 to vanilla Transformer [40] and Compositional Transformer [20), designed to enhance the systematic generalization capabilities multi-head self-attention methods. Table 10 shows that the Linear Transformer significantly degrades systematic generalization performance compared to the vanilla Transformer and the Compositional Transformer. While D3 improves the performance of the Linear Transformer from TPR perspective, it still shows limited performance reasoning the relationships between multiple objects (Binary and Ternary) compared to the vanilla Transformer and Compositional Transformer.
WikiText-100 task. We compared D3 the Delta Network [31], which introduced a delta updating rule instead of the additive outer product- based updating rule in the Linear Transformer. Table 11
17

---

Table 10: The mean accuracy [%]] on additional comparison of the sort-of-CLEVR task for seeds.
Model
Deode
Unary (f)
Binary
Termary
Transformer
97.4 3.
84.3 433
62.7 1:39
Compositional Transformer
98.9
88.4 +
66.5
Linear Transformer
69.3 + 148
75.5 +:
56.4
98.9
78.6
63.7
D3 (w/o F)
128
73.9 16.5
77.2
57.3
256
73.7 + 16.5
77.8 2.
57.9 5
D3 (w/ F)
128
98.9 4: 0.2
79.5 0.08
63.1
256
99.0
0.3
82.1 2:24
68.8
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.
Table 11: Perplexity additional comparison of the Wiki'Tex- 103 task.
Model
Dende
Valid (1)
Test ())
Delta Network
35.640
36.659
Linear Transformer
36.473
37.533
AID
36.159
37.151
D3 (w/o F)
32
36.061
37,220
64
35.975
37.009
D3 (w/ F)
32
36.630
37,620
64
36.220
37.128
E
Additional Ablation Study
In this section, we extend our ablation studies to investigate the effects of varying the number keys in the codebook and the impact of removing either the residual connection the codebook from the D3 layer.
The Efffet of Varying the Number of Codebook Keys. Fig. 7 shows that even with 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 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 layer on the SAR and sys-bAbI tasks. Specifically, removing the codebook means that the components are generated solely by the shared -hed-forwwrdd networks (layer, residual and layer while removing the residual connection implies that the components are derived solely from the codebook values.
Accuracy
60
<0
24
120
16
32
54
128
Number of Keys
Figure 7: The mean accuracy on the SAR task for 10 seeds in the ablation study for the effect of varying Noute from 128 with top-k constant.
18

---

The Effect of Residual Connection. Fig. 8 shows that without the residual connection, the generalization performance of dramatically degrades. This result indicates that the residual connection is crucial for effectively training the D3 layer.
(b)
100
Residual
Error Rate
Residual
Accuracy
60
False
10
False
True
True
5
20
w/o sys
w/ sys
Figure 8: Ablation study for the effect of the residual connection on (a) the SAR task and (b) the sys-bAbl task for 10 seeds.
The Effect of Codebook. Table 12 shows that even without the codebook ("w/o 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, which may explain why the model maintains robust performance even with fewer keys.
However, it is important to that without the codebook, the D3 layer does not achieve near-perfect accuracy on the SAR task (as shown in Table 12) and fails to significantly enhance the systematic generalization the baseline model on the sys-bAbI task (as shown in Table 13). These results demonstrate that the codebook plays crucial role in enhancing the model's 's overall performance and generalization capabilities, especially tasks requiring systematic generalization.
Furthermore, we experiment with Noude on the SAR task, where the codebook may act as a bias term. The results in Table 12) show that using single codebook element leads to degraded generalization performance compared to the "w/o codebook" configuration, indicating that multiple codebook elements are essential for achieving optimal results.
Table 12: Ablation study for the effect of the codebook on the SAR task for 10 seeds.
Model
D
Nox
top-k
Accuracy (1)
FWM
44.90 3:.5
D3
87.38 11.10
+ D3
32
64
99.27 + 0.8
+ D3 (w/o codebook)
89.00 .9.02 d: 4,56
+ D3
89.10 10 +: 7.99
94.44 2.35
64
D3
64
94.29
D3 (w/o codebook)
91.65
Table 13: Ablation study for the effect of the codebook on the sys-bAbl task for 10 seeds.
Model
w/o sys diff (4)
sys diff (4)
Gap (4)
FWM
0.79 0.14
2.85 :: 1.11
2.06
0.75 0.17
1.96 0.88
1.21
(w/o codebook)
1.19
3.55 1.04
2.36
Discussion. Our ablation study on codebook the SAR task (Table 12) indicates that the shared residual networks within the D3 layer significantly enhance generalization performance. However, the results from the sybl task (Table 13) suggest that while these networks improve performance, they alone struggle to generalize more structured data.
19

---

The ablation studies in Tables 12 and 13 demonstrate that incorporating the codebook mechanism leads nearly 100% accuracy on the SAR task and significantly improves the systematic generalization of models on the sys-bAbl task. However, 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 layer significantly improves the systematic generalization capabilities of TPR-bas based models.
F
Additional Qualitative Analysis
Comparison to Baselines
We conduct an orthogonal analysis for the baseline models (FWM and AID) similar to the analysis presented in Section 4.3.1. Figs. and indicate that the D3 model generates more structured and orthogonal representations than the baseline models, FWM and AID, demonstrating its effectiveness.
(a) FWM
(b) AID
(c) D3
Discovery Phose
H. Discovery Phase
Discovery Phase
Discovery Phase
Figure 9: The heatmap displays the cosine similarity between the roles during the discovery phase for the SAR task.
(a) FWM
(b) AID
(c) D3
Inference Phase
Discovery Phase
Discovery Phase
Discovery Phase
Figure 10: The heatmap displays the cosine similarity between the roles (x-axis) during the discovery phase and the unbinding operators (y-axis) during the inference phase for the SAR task.
20

---

F.2
Qualitative Analysis for Different Seeds
Additionally, we present the results of the qualitative analysis for different seeds in the SAR task.
F.2.1
Neode: 64, Dodee 32, top-k: 8, seed: 3333
(a) query
(b) code
(c) component
Discovery Phase
Discovery
Discovery Phase
Discovery
Figure 11: The heatmap displays the cosine similarity between the generated representations during the discovery phase for the SAR task. We explore the similarity across different types of representations: (a) queries of roles, (b) codes of roles, and (c) the roles themselves.
(a) query
(b) code
(c) component
Inference Phase
Discovery
Discovery
Discovery Phase
The heatmap displays the cosine similarity between the generated representations during the discovery phase (represented the x-axis) and the inference phase (represented the y-axis) for the SAR task. We explore the similarity across different types of representations: (a) queries of roles and unbinding operators, (b) codes of roles and unbinding operators, and (c) the roles and unbinding operators themselves.
Figure 12:
(a) Codebook Key
(b) Codebook Value
8.5
035
8.0
059
Figure 13:
The heatmap visualizes the cosine similarity of the learned codebook features for the SAR task. There are two parts to each heatmap: (a) the similarity among codebook keys, denoted
(k, Noode and (b) the similarity among codebook values, denoted {V, {V:) Noode For better visualization, the heatmap values are reordered to reflect the cluster of similar codebook keys.
21

---

F.2.2
Neode: 64, Dedde: 32, top-k: 8, seed:
4444
(a) query
(b) code
(c) component
Discovery Phase
Discovery Phase
Discovery Phase
Discovery Phase
Figure 14:
The heatmap displays the cosine similarity between the generated representations during the discovery phase for the SAR task. We explore the similarity across diffferent types of representations: (a) queries of roles, (b) codes of roles, and (c) the roles themselves.
(a) query
code
component
Inference Phase
Linn Discovery Phase
Discovery Phase
Discovery Phase
The heatmap displays the cosine similarity between the generated representations during the discovery phase (represented on the x-axis) and the inference phase (represented on the y-axis) for the SAR task. We explore the similarity across different types representations: (a) queries of roles and unbinding operators, (b) codes of roles and unbinding operators, and (c) the roles and unbinding operators themselves.
Figure 15:
(a) Codebook Key
(b) Codebook Value
Figure 16: The heatmap visualizes the cosine similarity of the learned codebook features for the SAR task. There are two parts each heatmap: (a) the similarity among codebook keys, denoted as [ Naude and (b) the similarity among codebook values, denoted as {v; | Neade For better visualization, the heatmap values are reordered to reflect the cluster of similar codebook keys.
22

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
This puper includes the paper's contributions and scope the abstract and introduction, as follows. This paper tackles the decomposition problem inherent in the TPR-based approaches. To address this, this paper proposes discrete dictionary-baser decomposition (D3) layer designed to enhance the decomposition capabilities of the TPR- based models.
Justification:
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 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 is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: This paper discusses the limitations of the work performed by the authors in Section 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 encouraged to create separate "I "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results 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 the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
23

---

3. Theory Assumptions and Proofs
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 cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that 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 reproduce the experimental results. This paper explains the mechanism of the proposed model and how it is applied existing baseline models in Section and presents the experiment details and hyper- parameter settings in Appendices A and B.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
24

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
This paper provides supplementary materials to reproduce all experimental results of the proposed method, including source codes about our model implementation, data processing, scripts for execution,
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm cc/ Batcc/guides/Codeeee
www.ssssiiiiicyyllc for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
/ /nips. /twpsscc/ppubliiie wwtpccccwpppppccccc.eeeccoccc gww./// CCooeooooiisseeeeeeeeo ssssnnoiccy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: This paper presents experiment details and hyper-parameter settings in Appendices A and B.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to a level detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: This paper reports the mean and standard deviation values in the experimental results conducted using fixed 10 different random seeds.
25

---

Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: This paper provides the computer resources used in experiments and the time took to learn each task in Appendix A.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https htap:::///euree http:///eeuu.peeeeeee wwtpublicc/thieeetiiiiiiie
Answer: [Yes]
Justification:
The research conducted in this paper conforms to the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
26

---

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 why the paper does not address societal impact.
Examples of negative socictal impacts include potential malicious unintended uses (c.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 being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper does not have a high risk for misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. 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 existing assets.
27

---

The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license ge.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.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thwcee com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-pase 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 s creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: This paper provides supplementary materials with source code, license, and README md files. The README. files cite the code packages utilized in this paper and provide all the instructions to reproduce the experimental results.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mooel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
28

---

Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
29


Reviews:
[{"review_id": "6UtWPwOiTq", "reviewer": "Reviewer_KCxS", "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.", "questions": "The questions naturally arise from the weaknesses above.", "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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Originality: The D3 method presents a novel approach to addressing the decomposition problem in TPR-based models. The use of discrete, learnable dictionaries for this purpose is innovative.\n\n* Quality: The paper demonstrates thorough experimentation across multiple tasks and provides detailed ablation studies to support its claims.\n\n* Clarity: The method is explained clearly, with helpful visualizations and step-by-step descriptions of the D3 layer.\n\n* Significance: Improving the decomposition capabilities of TPR-based models has potential implications for enhancing systematic generalization in neural networks, which is an important goal in AI research.", "weaknesses": "* 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.\n* 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.\n* While the paper provides some analysis of the learned representations, a deeper theoretical analysis of why D3 works well could enhance the contribution.\n* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer for their constructive feedback and insightful suggestions. We will ensure that we reflect on our responses in our revised manuscript in the future.\n\n> W1: Comparison to AID\n\nWhile our work and AID focus on enhancing the systematic generalization of TPR-based models by addressing a decomposition problem, our contributions differ from AID's in several important aspects. Our main contribution is **introducing a novel decomposition method for the decomposition of TPR**, which stores discrete features during training and maps given data to these learned discrete features via sparse activation.\n\nIn detail, AID employs a competitive attention mechanism between input features and intermediate TPR representations, refining these representations iteratively. However, as shown in Fig. 2, AID has difficulties generalizing to unseen compositions of known symbols. Additionally, AID does not explicitly leverage observed structural information during training. To address these limitations, we propose the D3 layer, which maps input data to the nearest pre-learned discrete features within dictionaries, generating structured TPR representations. Our comprehensive experiments demonstrate that the D3 layer significantly improves the systematic generalization of TPR-based models.\n***\n\n> W2: Design choices in the D3\n\nOur work focuses on enhancing the systematic generalization of TPR-based models. Since systematic generalization aims at generalizing unseen data composed of known components observed during training, the decomposition of TPR can be thought of as mapping input data to observed TPR features during training. This insight inspired our design choices for the D3 layer.\n\n**Our primary design choice was to employ a discrete representation-based decomposition module, which maps input data to discrete features learned during training**. Inspired by a prior key-value architecture [1], we introduced separate key-value-based dictionaries, each linked to specific TPR components (*role, filler, and unbinding operator*). This design ensures that each dictionary captures the symbolic information of the specific TPR components. The D3 layer generates *queries* based on the input data, which are then used to search and identify discrete features associated with the input data. Once relevant discrete features are identified, the codebook values are aggregated to generate *codes*. These *codes and queries* are then utilized as components for TPR operations to solve downstream tasks.\n\nIn summary, each design choice in the D3 layer—from the multiple discrete dictionaries to the generation and utilization of *queries and codes*—was motivated by **the goal of mapping input data to discrete, learned features that facilitate systematic generalization in the decomposition operation of TPR**.\n\n[1] Large memory layers with product keys, NeurIPS’19.\n\n***\n\n> W3: A deeper theoretical analysis\n\nWe appreciate the suggestion to provide a deeper theoretical analysis of why D3 works well.\n\nOur D3 employs a separate key-value-based discretization mechanism, the robustness of which against distributional shifts was theoretically investigated in prior work [2]. **The D3 layer enables TPR-based models to mitigate errors in generating structured representations by mapping input data to pre-learned discrete features**. Figs. 3 and 4 empirically demonstrate how D3 mitigates errors while generating structured representations. Initially, generated *queries* fail to satisfy the TPR conditions, potentially causing inaccuracies in TPR operations. D3 addresses this by leveraging the *codes* derived through the discretization and generates near-ideal structured representations. These results imply that D3 improves the robustness of TPR-based models to unseen data comprising known symbols, thus enhancing the model's overall performance.\n\nFurthermore, in response to Reviewer FjyM, we have included further orthogonal analysis for the baseline models. Please see the global response. Figs. 14 and 15 show the superior quality of representations produced by D3 compared to the baseline models, further demonstrating the effectiveness of D3.\n\n[2] Discrete key-value bottleneck, ICML’23\n\n***\n\n> W4: Additional experimental comparison\n\nOur experiments evaluated the enhancement of systematic generalization in TPR-based models achieved by D3. We compared D3 with several TPR-based baselines across various systematic generalization tasks. Additionally, to assess the effectiveness of D3 in more complex scenarios, we extended our evaluation to the WikiText-103 task. **While those tasks considered in our study are relatively simple, the results consistently demonstrate that D3 significantly improves the generalization performance of TPR-based models**. We acknowledge the comment regarding the need for testing on more complex, real-world tasks. Although our current evaluation provides strong evidence of D3's benefits in systematic generalization tasks, further investigation is needed to confirm its effectiveness in real-world scenarios. We plan to address this limitation in future work.\n\nIn response to your suggestion, we have expanded our comparisons to include a broader range of state-of-the-art methods. Please see the global response.\n\n***\n\n> L1: Scalability of D3\n\nThe scalability of our approach 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 [1]. However, this approach requires further investigation and experimentation, which we plan to research in future work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jKYiuCi4Rq", "reviewer": "Reviewer_QJJM", "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.", "questions": "Please check out the Weakness section first. I listed additional questions as follows:\n- 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?\n- 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": "Please check out the Weakness section.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- S1: Generating components in the proposed method is analogous to transformer blocks with Top-K selection, which is straightforward and intuitive. \n- S2: Experimental results and ablation studies are well-organized and highlight the effectiveness of the proposed method.", "weaknesses": "- 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. \n- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer for their constructive feedback. We will ensure that we reflect on our responses in our revised manuscript in the future.\n\n> W1: Technical novelty over AID and scalability issue\n\nWhile the prior work introduced an iterative competitive attention-based decomposition module, called AID, which refines structured representations iteratively, it has limitations in generalizing to unseen compositions of known symbols even in simple synthetic tasks, as shown in Fig. 2. AID lacks an explicit mechanism to leverage observed structural information during training for decomposition of TPR.\n\nTo address these limitations, **we propose a novel discrete representation-based decomposition method for TPR-based models, which stores discrete features during training and maps given data to these learned discrete features via sparse activation**.\n\nOur experimental results in Tables 2 and 3 demonstrate that D3 achieves comparable performance to AID with a similar parameter setting ($D_\\text{code}$ of 128 on the CLEVR and 32 on the Wiki). Moreover, D3 performs better with increased parameters (256 on the CLEVR task and 64 on the Wiki). **These results highlight our method's technical novelty and effectiveness compared to AID**.\n\nThe scalability of our approach 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 [1]. However, this approach requires further investigation and experimentation, which we plan to research in future work.\n\n[1] Large memory layers with product keys, NeurIPS’19.\n\n***\n\n> W2: How can you ensure that the shared dictionary satisfies these properties during training to satisfy near orthogonality between roles or between unbinding operators?\n\nThe TPR framework operates by explicitly decomposing data at the representation level into distinct symbols, such as *role-filler* pairs, which are then encoded through the tensor product of *role* vectors and *filler* vectors, $T= filler \\otimes\\textit{role}$. By doing so, this encoding method preserves the symbolic structure of the data. During decoding, the framework retrieves specific *fillers*—essential for solving the given task—via matrix multiplication using *unbinding operators* associated with particular *roles*, $filler = T \\cdot \\textit{unbind}$. These TPR characteristics force TPR-based models to learn to generate structured representations satisfying TPR conditions through supervised training.\n\n**When integrating with TPR-based models, our D3 layer is trained to produce structured representations that satisfy these TPR conditions**. D3 employs a sparse key access mechanism during the generation of these representations. This mechanism ensures that individual discrete features within dictionaries become specialized to specific latent data features. Consequently, this specialization allows each discrete feature to learn and indirectly satisfy the near orthogonality requirements of TPR properties, thereby maintaining accurate TPR operations.\n\n***\n\n> Q1: Difference in experimental settings of the SAR task between our work and the AID paper\n\nIn the original AID paper, the SAR task performance was evaluated with varying levels of the task parameter $p$ with different values (0.0, 0.5, and 1.0), which adjusts the combinations of symbols observed during training. The AID achieved about 90% accuracy for $p$ values of 0.5 and 1.0, which are less challenging settings.\n\n**Our experiment adopts the SAR task's most challenging setting ($p=0.0$)**. This setting results in a more rigid evaluation environment than the other configurations. Consequently, the AID method does not achieve the same level of performance in our setup as it did in the original paper.\n\nTo ensure clarity and prevent potential misinterpretation, we will include a detailed description of these differences between our work and the AID paper in the revised manuscript. This will help readers understand the context and rationale behind the observed performance disparities.\n\n***\n\n> Q2: How can these redundant properties of keys not violate the TPR operation conditions?\n\nThe D3 layer generates structured TPR representations by mapping input data to pre-learned discrete features within dictionaries through three steps: (1) *query* generation, (2) sparse key access, and (3) aggregation of code values. Therefore, the keys within dictionaries are intermediate features that assist the D3 layer in generating structured representations that satisfy the TPR conditions rather than needing to satisfy those conditions themselves.\n\nWhile some codebook keys in Fig. 5(a) show strong similarities, their corresponding codebook values in Fig. 5(b) show near-orthogonal patterns. This implies that **when similar *queries* access the dictionaries via these codebook keys, the resultant codes (read from the dictionaries) are orthogonal**. Figs. 3 and 4 provide further evidence supporting this claim.\n\n**As illustrated in Figs. 3 and 4, even though the *queries* are similar, the *codes* (retrieved from dictionaries via the *queries*) show more orthogonal patterns than the *queries* themselves**. Additionally, Figs. 3(c) and 4(c) demonstrate that D3 successfully generates *components* (which are our main focus) that meet the TPR conditions using these intermediate features (*queries and codes*).\n\nWe believe this comprehensive approach ensures that the redundant properties of the keys do not violate TPR operation conditions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "08thuCRA3j", "reviewer": "Reviewer_3Nfz", "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.", "questions": "-\tHow is the TPR-with-D3-based models compared with the SOTA method on their individual tasks?\n-\tDoes the use of the D3 layer increase the interpretability of the models? \n-\tHow do the dictionaries trained with D3 correspond to three components (roles, fillers, and unbinding operators) in TPR?", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "-\tThe use of discrete dictionary-based method in TPR models seems novel. I am not working in the area. This novelty is claimed in the paper.\n-\tThe related work section is well structured and clearly mentions the difference between the proposed method and existing ones.\n-\tThe experiments were conducted with multiple TPR-based models on multiple tasks and the proposed method achieved the best results on most cases (except for one case in Table 1 in which FWM+AID is the best).", "weaknesses": "-\tThe 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. \n-\tD3 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. \n-\tThe 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. \n-\tI 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.\n-\tIt 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.\n-\tThere 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer for their constructive feedback and insightful suggestions. We will ensure that we reflect on our responses in our revised manuscript in the future.\n\n> W1: More explanation for some key concepts\n\nWe acknowledge that the initial description may have been challenging for readers unfamiliar with the TPR framework. We provide more detailed explanations of the key terms to enhance the comprehensibility of our paper as follows.\n\nThe TPR is a general method for representing the symbolic structure of data using distributed representations. **This framework operates by explicitly decomposing data at the representation level into distinct symbols, such as *role-filler* pairs, which are then encoded through the tensor product of *role* vectors and *filler* vectors**, $T= filler \\otimes\\textit{role}$. By doing so, this encoding method preserves the symbolic structure of the data. The *roles and fillers* are dependent on the task at hand.\n\nFor instance, in a tree structure, the *role* corresponds to the position within the tree, while the *filler* represents the associated label with that position [1]. In associative memory, the *role* is analogous to an associative key (or address), and the *filler* corresponds to the associative value (or contents) [2, 3].\n\nDuring the decoding phase, **the TPR framework retrieves specific *fillers*—essential for solving the given task—**from the encoded TPR representation via matrix multiplication using *unbinding operators* associated with particular *roles*, $filler = T \\cdot \\textit{unbind}$.\n\nIn TPR-based neural networks, during training, **these TPR characteristics force the models to learn to generate structured representations satisfying TPR conditions through supervised training**. This process ensures the models can perform correct TPR operations to solve tasks accurately.\n\n[1] Differentiable tree operations promote compositional generalization, ICML’23.\n\n[2] TPR-RNN, NeurIPS’18.\n\n[3] FWM, ICLR’21.\n\n***\n\n> W2 & Q1: Comparison with the SOTA models\n\nOur primary objective was to evaluate **how D3 enhances the systematic generalization capabilities of TPR-based models**. Therefore, **our initial experiments compared D3 with several TPR-based baselines across multiple tasks**. However, we acknowledge the importance of comparing our D3-enhanced TPR models with SOTA methods for a comprehensive evaluation.\n\nIn response to your inquiry, we have expanded our comparisons to include a broader range of state-of-the-art methods. Please see the global response.\n\n***\n\n> W3 & Q2: Interpretability of D3\n\nThe 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 of TPR improves the interpretability of models** because the relationship between *the roles and the unbinding operators* explains which parts of the input the model focuses on to predict the answer. However, **this interpretability is accurate only when the generated structured representations satisfy the TPR conditions**.\n\nIn this context, **the D3 layer enhances the interpretability of the models by providing structured representations that more effectively satisfy the TPR conditions compared to baseline models**. In response to Reviewer FjyM, we have included experimental results of the orthogonal analysis for the baseline models. Please see the global response. Figs. 14 and 15 demonstrate that the generated representations by the D3 better confirm the TPR conditions than other baseline models, supporting our claim that the D3 layer contributes to increased interpretability.\n\n***\n\n> W4: The term \"symbolic\" features in the manuscript\n\nThank you for your valuable input. Although we intended to convey that the codebook features within the dictionaries capture information that could be associated with symbolic components, it is more accurate to refer to these features as \"discrete\" rather than \"symbolic\" since they remain in vector form and are numeric. **We will revise the term \"symbolic feature\" to \"discrete feature\"** to prevent any potential misunderstanding (specifically on lines 13, 60, 61, 71, and 116).\n\n***\n\n> W5 & Q3: How do the dictionaries trained with D3 correspond to TPR components?\n\nWe appreciate the opportunity to clarify the correspondence between the dictionaries trained with D3 and the TPR components in TPR.\n\nAs discussed in response to W1, during the training phase, the decomposition module in TPR-based models learns to generate TPR components to perform TPR operations. **Each dictionary in D3 is explicitly linked to a specific TPR component, ensuring that each dictionary is responsible only for generating its corresponding component**. The generated components are then utilized in pre-defined TPR operations of the TPR-based models. **This setup ensures that each dictionary is trained to be specialized to a specific TPR component**.\n\nMoreover, the number of dictionaries directly correlates with the TPR operations of the baseline models. For instance, the TPR operations of the FWM require two types of roles and one filler for encoding ($T= filler \\otimes\\textit{role}_1 \\otimes \\textit{role}_2$) and two types of unbinding operators for decoding ($filler = T \\cdot \\textit{unbind}_2 \\cdot \\textit{unbind}_1$). We thus set up three distinct dictionaries when integrating D3 with FWM (one for $filler$, another for $role_1/unbind_1$, and the other for $role_2/unbind_2$).\n\nWe opted for the configuration \"*with roles/unbinding operators*\" to allow the D3 layer to contribute to generating structured representations that satisfy the TPR conditions, regardless of the number of dictionaries. Your suggested option, \"*w/o roles/unbinding operators but with fillers*,\" presents an interesting alternative for examining the design of D3 and could be considered for future studies.\n\n***\n\n> W6: Small errors in Figs. 6(b) and (c)\n\n**We will correct the captions**.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LUk2vBJYfJ", "reviewer": "Reviewer_FjyM", "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.", "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?\n\nOn 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.\n\nCalling WikiText 103 a large-scale language modeling dataset in 2024 is a stretch. Please just call it \"language modeling\" dataset.", "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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors present an interesting, novel, discrete representation of learning. Learning such representations is notoriously difficult. Their method is the only one that can solve the systematic split on the SAR task. The evaluation is solid. The authors have analyzed the learned roles and unbinding operators and showed that they are learning orthogonal representations without explicit regularization.", "weaknesses": "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. \n\nAdditionally, 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. \n\nI believe that this additional analysis could deepen the understanding of why and how D3 works and significantly strengthen the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer for their constructive feedback and insightful suggestions. We will ensure that we reflect on our responses in our revised manuscript in the future.\n\n> W1: Ablation study for the effect of the residual connection\n\n**The *components* represent the vector representations of the symbolic components of TPR**, including roles, fillers, and unbinding operators, which are essential for TPR operations in TPR-based models. Specifically, D3 uses *queries* to access the codebook and generates *codes* based on the values in the codebook, as detailed in Eq. 3. Subsequently, it produces the *components* by residually connecting the *queries* to the *codes*, as described in Eq. 4. This mechanism implies that without the residual connection, the *components* are derived solely from the codebook values.\n\nIn response to your suggestion, **we conducted an ablation study on the SAR and sys-bAbI tasks** to examine the impact of removing the residual connection. Please see the global response. **Fig. 13 shows that the residual connection is crucial for effectively training the D3 layer**.\n\n***\n\n> W2: Additional orthogonal analysis on baseline models\n\nThank you for your valuable input. In response, **we have conducted an orthogonal analysis for the baseline models** (FWM and AID) similar to the analysis presented in Section 4.3.1. Please see the global response. **Figs. 14 and 15 indicate that the D3 model generates more structured and orthogonal representations than the baseline models**, FWM and AID, demonstrating its effectiveness.\n\n***\n\n> Q1: Could the authors provide an additional explanation for why discretization aids the decomposition of unseen data, as opposed to denoising the representations?\n\nThank you for your insightful comment. Our work focuses on enhancing the systematic generalization of TPR-based approaches. Since systematic generalization aims at generalizing unseen data composed of known components observed during the training phase, the decomposition operations of TPR can be thought of as mapping unseen data to observed TPR components during training.\n\nThe discretization technique enhances this decomposition operation by preserving the discrete features learned during training and enabling the mapping of given data to discrete representations. Within the D3 layer, **these discrete representations facilitate the generation of structured TPR representations that satisfy the TPR conditions**, as demonstrated in Figs. 3 and 4.\n\nMoreover, Figs. 3 and 4 illustrate that **the D3 layer also plays a role in denoising representations** by summing the *codes* to the *queries*, **effectively reducing the noise of the *queries* in the *components***. These characteristics allow the D3 layer to enhance the systematic generalization of TPR-based models by leveraging the discrete features to their decomposition operations of TPR.\n\n***\n\n> Q2: Ablation study for the effect of varying the number of keys\n\nThank you for your valuable input. In response to your comment, **we have expanded the range of our ablation study** to include additional experimental results for varying the number of keys in the codebook. Please see the global response. **Fig. 16 shows that even with a significantly reduced number of keys, the model with D3 maintains high accuracy on the SAR task**. This observation highlights the significance of the architectural inductive bias introduced by the D3 layer, effectively addressing the decomposition problem inherent in TPR-based models. These findings enhance our understanding of the codebook's impact on the model's performance.\n\n***\n\n> Q3: The term \"large-scale language modeling\" in the manuscript\n\n**We will update the terminology to \"language modeling\"** to reflect its current standing better.\n\n***\n\n> L1: 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.\n\nWe appreciate the opportunity to address the concern regarding the transferability of improved generalization from toy tasks to real-world scenarios, a common limitation in TPR-related research.\nPrevious work [1] has demonstrated the equivalence between TPR (or fast weight) and the linear attention mechanism, highlighting the potential for TPR to enhance linear attention. Additionally, recent works [2, 3] have proposed hybrid quadratic-linear attention methods that mitigate the computational complexity of self-attention while improving the generalization capabilities of LLMs.\n\nGiven these findings, **we believe that the techniques developed in prior TPR-related research could address the limitations of linear attention within these hybrid methods, potentially improving the generalization of LLMs**. We plan to explore this promising direction in our future work.\n\n[1] Linear Transformers Are Secretly Fast Weight Programmers, ICML’21.\n\n[2] Transformer quality in linear time, ICML’22.\n\n[3] When Linear Attention Meets Autoregressive Decoding: Towards More Effective and Efficient Linearized Large Language Models, ICML’24", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
## Global response

We thank all the reviewers for their constructive feedback and insightful suggestions. We believe that the additional ablation studies and comparisons they recommended provide a clearer understanding of D3's strengths and limitations. Our revised manuscript will thoroughly reflect these improvements.

We have attached a one-page PDF that includes the additional ablation studies suggested by reviewer FjyM, which are as follows:
- Ablation study on the effect of the residual connection (Figure 13)
- Additional orthogonal analysis on baseline models (Figures 14 and 15)
- Ablation study on the effect of varying the number of keys (Figure 16)

Moreover, as reviewers KCxS and 3Nfz suggested, we have expanded our comparisons to include a broader range of state-of-the-art methods, as detailed below.

**sys-bAbI task**: We compared D3 to state-of-the-art methods (DAM [1] and STM [2]) on the original bAbI task. The results in the table below show that existing memory networks struggle with the sys-bAbI task, highlighting the efficacy of D3 compared to these state-of-the-art memory networks.

| **sys-bAbI task** | *w/o sys diff* | *w/ sys diff* | Gap |
| --- | :---: | :---: | :---: |
| TPR-RNN | 0.79$_{\pm0.16}$ | 8.74$_{\pm3.74}$ | 7.95 |
| TPR-RNN (+AID) | 0.69$_{\pm0.08}$ | 5.61$_{\pm1.78}$ | 4.92 |
| TPR-RNN (+D3) | 0.65$_{\pm0.25}$ | 3.50$_{\pm2.07}$ | 2.85 |
| FWM | 0.79$_{\pm0.14}$ | 2.85$_{\pm1.61}$ | 2.06 |
| FWM (+AID) | 0.45$_{\pm0.16}$ | 1.21$_{\pm0.66}$ | 0.76 |
| FWM (+D3 *w/ F*) | 0.75$_{\pm0.17}$ | 1.96$_{\pm0.88}$ | 1.21 |
| --- |
| DAM [1] | 0.48$_{\pm0.20}$ | 5.25$_{\pm1.64}$ | 4.77 |
| STM [2] | 0.49$_{\pm0.16}$ | 4.19$_{\pm1.53}$ | 3.7 |

**Sort-of-CLEVR task**: We included the Compositional Transformer [3], designed to enhance the systematic generalization capabilities of multi-head self-attention methods. (our experiment was performed under identical experimental settings as this prior work [3].) The results show that the Linear Transformer significantly degrades systematic generalization performance compared to the Transformer. While D3 improves the performance of the Linear Transformer from a TPR perspective, it still shows limited performance in reasoning the relationships between multiple objects (*Binary* and *Ternary*) compared to the Transformer and Compositional Transformer.

| **Sort-of-CLEVR task** | $D_\text{code}$ | *Unary* | *Binary* | *Ternary* |
| --- | :---: | :---: | :---: | :---: |
| Linear Transformer | - | 69.3$_{\pm14.8}$ | 75.5$_{\pm1.3}$ | 56.4$_{\pm4.3}$ |
| Linear Transformer (+AID) | - | 98.9$_{\pm0.2}$ | 78.6$_{\pm0.3}$ | 63.7$_{\pm1.2}$ |
| Linear Transformer (+D3 *w/ F*)| 128 | 98.9$_{\pm0.2}$ | 79.5$_{\pm0.8}$ | 63.1$_{\pm1.9}$ |
| Linear Transformer (+D3 *w/ F*) | 256 | 99.0$_{\pm0.3}$ | 82.1$_{\pm2.4}$ | 68.8$_{\pm1.2}$ |
| --- |
| Transformer | - | 97.4$_{\pm3.5}$ | 84.3$_{\pm4.3}$ | 62.7$_{\pm3.9}$ |
| Compositional Transformer [3] | - | 98.9$_{\pm0.2}$ | 88.4$_{\pm1.4}$ | 66.5$_{\pm1.9}$ |


**WikiText-103 task**: We compared D3 to the Delta Network [4], which introduced a delta updating rule instead of the additive outer product-based updating rule in the Linear Transformer. (our experiment was performed under identical experimental settings as this prior work [4].) The results indicate 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.

| **WikiText-103 task** | $D_\text{code}$ | Valid | Test |
| --- | :---: | :---: | :---: |
| Linear Transfomer | - | 36.473 | 37.533 |
| Linear Transfomer (+AID) | - | 36.159 | 37.151 |
| Linear Transfomer (+D3 *w/o F*) | 32 | 36.061 | 37.220 |
| Linear Transfomer (+D3 *w/o F*) | 64 | 35.975 | 37.009 |
| --- |
| Delta Network [4] | - | 35.640 | 36.659 |

[1] Distributed associative memory network with memory refreshing loss. *Neural Networks*, *144*, 33-48.

[2] Self-attentive associative memory, ICML’20.

[3] Compositional Attention: Disentangling Search and Retrieval, ICLR’22.

[4] Linear transformers are secretly fast weight programmers, ICML’21.

Meta-review: 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.
Decision: Accept (poster)