Paper: On the generalization capacity of neural networks during generic multimodal reasoning (ICLR-2024)

Reviewer comment (Experiments):
While the paper’s studies show that certain designs (e.g. cross-attention) seem to confer multi-modal generalization, there are still some key questions that can be more thoroughly studied to uncover the reasons why this is the case.

Author response:
In response to the Reviewer’s concerns (and the related comment by Reviewer a4Su), we have now performed additional experiments that focus on how model scale and complexity can influence multimodal generalization. While the original manuscript was focused on understanding how a class of base neural architectures would fare on multimodal generalization, we agree that it is important to understand how choice of hyperparameters, such as number of attention heads and layer depth, can influence generalization. 

As such, we have performed a systematic experiment on a standard encoder-only transformer (with the same architecture as BERT; Devlin et al., 2019). We manipulated the number of layers (1, 2, 3, 4 layers), and number of attention heads (1, 4, 8 heads) and assessed the corresponding generalization performance across all splits. (Note that 4 transformer encoder layers and 8 attention heads match the BERT-small architecture.) Indeed, we found that increasing encoder depth significantly improved distractor and systematic generalization. Increasing attention heads also improved distractor and systematic generalization, though to a lesser extent. Nevertheless, neither of these modifications influenced productive generalization. Our conclusion from this is that some tasks (e.g., systematic generalization) require more abstractions; a single layer of attention that handles multimodal inputs is insufficient. Cross-attention mechanisms offer a more targeted and efficient solution (I.e., fewer number of parameters and faster to train) that explicitly integrate visual stimuli (keys and values) with the instructions (queries). However, simply adding encoder/attention layers and parameters can suffice for some types of generalization. 

We have now included discussion of some of these experiments and insights into the manuscript, with a few key revisions copied below. The new results are depicted in Figure 8 (Appendix). Below is the caption for Figure 8; please see the updated PDF submission for the actual figure: “Figure 8: Evaluating generalization splits on BERT-like single-stream transformer models with varying layers and attention heads. We manipulate a generic encoder-only transformer based on the BERT architecture, evaluating the influence of the number of encoder layers (1, 2, 3, and 4 layers), and the number of attention heads per encoder layer (1, 4, and 8 heads). Overall, increasing layers improves generalization across distractor and systematic generalization, but not productive generalization. Increasing attention heads also marginally improves distractor and systematic generalization, but to a lesser extent than adding depth. A) Evaluation on distractor generalization across all model parameters. B) The effect of adding additional encoder layers on distractor generalization performance (averaged across all attention head configurations). C) The effect of adding attention heads on distractor generalization performance (averaged across all layer depth configurations. D-F) Evaluation on systematicity for depth 1 tasks (identical to generalization split in Fig. 4a). G-I) Evaluation on systematicity for depth 3 tasks (identical to generalization split in Fig. 4d). J-L) Evaluation on productivity split (identical to generalization split in Fig. 5a).” 

Updates to Contributions Section (1.2): “2. A comprehensive evaluation of commonly-used base neural models (RNNs, GRUs, Transformers, Perceivers) on distractor, systematic, and productive generalization splits. We find that for distractor and systematic generalization, including a cross-attention mechanism across input modalities is important. However, all models fail on the productivity split. In addition, we include experiments demonstrating the impact of transformer depth and attention heads on all generalization splits in an encoder-only Transformer model."

Reviewer comment (Evaluation):
Important discussions such as why the (cross-attention) transformers might fail at productive generalization is lacking.

Author response:
This is a challenging question to tackle. Our ongoing hypothesis is that productive generalization is a fundamentally distinct type of generalization relative to systematic compositional generalization. We have now included a brief discussion in the Results section of Productive Compositional Generalization (Section 3.3) that addresses these challenges: “One possible explanation for the disparity between systematic and productive generalization in neural models is that systematicity requires the ability to exchange semantics (or tokens) from a known syntactic structure (e.g., a tree of certain depth). In contrast, productive generalization requires generalizing to an entirely new syntactic structure (e.g., a task tree of different size or depth). This requires understanding the syntax -- how to piece together syntactic structures on-the-fly -- requiring another level of abstraction. To our knowledge, there is nothing in the current set of mechanisms in the Transformer that would enable this. Thus, productive compositional generalization remains a difficult capability for purely neural models to achieve.”

Reviewer comment (Theory):
What is the key architectural difference between dual stream transformer and transformers with cross attn that can explain their generalization performance? Is it only the lack of a cross attention between the different modalities?

Author response:
The short answer is yes. When comparing the Dual Stream Transformer with the models with cross attention, indeed, the only distinction is the lack of an attention mechanism to explicitly integrate outputs from the two input streams.  

The longer answer, which became clear after performing the new experiments (that scaled up to the BERT-small architecture), is that at minimum, a second attention layer is required to systematically abstract token information from each of the inputs. While applying a large self-attention matrix twice to simultaneously presented visual and language instructions is computationally inefficient (our cross-attention models show that it is unnecessary), it provides the necessary base structure to allow for good systematic generalization.

Reviewer comment (Writing):
Possible typo: “Finally, we included a Perceiver-like model (Jaegle et al., 2021), an architecture designed to generically process multimodal inputs (Fig. 2f).”:  (Fig. 2f) > (Fig. 2e).

Author response:
We thank the Reviewer for spotting this error. The manuscript has now been updated.

Reviewer comment (Experiments):
I'm concerned about the strength of the baselines used in the paper (see my related questions below). While the primary contribution of the paper is the dataset, it is also important to establish strong baselines for this new dataset and to ensure that the conclusions from the empirical results are valid. The appendix states that only a *single Transformer layer* with a *single attention head* was used. This is almost certainly not an optimal depth and number of attention heads. Relatedly, it looks like some models are potentially underfit, according to the figures. With >5M training examples and a relatively simple input space, I would have expected a reasonably sized Transformer model to achieve low training loss and reasonable IID generalization. If these models could have been applied to similar tasks such as gSCAN (even using symbolic tokens to represent the scene context), where they could be compared with comparable baselines from prior work, this would have helped establish that these are indeed reasonably strong baselines that have been well tuned.

Author response:
We thank the Reviewer for their thorough and thoughtful feedback. Below, we have worked to address some of the weaknesses the Reviewer raised, particularly the strength of the baselines. We have included new experiments to directly address these concerns, taking into consideration this Reviewer’s suggestion. Below, we also address all questions directly.

However, before addressing the Reviewer’s specific questions, we first wanted to address the two weaknesses the Reviewer raised. The first weakness was the lack of evaluation of deeper networks with more attention heads. We have now directly addressed this by performing additional experiments on a range of transformer layer depths (1, 2, 3, 4 layers) and attention heads (1, 4, and 8 heads), up to the size of BERT-small (4 layers, 8 heads; 12 model experiments in total). We have now included figures that detail the performance across this parameter sweep, with the primary conclusion that adding layers indeed aids with distractor and systematic generalization, but not productive generalization. We provide additional details below, and have incorporated the results for these new experiments in Figure 8 (in the Appendix), and have referenced these results in the manuscript's main text. In addition, we wanted to address another comment the Reviewer raised: “*I would have expected a reasonably sized Transformer model to achieve low training loss and reasonable IID generalization.*”  

We thank the Reviewer for their suggestion on expanding our evaluations. The Reviewer was correct in their intuition – deeper models help standard single stream transformers achieve improved IID generalization across all generalization splits (Fig. 8). We have now included the following text to the main manuscript (Results section 3.2) to reference these results: "(Note, however, that increasing depth (encoder layers) to Transformers improves IID generalization on these splits; Fig. 8)."

Reviewer comment (Novelty):
The qualitative difference between gCOG and datasets from prior work such as gSCAN was not very clearly described. For example, one of the key claims seemed to be gCOG "employs generic feature sets that are not tied to any specific modality". However, it seems like it is a useful property for a multimodal dataset to have a clear relation to real-world multimodal tasks. Indeed, the authors provide interpretations of their tasks in the form of natural language instructions and visual scenes (e.g. in Figure 1), and these are very useful for understanding the task. Representing this dataset using familiar modalities (e.g. vision, natural language) could enable future work to study different research questions, e.g. the impact of pre-training. The ability to alternatively represent the task input as a sequence of tokens is also reasonable for studying certain research questions, but this also seems possible for datasets from prior work. For example, I understand that gSCAN includes both symbolic descriptions as well as visual renderings. Anyways, I think clarifying the motivation for this dataset (e.g. increasing diversity of available benchmarks, focusing on different generalization challenges, etc.) separately from how inputs are represented for the experiments in this paper (e.g. token sequence vs. images and natural language) would be useful.

Author response:
The second weakness was the lack of clear distinction between our presented task, gCOG, and prior tasks such as the gSCAN task. We thank the Reviewer for this comment, and have now emphasized the primary distinctions between the two tasks. In brief, the two tasks require different neural network architectures. In terms of transformers, gSCAN is a generation task (requiring a decoder architecture) and gCOG is a classification task, which requires only an encoder Transformer. While models evaluated on gSCAN can incorporate encoder models to the architecture, as illustrated in one of the articles the Reviewer cites (Qiu et al., 2021), generating navigation instructions autoregressively adds complexity to studying compositional productivity, since autoregressive models are susceptible to exposure bias (Wang & Sennrich, 2020). Additionally, gCOG includes a distractor generalization split; our dataloaders are organized such that distractor generalization splits can interact with either systematic and/or productivity splits. We have now emphasized these differences in the manuscript, and have included citations to the studies mentioned by the Reviewer in the Related work sections. 

Related work revisions: 

“Our approach to constructing arbitrarily complex compositions of simple tasks is similar to (Ruis et al., 2020). However, it differs in three key ways. First, we focus on question-answer tasks (which require encoder-only architectures), rather than sequence-to-sequence learning tasks (which require decoder architectures, e.g., Qiu et al., 2021). Sequence decoding tasks introduce the added complexity of requiring autoregressive responses, which are susceptible to fundamental statistical challenges, such as exposure bias (Wang & Senn, 2021). Second, gCOG includes a distractor generalization split, in addition to systematic and productive compositional generalization splits. Finally, we methodically characterize different forms of generalization using simpler underlying abstractions (i.e., without the explicit use of image pixels).” 

We thank the reviewer for bringing to our attention Furer et al. (2021), and Shaw et al, (2021). References to these two studies are now included in the preceding paragraph of the above quoted passage.

Reviewer comment (Presentation):
Appendix A.2.1 - Maybe reference Tables 8 and 9 where you discuss different positional embeddings.

Author response:
*Position embeddings - Since you are representing 10x10 grids as 1D sequences, 1D relative positions may not capture this structure well. On the other hand, absolute position embeddings seem potentially problematic in the case of the SSTrfmr model, since they will not be consistently assigned to the same grid position if the text sequence is first and has varying length. Mitigating this may be important to provide for a fairer comparison with the SSTrfmr model.* 

We agree that a downside of the SSTfmr is that there is ambiguity in how different modalities interact with fixed or 1d positional encodings. This was part of the motivation for comparing that model with the DSTfmr and the CrossAttn models, where each modality would have its own set of positional encodings, and then lower-level features would then be integrated with either a shared MLP or cross-attention mechanisms. 

However, the additional experiments we performed in response to the Reviewer’s question 1 should demonstrate that adding more layers ameliorates the poor performance in the SSTfmr. Nevertheless, we agree that this is a potential issue for other problems that may arise, and could potentially be solved by either using separate positional encodings (for different modalities), and/or learnable positional encodings. However, we believe that this is beyond the scope of the current project, since adding additional encoder layers suffices for the systematic generalization split of this task.

Reviewer comment (Novelty):
Consider discussing [3] in related work. [3] demonstrated the importance of cross-modal attention for gSCAN, and similarly studied the relative difficulty of various aspects of generalization, including distractors.

Author response:
We have additionally included discussion of Qiu et al., 2021 in the results section reporting that distractor generalization is improved using cross-modal attention: “While all models performed IID generalization well, only models that contained cross-attention mechanisms (CrossAttn and Perceiver models) exhibited excellent OOD distractor generalization (Fig. 3d). A related result was also reported in Qiu et al., (2021) using cross-modal self-attention.”

Reviewer comment (Experiments):
**Pre-trained models** The paper focuses on models trained from scratch rather than pre-trained. This could be a strength and a weakness. On the one hand, it allows for isolating the contribution of the architectural choices from other factors of optimization, and training data. On the other hand, it has been observed that by training models at large enough scales enables the emergence of generalization capabilities, which we don’t see in smaller scales. I think it will be critical to also analyze the performance of pretrained models on the benchmark, in order to strengthen the paper.

Author response:
Weakness 1: Lack of evaluation using pre-trained models. We agree with the Reviewer that there is utility in assessing how a pretrained model performs on a new task to the literature. However, when trying to address this question regarding our specific task, we realized that most models would not be able to perform this task out-of-the-box without any fine-tuning, since tokens in the current task set up have no intrinsic meaning. Since we could not directly test mainstream pretrained transformer models (without devising a specific way of first aligning domains through finetuning), we addressed a related question raised by the Reviewer: How would mainstream architectures (such as BERT-like models; Devlin et al., 2019) fare on our benchmark? Moreover, much of the motivation for this manuscript was to focus on questions such as: What architectural components are essential for distractor, systematic, and productive generalization? To this end, we performed additional experiments to evaluate how standard architectures – such as BERT-like architectures -- generalize appropriately. These experiments focused on investigating how increasing encoder-layer depth, and increasing attention heads in each encoder layer, influenced generalization performance. These results have now been included as Figure 8 in the Appendix. Nevertheless, we agree that evaluating the out-of-the-box performance of pretrained models on a pixel and natural language variant of this task is important for future work to explore (once a sensible fine-tuning procedure to align the two is agreed upon).

Reviewer comment (Presentation):
**COG task**: It will be useful to discuss the COG task (rather than just mentioning it) before describing the new gCOG one, so that it will be clearer to the reader what are new contributions of the new benchmark compared to COG and the degree of their importance. In the overview diagram I would also recommend showing a sample also from COG to make the differences clearer.

Author response:
*COG task: It will be useful to discuss the COG task (rather than just mentioning it) before describing the new gCOG one, so that it will be clearer to the reader what are new contributions of the new benchmark compared to COG and the degree of their importance. In the overview diagram I would also recommend showing a sample also from COG to make the differences clearer.*

We thank the Reviewer for the suggestion. We have now clarified the explicit differences between COG and gCOG in the Experimental Design section (2.1): “gCOG is a configurable question-answer dataset, originally inspired from COG (Yang et al., 2018), that programmatically composes task instructions, and then generates synthetic stimuli to satisfy those instructions on-the-fly (Fig. 1). The primary modifications in gCOG are 1) differences in the set of task operators, 2) the ability to use categorical tokens to allow for generic testing of multimodal reasoning, and 3) the ability to allow for arbitrarily long task trees to assess productive compositional generalization, in addition to distractor and systematic generalization (e.g., see Appendix Fig. 7). Importantly, the original COG task did not allow for tasks with more than a single conditional statement, e.g., a task tree of depth 3, making it ill-suited to evaluate productive compositional generalization... We additionally provide functionality in the dataset that allows the choice to load samples using either a categorical task encoding, or a task encoding with image pixels and natural language task instructions."

Due to copyright and the license of the original COG paper, we were unable to include a figure from the original COG task/paper (Springer publishing). However, we have included sample questions/queries from the COG task in the Appendix: “A few example queries from the original COG task include: `What is the color of the latest triangle? Point to the latest red object. If a square exists, then point to the current x, otherwise point to the last b.’”

Reviewer comment (Presentation):
**Figures**: Would be good to increase the size of the plots in Figure 3b. It will also be good the increase the distance and visual separation between the sub-figures in each figure throughout the paper.

Author response:
*Figures: Would be good to increase the size of the plots in Figure 3b. It will also be good the increase the distance and visual separation between the sub-figures in each figure throughout the paper.* 

We have now increased the size of the plots in Figure 3b, splitting panel 3b in to 3b and 3c. We have also worked to increase the amount of visual separation between panels in Figure 2, but note for some figures this was difficult due to space constraints.