Paper: Adaptivity and Modularity for Efficient Generalization Over Task Complexity (ICLR-2024)

Reviewer comment (Writing):
The task definition is unclear.

Author response:
We would appreciate if you could tell us which parts of the task definition you found unclear. We will do our best to clarify this as much as possible in the camera ready version (Please see appendix G in the updated version in updated version of the paper).

Reviewer comment (Experiments):
In the experiment, the authors do not compare the state-of-art methods.

Author response:
The focus of the paper has been to study the ability of transformers to generalize to example of higher complexity while allocating compute proportionally. Hence, our baseline is a vanilla transformer and we compare how augmenting the vanilla transformer with adaptive compute time and modularity enhances their generalization and efficiency. We think it makes sense to study transformers since they seem to be the main building block of current state-of-the-art models.  After showing that it is challenging for vanilla transformers to generalize in the specific setting of the C-PVR task, we show that adaptivity and modularity have complementary effects on improving the generalization and efficiency of transformer models. 
If you could clarify what are the other methods/models that you would like us to add to our comparisons, we will try to include them in the future versions of the paper.

Reviewer comment (Theory):
The details on how and where ACT is being used in Hyper-UT/HyperU-ViT is lacking. It is also similar for Hyper-Module. While there's diagram that shows the attention-based router, it is unclear how attention-based router is formulated (i.e., there's no equation).

Author response:
Thank you very much for the review and the interesting question. Please see our shared response. Also, here are the response to your question and clarification on your remark on the lack of details. 

`Lack of details for ACT and Hyper-Module in Hyper-UT architecture:`

We will explain the ACT and the Hyper-Module and their corresponding hyper-parameters with more details in the appendix.

**ACT mechanism:**
We apply the ACT mechanism per token, and for the details of the method we refer to the original paper (Graves, 2016), and for how it is incorporated into the transformer architecture we refer to the universal transformer paper (Dehghani et al, 2019). Basically, there is a dense layer with sigmoid activation (ACT unit), that given the representation of each token at the current layer, predicts a score (the halting score). The model will halt at a given layer for a given token when the sum of all the halting scores for the token in the previous layer is bigger than $1 - \epsilon$. We could also do this per example, but the current experiments in the paper report the results for when ACT is applied per token.  Also, the objective of the model is augmented with the ACT loss, which is: 
<number of steps> + (1 - <final halting_score>). 

**Attention based router:**
For the attention based router, we use a standard attention layer (multi-head cross attention, in our experiments with just one head). Where the query is the example embedding at the current layer, and the key/values are the embeddings in the module embedding pool.

Reviewer comment (Novelty):
Hyper-UT is essentially an alternative parameterization of HyperNetworks, despite the dynmaic computation is placed in transformer blocks. I don't find it particularly different from HyperNetworks and its numerous variants, say [1,2,3,4,5].

Author response:
`“Hyper-UT is essentially an alternative parameterization of HyperNetworks...”`
Please see the shared response, the section on the concerns over novelty. Also, thank you very much for pointing out the additional references.

Reviewer comment (Experiments):
The experiment on C-PVR is not particularly convincing, as the task is a synethic one rather than a real problem. Could the authors find some other tasks that are more representative and yet realistic?

Author response:
`“The experiment on C-PVR is not particularly convincing, as the task is a synethic one rather than a real problem. Could the authors find some other tasks that are more representative and yet realistic?”`
That is a very good point. We think the experiments on the C-PVR task are still interesting, as it allows us to study the phenomena of interest more rigorously through controlled experiments.  The phenomena being how well the model can generalize to deal with examples that are more complex (requiring more sequential steps) than ones seen during training abstracting away from other factors that could influence the model. We also find it important to show this on more realistic use-cases. The main issue is that for real benchmarks it is really hard to come up with an automated way of annotating the complexity of the examples. Also, the space of potential solutions that the model can converge to that is not necessarily aligned with a defined complexity measure is much wider, making it very challenging to measure phenomena like generalization over complexity in these benchmarks. Nevertheless, we agree it is very important to build such benchmarks or find ways to show generalization over complexity on more realistic datasets.  We will discuss this point in the updated version of the paper. It is worth noting that we do have some experiment on ImageNet to show the generality of the suggested architecture.

Reviewer comment (Experiments):
In Section 4.2, a comparison with the T5 model trained from scratch is absent. To explore the influence of pre-training with language modeling on the generalization capacity, it is essential to include a comparison between the pre-trained T5 model and a T5 model trained from scratch specifically for the C-PVR task. This comparison will provide valuable insights into the impact of pre-training on the generalization.

Author response:
`“A comparison with the T5 model trained from scratch is absent.”`
We agree that the results with the T5 model are not directly comparable to transformers from scratch that we have in Figure 7. However, the main goal in these experiments was to show that the task remains challenging even for transformers  pre-trained on a language modeling objective. We will include experiments with the T5 architecture trained from scratch in the next version of the paper (the camera  ready version if the paper gets accepted). We speculate the results not to be very different than transformer encoders trained from scratch when there is no scratch-pad (Figure 7), since the decoding part would only have one decoding step. 
For the setting with scratch-pad we would expect better results (compared to the vanilla version) since the scratch-pad is basically augmenting the model with a mechanism for adaptive compute. While this effect has already been studied in other works, it would be interesting to see how scratch-pad or CoT helps in the context of the C-PVR task specifically (when there is no pretraining).

Reviewer comment (Presentation):
Comparing models with and without the scratch-pad in Figure 5 is challenging. It would be more helpful to combine Figure 5 (a) with Figure 5 (b), and Figure 5 (c) with Figure 5 (d).

Author response:
`Comparing models in Figure 5.`
We intentionally separated the plots to make them easier to read, but we agree that this makes the comparison a bit more difficult. We will update the plots and will incorporate this suggestion.

Reviewer comment (Writing):
What do the notations "32x3" and "128x2" mean in Table 2 and Table 3 in the appendix?

Author response:
`“What do the notations "32x3" and "128x2" mean in Table 2 and Table 3 in the appendix?"`
 Apologies for the typo. It is meant to be 32x128 and 128x256. It is <number of model embeddings> \times <size of the embeddings>.

Reviewer comment (Experiments):
Is it possible to apply multi-head attention to the attention-based router? How it will affect the generalization capacity compared to the single-head one?

Author response:
`“Is it possible to apply multi-head attention to the attention-based router? How it will affect the generalization capacity compared to the single-head one?”`
Intuitively, this might help for more complicated tasks, but the effect might be less significant for a task like C-PVR. There might also be some interaction between the number of heads and the size of the embeddings. We will include an ablation experiment for the number of heads in the next version of the paper (So far we have not seen an improvement in efficiency or accuracy when increasing number of heads on the C-PVR task). Nevertheless, we think it is important to continue research in this direction to find the optimal and most efficient way of incorporating modularity into adaptive models.