Adaptivity and Modularity for Efficient Generalization Over Task Complexity

Thank you for reviewing our paper. Please have a look at our shared response. And here are our comments/questions about some of your remarks about the weaknesses of the paper.

Task definition:

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).

Comparison with the state of the art:

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.

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.

Can this be applied to NLP Transformers (e.g., LLMs)? What's the reason of using ViT/image classification as the testbed? Of course, we just wanted to show that the idea generally applies and can be useful even in system-1 like tasks where there is not much need to a lot of sequential processing. Hence we chose image classification. We agree it would be very interesting to see the results of applying these to LLMs. We might even see more benefits on language modeling since, it sounds like a much more complex task. While we have deferred this to future work, we would like to point out to a concurrent work that applies a similar idea to NLP tasks (Tan et al, 2023).

[1] Shawn Tan, Yikang Shen, Zhenfang Chen, Aaron Courville, and Chuang Gan. Sparse universal transformer. arXiv preprint arXiv:2310.07096, 2023.

Thank you for the feedback. We have responded to your comments about novelty and ablation experiments in the shared response. Here are our comments on some of the remaining points:

“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.

“The motivation behind Hyper-UT is not clear to me....It seems that it is simply yet another hypernetwork.” As mentioned in the shared response, having an adaptive depth mechanism is a crucial part of this architecture. To reiterate the main point we are trying to make in this paper, we show how mechanism for adaptive depth and modularity can help transformers (1) generalize faster and better to examples of higher complexity (2) allocate compute more fairly. The particular design choice for the hyper module worked well empirically, and the precise overall scheme was intended to combine adaptive depth with modular reuse. The general intuition that leads to such architecture is that, we want a model that can allocate both type and amount of compute conditioned on the given example. From another point of view, we want to benefit from inductive biases of parameter sharing and recurrence in depth without sacrificing capacity.

“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.

Thank you for the thorough review and all the great questions and comments.

“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).

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.

“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>.

“Is Hyper-UT also parameter-efficient compared to other methods?” Intuitively we think, at larger scales, Hyper-UT can also be more parameter-efficient but in this paper our main focus has been on the flop efficiency and we have not investigated this aspect. In the current experimental results in this paper, Hyper-UT has actually more parameters compared to the vanilla transformer / vision transformer and certainly universal transformer with full parameter sharing but it does converge to use fewer flops per example (more specifically less flops for simpler example and more flops for more complex examples).

“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.