Paper: Associative Transformer is a Sparse Representation Learner (ICLR-2024)

Reviewer comment (Experiments):
Allowing samples across batches to compete via the bottleneck attention is inelegant -- inherently, this competition means that the ViT's prediction on an input is dependent on what samples are in the same batch. What is the performance if we use batch size = 1 during evaluation? What about different batch sizes? Unless I am misunderstanding something, it would be necessary to report standard deviation across different evaluation batch sizes when reporting metrics for this kind of network.

Author response:
**Q1: Allowing samples across batches to compete via the bottleneck attention is inelegant.**

ViT's prediction for an input is not solely dependent on the samples within the same batch. To elaborate on this, the memory bank update utilizes Exponential Moving Average (EMA) with a smoothing factor $\alpha = 0.9$ (that determines the rate at which older observations decay) in Equation 4, accumulating information from both old $\gamma^{t}$ and new memories $\hat{\gamma}^{t}$. 

$\gamma^{t+1} = \alpha \cdot \gamma^{t} + (1 - \alpha) \cdot \hat{\gamma}^{t}$.

This memory update strategy with a high smoothing factor $\alpha = 0.9$ based on EMA ensures the stable performance of the proposed architecture with varying batch sizes. The learned explicit memory encompass not only information from the current batch but also from all batches observed thus far. Therefore, given that the overall training data distribution remains constant, using different batch sizes is not expected to significantly impact model behavior.

The only factor that might affect the model's behavior concerning batch size is that, intuitively, a larger batch size may introduce stronger competition among patches with a constant bottleneck size. In our experiments, we employed grid search to tune the patch size, batch size, and bottleneck size. In P13 Appendix A.3 "Experimental Settings and Hyperparameters", we provided detailed settings. 

We aimed to maintain a ratio between the bottleneck size and the total patch numbers, as shown to be effective in previous studies on the Coordination method. For the CIFAR and Pet datasets, we used a bottleneck size of 512, selected from a pool of 32.8k/25.1k patches, respectively. For the Triangle dataset, we used a bottleneck size of 64 from a pool of 2.0k patches. For the relational reasoning tasks, we used a bottleneck size of 256, selected from a pool of 14.4k patches. The difference with the Coordination method is that we introduced the competition among different samples. This allows a sample's classification task, for instance, to benefit from other patches within the batch input.

When using a batch size of 1 during evaluation, the proposed bottleneck attention with the squash layer degrades to the traditional sample-wise bottleneck in the Coordination method. We could adjust the bottleneck size for evaluation to a smaller value, for example 1, when the task is CIFAR10 with a patch size of 4. Alternatively, using a fixed memory learned during training for evaluation is possible. The memory remains static during evaluation, and the batch is directly sent to the Hopfield network for reconstruction based on the fixed memory.

Reviewer comment (Reproducibility):
I do not find the attention maps in Fig 2 meaningful. Which layer's memory bank was used to display this attention map? How is the attention heatmap smooth if it operates on patches (not pixels)? These questions are not answered in the paper.

Author response:
**Q2:  I do not find the attention maps in Fig 2 meaningful. Which layer's memory bank was used to display this attention map? How is the attention heatmap smooth if it operates on patches (not pixels)? These questions are not answered in the paper.**

Thank you for pointing this out. The attention map pertains to the memory bank of the first layer, representing the attention scores across different **patch** locations (8 rows and 8 columns, totaling 64 locations, with one pixel for each). The attention maps are computed by averaging the attention scores across different samples, resulting in the smooth visualization shown in Figure 2.

Reviewer comment (Evaluation):
I suspect there are many more FLOPS needed for this network if the Hopfield Network component of the GWT sub-block employs recurrence. If true, this should be reported and quantified as a limitation of the model.

Author response:
**Q3: I suspect there are many more FLOPS needed for this network if the Hopfield Network component of the GWT sub-block employs recurrence. If true, this should be reported and quantified as a limitation of the model.**

Regarding the functioning of the Hopfield network, it does not employ recurrence during the end-to-end training and features a decoupled memory (the explicit memory bank) for efficient training.

In detail, Equation 7 describes an iteration that could be applied multiple times, but typically, we use just a single step for forward and backward computation during end-to-end training. Here, $t$ represents the batch time step, and the iteration time step is implicit. The Hopfield network utilizes a single forward step to recall patterns in the memory bank. In every retrieval, the global energy is guaranteed to decrease over time step(s), as illustrated in Figure 10.

Moreover, the Hopfield network does not need to compute and store inner states for every batch input, which could increase the FLOPS. Instead, we utilize the explicit memory bank (Figure 1.a) for the update of memory slots. Alternatively, the Hopfield network can also incorporate an implicit memory by introducing trainable parameters. In this case, the patterns are stored in the weights of these parameters requiring more FLOPS. Differing from the implicit memory, in the proposed architecture, we employ the same memory architecture for learning both the priors for the bottleneck attention and the attractors for the associative memory in the Hopfield network.

Reviewer comment (Presentation):
Caption of Fig 1 says: "...squashed into vectors $\mathbb{R}^{(B \times N) \times E}$". Is this not a matrix, where you combine the batch and token dimension? Perhaps this means a collection of vectors, but then the figure itself is confusing because we see only one vector passed into the bottleneck attention.

Author response:
**Q4.1: Caption of Fig 1 says: "...squashed into vectors $\mathbb{R}^{(B\times N)\times E}$". Is this not a matrix, where you combine the batch and token dimension? Perhaps this means a collection of vectors, but then the figure itself is confusing because we see only one vector passed into the bottleneck attention.**

Yes, it is a collection of vectors. We revised the Figure 1 where images are represented by different colors. Patches from the same image are distinguished by a gradient with each square representing the vector $\mathbb{R}^{E}$ of a patch.

Reviewer comment (Presentation):
"MASK" is not defined in Fig 1.

Author response:
**Q4.2: "MASK" is not defined in Fig 1**

The "Mask" in Figure 1 was the top-$k$ hard attention. We revised the figure and replaced the Mask accordingly.

Reviewer comment (Presentation):
In Fig 1, there are two blocks of "Explicit Memory" -- do these represent the same matrix?

Author response:
**Q4.3: In Fig 1, there are two blocks of "Explicit Memory" -- do these represent the same matrix?**

Yes, the two blocks of "Explicit Memory" represent the same matrix, but at different time steps. The first explicit memory in Figure 1 corresponds to the memory slots $\gamma^t$ at the time step $t$, while the second explicit memory represents the updated memory $\gamma^{t+1}$ based on the bottleneck attention. The illustration of the second explicit memory block is intended to demonstrate that the updated memory $\gamma^{t+1}$ is employed as the input to the Hopfield network.

Reviewer comment (Theory):
In Fig 1, why are there three inputs to the Hopfield Network `(Explicit Memory, LT(explicit memory), state)`? Where is the update equation for the Hopfield Network?

Author response:
**Q4.4: In Fig 1, why are there three inputs to the Hopfield Network (Explicit Memory, LT(explicit memory), state)? Where is the update equation for the Hopfield Network?**

Thank you for pointing this out! The inputs to the Hopfield network are twofold, instead of the three that were shown in Figure 1. We revised the figure 1 and the two inputs are the output of the linear transformation of the explicit memory $f_{\text{LT}}(\gamma^{t+1})$, and the current patch representations (the input state) $\Xi^t$.

The update equation of the Hopfield network is based on the modern Hopfield proposed in [1]. The update rule for a patch representation $\xi^t \in \Xi^t$ in the Hopfield network is $\hat{\xi}^t=f_{\text{LT}}(\gamma^{t+1}) \text{softmax}(\beta (f_{\text{LT}}(\gamma^{t+1}))^T \xi^t)$, which can be derived from Equation 7. A skip connection is employed to obtain the final output $\Xi^{t+1}=\hat{\Xi}^t+\Xi^t$. 

[1] Hubert Ramsauer, Bernhard Schafl, Johannes Lehner, and et al. Hopfield networks is all you need. ICLR 2021.

Reviewer comment (Experiments):
The paper does not include any ablation studies.

Author response:
**Q1: Ablation studies**

We performed several ablation studies in P7 Section 5.3 Ablation study and P8 Section 5.4 Comparison with the coordination method. In particular, the Table 2 "Comparison based on an ablation study" and the Figure 4 "Model size vs. accuracy for configurations" were discussed in details.

| Models               | CIFAR10 | CIFAR100 | Triangle | Average   |
|----------------------|---------|----------|----------|-----------|
| AiT                  | **83.34** | **56.30** | **99.47** | **79.70** |
| Reset memory         | 81.94 | 55.96 | 99.46 | 79.12 |
| W/O Hopfield         | 81.03   | 54.96    | 99.44    | 78.48     |
| W/O memory (ViT)     | 79.53   | 53.19    | 99.47 | 77.40 |
| Dense networks       | 77.78   | 53.14    | 99.46    | 76.79     |
| W/O bottleneck       | 75.40   | 46.53    | 93.33    | 73.75     |
| W/O SA               | 72.72   | 47.75    | 99.46    | 73.31     |
| W/O FF               | 69.51   | 40.89    | 97.61    | 69.34     |

Reviewer comment (Experiments):
The experimental evaluation of the model is limited to a few vision datasets.

Author response:
**Q2: Datasets**

Following the previous studies, we experimented on the datasets of the Triangle, CIFAR10, CIFAR100, Oxford-IIIT Pet, and Sort-of-CLEVR. The detailed explanation of these datasets was clarified in P12 Appendix A.1 Datasets.

Reviewer comment (Experiments):
How does the model's performance vary when different types of priors are used in the workspace memory?

Author response:
**Q4: How does the model's performance vary when different types of priors are used in the workspace memory?**

Thank you for pointing this out! All priors are used during the training however with different expertise to process different patches. In Figure 2, we visualize the attention maps generated by each prior. These maps highlight the same specific locations across all input samples during the selection of relevant image patches for different priors.

Additionally, in P15 Figure 9, the examples of the selected image patches by the prior-guided attention are visualized. As training progresses, a more diverse selection of patches can be obtained.

Reviewer comment (Experiments):
How does the model's performance vary when different types of attention mechanisms are used?

Author response:
**Q5: How does the model's performance vary when different types of attention mechanisms are used?**

We provided a comprehensive evaluation of the proposed attention mechanism, both in comparison to existing models, and through internal performance analysis in the context of ablation studies. First, the comparison between the proposed global workspace layer (Figure 1.a) and conventional sparse attention-based workspace methods is shown in Table 1, including the Coordination, Perceiver, Set Transformer, and Luna methods. Moreover, we conducted an ablation study regarding the attention mechanism to understand how the model's performance varies when different types of attention mechanisms are used. In Table 2, under the row labeled "W/O Hopfield," we measured the model's performance when replacing the Hopfield network with another cross-attention mechanism. The added cross-attention has the same architecture as the multi-head attention in Figure 1.a, but without applying the bottleneck attention. Unlike the multi-head attention in the global workspace layer, the added cross-attention in the ablation study takes the current input state as the query, using the contents in the explicit memory as the keys and values. The output of the cross-attention is used as the output of this layer. The ablation study in Table 2 demonstrated the efficacy of the proposed method by investigating different attention mechanisms of the cascaded cross-attention and the self-attention components.

Reviewer comment (Evaluation):
What are the limitations of the proposed model?

Author response:
**Q6: What are the limitations of the proposed model?** 

We discussed several limitations in Section 9. First, the intricacies of managing multiple modalities simultaneously pose a current challenge in the methodology, although AiT demonstrates superior performance in its capability to attend to and retrieve resembling patches from the shared workspace (please refer to Section 5.8 Relational Reasoning). We aim to investigate multi-modal competition within the shared workspace, enabling tasks to benefit from the cross-modal learning of distinct perceptual inputs. For example, when the input vectors are from both image patches and text tokens, the bottleneck attention could adaptively select vectors from different modalities and learn cross-modal priors that guide attention.

Additionally, in Section 5.7, we investigated the effect of the inverse temperature, $\beta$, on information retrieval based on the Hopfield networks in Figure 5. The limitation arises in the need for tuning the $\beta$ value to reach optimal performance. We aim to study a mechanism to adjust the $\beta$ adaptively in the future, addressing this limitation and potentially further improving performance.

Reviewer comment (Experiments):
The baseline seems not include SOTA models in vision tasks.

Author response:
**Q7: The baseline seems not include SOTA models in vision tasks.**

A crucial aspect of this study is its exclusive comparison of models that exhibit most of the properties outlined in the Global Workspace Theory. In Appendix A.2 'Comparison of Work Related to Global Workspace Theory', Table 5 demonstrates the comparison of attention architectures based on the properties of the Global Workspace Theory and how existing studies are related to these properties. The Vision Transformer (the first row) serves as an example to illustrate its inapplicability to implement the shared workspace. Therefore, other related conventional Transformer models typically do not fall within the scope of sparse attention-based workspace approaches.

The SOTA method in the context of sparse attention workspace-based approaches is the Coordination method by Goyal et al. (2022) [1], where a shared coordination space was proposed to correlate different information fragments. To investigate further, we conducted an additional ablation study in Section 5.4. The detailed comparison with the coordination method is presented in Figure 4, where we showed that integrating all three components (C+LM+HN+SA) resulted in a competitive accuracy of 71.49\% with a compact model size of only 1.0M. These results highlight that the proposed AiT is an efficient sparse representation learner, outperforming the existing SOTA method of Coordination. Additional models compared in this study include the Set Transformer, Luna, Perceiver, and BRIMs, all of which exhibit most of the properties outlined in the Global Workspace Theory. We are pleased to provide clarification if any points may be unclear or if additional results are necessary for understanding the AiT architecture and its contributions to the existing literature.

[1] Anirudh Goyal, Aniket Rajiv Didolkar, Alex Lamb, and et al. Coordination among neural modules through a shared global workspace. ICLR 2022.