Discrete Dictionary-based Decomposition Layer for Structured Representation Learning

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.

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. Additionally, I would like to see analysis like Figures 3 and 4, but for the baselines without D3. Currently, in 282, the authors claim "indicating the effectiveness of D3", but it is unclear how orthogonal representations are the baselines' representations. I believe that this additional analysis could deepen the understanding of why and how D3 works and significantly strengthen the paper.

Questions

In line 115, the authors write: "This distinctive characteristic allows D3 to leverage the learned symbolic features to decompose unseen data after training.". Could the authors provide an additional explanation for why discretization aids the decomposition of unseen data, as opposed to denoising the representations? On Fig 5 (b), the authors show an ablation over the number of keys in the codebook. Even with the smallest key size, the model works well. I would like to see a similar plot going down in the number of keys as long as the performance doesn't suffer significantly or the number of keys doesn't reach 1. This would help to understand better the effect of the codebook. Calling WikiText 103 a large-scale language modeling dataset in 2024 is a stretch. Please just call it "language modeling" dataset.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

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.

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.

Strengths

- The 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. - The related work section is well structured and clearly mentions the difference between the proposed method and existing ones. - The 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

- The paper assumes that the readers are familiar with terminologies used in the paper. It would be better if the paper provided definitions of some key concepts frequently used in the paper (such as roles, fillers, etc.). I had to look into references to fully understand the concepts. - D3 is evaluated with multiple TPR-based models on multiple tasks (such as text understanding and reasoning, visual relational reasoning and large language modeling). But it is unclear how these TPR-with-D3-based models compares with the SOTA models for these individual tasks. Without such a comparison, the significance of the work is unclear. - The paper mentioned at the beginning that using discrete representation can help with the interpretability of the models. However, there is no discussion or evaluation of this aspect of the model as the result of using D3. - I am not sure if the resulting “discrete” representation of the input data can be called “symbolic” features (which the paper uses a few times). The representations are still in the vector form, which is numeric. They may be considered “discrete”, but they are not symbols. - It is not clear how the dictionaries trained with D3 correspond to three components (roles, fillers, and unbinding operators) in TPR. Each component seems to be trained in the same way. There do not seem to be distinctions among them according to the descriptions in Section 3.1. In the experiments, only two or one dictionary is used, with one dictionary corresponding to roles and unbinding operators and an optional one for the fillers. If only the number of dictionaries matters, then the “w/o F” option could be considered “w/o roles/unbinding operators but with fillers”. More explanation on this would be beneficial. - There are some small errors in the presentation. For example, Subfigues (b) and (c) in Figure 6 should be switched to match their references in the text.

Questions

- How is the TPR-with-D3-based models compared with the SOTA method on their individual tasks? - Does the use of the D3 layer increase the interpretability of the models? - How do the dictionaries trained with D3 correspond to three components (roles, fillers, and unbinding operators) in TPR?

Rating

5

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

Yes.

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.

Strengths

- S1: Generating components in the proposed method is analogous to transformer blocks with Top-K selection, which is straightforward and intuitive. - 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. - W2: Although the proposed method is a drop-in layer applicable to other existing TPR methods, it raises the question of how such an attention key-query-value mechanism can meet the TPR operation conditions. In particular, how can you ensure that the shared dictionary satisfies these properties during training to satisfy near orthogonality between roles or between unbinding operators?

Questions

Please check out the Weakness section first. I listed additional questions as follows: - Q1: In Fig 2, the experimental results of the SAR task appear to show a huge difference in performance between the proposed method and AID. In the original AID paper, it achieved about 90% accuracy for the SAR task. The original experiment setup that AID ran on appears to be more difficult than yours, so why did AID perform significantly worse on your setup? - Q2 (related to W2): In Fig. 5(a), there are some strong similarities among specific keys, which indicates some redundancy of the learned keys. How can these redundant properties of keys not violate the TPR operation conditions?

Rating

6

Confidence

3

Soundness

2

Presentation

3

Contribution

3

Limitations

Please check out the Weakness section.

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.

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. * Quality: The paper demonstrates thorough experimentation across multiple tasks and provides detailed ablation studies to support its claims. * Clarity: The method is explained clearly, with helpful visualizations and step-by-step descriptions of the D3 layer. * 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. * There are design choices in the D3: dictionaries, query, code, components and the ways they are connected. There should be clear motivations for each of these choice. * While the paper provides some analysis of the learned representations, a deeper theoretical analysis of why D3 works well could enhance the contribution. * The experiments are primarily focused on synthetic or relatively simple tasks. Testing on more complex, real-world tasks would strengthen the paper's claims about generalization. The comparison to baselines could be expanded to include a wider range of state-of-the-art methods in compositional generalization.

Questions

The questions naturally arise from the weaknesses above.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

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.

Concerns about the usefulness codebook

We are thankful to the authors for their efforts in the rebuttal, and they answered most of my questions positively. However, I find the results of "Q2: Ablation study for the effect of varying the number of keys *extremely concerning*. The performance of the network doesn't change with even as low as 2 keys in the codebook. This begs the question of whether the codebook, which is the main component of the model, is doing anything, or are there other differences in the training pipeline or the model that cause the observed difference in Fig 2 compared to the baselines? This makes me question the validity of the approach. Do the authors have an explanation of how this is possible? Can you please run additional ablations with (1) the codebook removed, but otherwise use the same code and pipeline, (2) the number of codebook elements set to 1, which is equivalent to having a linear layer down-projecting to 1 elements and then projecting it back. I would appreciate any evidence that shows that the codebook is crucial, like the one I described above, or one that show that codebooks with very small number of entries do not work on other tasks.

Response to the reviewers' rebuttal

Thank you for the authors' clarification and for addressing my concerns. I have decided to increase the score.