Paper: Discovering modular solutions that generalize compositionally (ICLR-2024)

Reviewer comment (Theory):
The choice of linear hypernetworks as the base modular model somewhat limits the interpretability of the compositionality.

Author response:
This is indeed an interesting point to consider as it is not clear if parameter module composition as done by hypernetworks generally leads to interpretable modules.
One part of the problem is certainly that interpreting individual modules boils down to interpreting the neural network they individually parameterize which is an open research problem.
In general, it is not clear if the compositional structure discovered by hypernetworks reflects the abstract compositional structure of the task (e.g. compositional goals or compositional preferences).
In the case of the modular teacher-student and the extended hyperteacher setting, our linear identification result shows, that it is possible to decode the ground-truth modules given the learned embeddings.
We show that this is indeed possible in practice in Figures 3C and D.
Future work is needed to understand under what conditions parameter modules reflect the potentially more abstract compositional structure in more complex settings.

Reviewer comment (Experiments):
In Sec 4.1, it is not described how compositional generalization is measured. Specifically, it is unclear whether the agent trains the task-specific coefficients on the new task without training the basis vectors.

Author response:
Yes, this is correct!
We were indeed missing a paragraph on the evaluation during test time which we have now added to Section 2.
In short, we use a support set to allow the model to infer the task-specific coefficients while keeping the basis vectors fixed and measure performance on an independent query set.

Reviewer comment (Theory):
In Sec 4.1, the statement "certain overparameterization is necessary for meta-learning to succeed" is unclear — it appears to contradict the theory, raises the question of why it is necessary, and seems inconsistent with previous results that used the “right” number of modules. Additionally, all results in Fig. 3F appear to show that more module overparameterization is consistently better, making it unclear how this connects to the theory or the remaining claims.

Author response:
Thank you for raising this important point.
We have now added a sentence clarifying this point in Section 3.3.
In short, while the compositional generalization *when achieving 0 training loss* can only be guaranteed for non over-parameterized students, optimization, i.e. achieving 0 loss, is itself difficult for non over-parameterized models.
Fortunately, the empirical results suggest that gradient descent seems to find the generalizing solutions even when the student is slightly over-parameterized, up to some extent.
Interestingly, the other conditions, and in particular ensuring connected support, seem to be much more important.

Reviewer comment (Experiments):
The training setting for Sec 4.2 is very unclear. It is not stated whether there is a teacher, what the training process is, why "accuracy" is used as the metric in a sequential decision-making problem, how many total tasks there are, whether the task proportions make the training data size adequate, whether accuracy correlates with actual performance on the task, whether the agent compositionally generalizes to new RL tasks, and how “OOD generalization” is measured in this setting.

Author response:
Thank you for raising this issue.
We have added the following description of the task to the updated version of the manuscript:

We now consider a setting where an agent has compositional preferences shared across tasks but the functional mapping to its action-value function for each specific task is unknown.
In particular, we study the compositional preference grid world inspired by [1] where an agent can obtain rewards by collecting objects with different colors.
For each task, up to $K=3$ out of $M=8$ preference modules are sampled and summed to obtain a hidden preference vector which specifies the reward the agent obtains when moving to an object of a specific color (see Figure 4A).
We test the ability for compositional generalization by presenting the agent with tasks of unseen combinations of preference modules.
The agent then has to infer the current preferences from observing optimal trajectories and is evaluated to predict the action-value function for a new environment configuration where agent and object locations are resampled.
For more details on the task construction see Appendix C3.

We measure the accuracy with respect to the ground truth optimal policy.
This is due to our particular meta-learning setup where the model has to learn a composable behavior policy from demonstrations of optimal trajectories.
We clarified this point in the updated version of the manuscript.

Since the accuracy is measured with respect to the ground truth optimal policy, it is a conservative measure of actual performance on the task.
That is there could be suboptimal policies that obtain similar reward to the optimal policy but would still result in a lower accuracy.

Thank you for asking for this clarification, which we have added to the updated version of the manuscript.
We measure compositional generalization in the two grid world environments by withholding 25\% of the possible compositions.
In the compositional preferences environment this amounts to withholding certain combinations of the preference modules, in the compositional goals environment this amounts to withholding certain goal compositions.
For example the agent might during training never encounter tasks with the goal of "performing interaction 2 on the purple object in goal quadrant 3 in the maze layout 1" since they are used for OOD evaluation.

Reviewer comment (Theory):
The phrasing in the first lines of page 4 suggests that assumptions i–iii are sufficient conditions but not necessary conditions for identification, which contrasts with the authors’ stated aim in the introduction to understand “what properties of the data generating process allow to discover modular solutions enabling compositional generalization.”

Author response:
Thank you for pointing out that this part was not well communicated in the initial version of the manuscript.
Please see our response to your related point above.

Reviewer comment (Theory):
If assumptions i–iii are only sufficient and not necessary, then the work offers an understanding of only some of the properties that allow discovering modular solutions that compositionally generalize, rather than a full characterization.

Author response:
Thank you for pointing out that this part was not well communicated in the initial version of the manuscript.
Please see our response to your related point above.

Reviewer comment (Theory):
While the authors do empirically examine the “cost” of violating those assumptions and the assumptions themselves are fairly intuitive, this limitation in scope should be acknowledged.

Author response:
Thank you for pointing out that this part was not well communicated in the initial version of the manuscript.
Please see our response to your related point above.

Reviewer comment (Theory):
Other modular solutions, such as neural module networks [2], have very intuitive explanations of what each module does and how it composes with other modules. Linear parameter combinations within a hypernet do not seem to have such a straightforward intuition.

Author response:
This is indeed an interesting point to consider as it is not clear if parameter module composition as done by hypernetworks generally leads to interpretable modules.
One part of the problem is certainly that interpreting individual modules boils down to interpreting the neural network they individually parameterize which is an open research problem.
In general, it is not clear if the compositional structure discovered by hypernetworks reflects the abstract compositional structure of the task (e.g. compositional goals or compositional preferences).
In the case of the modular teacher-student and the extended hyperteacher setting, our linear identification result shows, that it is possible to decode the ground-truth modules given the learned embeddings.
We show that this is indeed possible in practice in Figures 3C and D.
Future work is needed to understand under what conditions parameter modules reflect the potentially more abstract compositional structure in more complex settings.

Reviewer comment (Experiments):
The authors’ attempt at understanding the kinds of composition the model can learn is presented in the experiments of Sec 4.2, but since those experiments are not clearly explained, they do not provide much intuition.

Author response:
Thank you for this suggestion.
We have tried to better motivate and explain the two environments in this light in our updated version of the manuscript.

Reviewer comment (Writing):
The authors are encouraged to include an intuitive description of what the compositionality is in the Sec 4.2 tasks and how linear hypernets should capture that compositionality.

Author response:
Thank you for this suggestion.
We have tried to better motivate and explain the two environments in this light in our updated version of the manuscript.

Reviewer comment (Theory):
My largest concern is that the framework is very constrained, and some constraints, e.g., linearity, may be essential to deriving the results.

Author response:
We would like to emphasize that we study a general class of compositional problems that can be captured as parameterized functions whose parameters are combinations of parameter modules.
While we assume that parameter modules are linearly combined to make our theoretical derivation in the modular teacher-student setup tractable, we have carefully designed our experiments to understand how our results extend to complex situations where these assumptions are violated.
We have tried to clarify this point in the introductory paragraphs to the experiments.

Specifically, we train both a linear and a nonlinear hypernetwork and we consider two tasks - the compositional goals and compositional preferences grid worlds - where the data generating distribution has no known closed-form expression as a linear hypernetwork.
Despite this setting not being linear, our theoretical insights extend to this setting as we can show that both compositional and connected support are required to achieve compositional generalization in the linear and nonlinear hypernetwork. 

We would also like to highlight that the MLP parameterized by the linear hypernetwork is nonlinear and we prove our theoretical results both for the ReLU nonlinearty and a general class of smooth activation functions.

Reviewer comment (Evaluation):
It indicates that there may be difficulty when generalizing the result to more complex situations.

Author response:
We would like to emphasize that we study a general class of compositional problems that can be captured as parameterized functions whose parameters are combinations of parameter modules.
While we assume that parameter modules are linearly combined to make our theoretical derivation in the modular teacher-student setup tractable, we have carefully designed our experiments to understand how our results extend to complex situations where these assumptions are violated.
We have tried to clarify this point in the introductory paragraphs to the experiments.

Specifically, we train both a linear and a nonlinear hypernetwork and we consider two tasks - the compositional goals and compositional preferences grid worlds - where the data generating distribution has no known closed-form expression as a linear hypernetwork.
Despite this setting not being linear, our theoretical insights extend to this setting as we can show that both compositional and connected support are required to achieve compositional generalization in the linear and nonlinear hypernetwork. 

We would also like to highlight that the MLP parameterized by the linear hypernetwork is nonlinear and we prove our theoretical results both for the ReLU nonlinearty and a general class of smooth activation functions.

Reviewer comment (Theory):
Linear assumption in hypernetwork.

Author response:
We agree that there are many ground-truth nonlinear combinations of parts in the real world that are very interesting and important to consider.
We would like to clarify however that this is distinct from linearly combining the parameters of a (nonlinear) neural network as is done in the linear hypernetwork we consider in the theory.
Our experimental results in Figure 4C, 4E and 4F suggest that it is indeed possible to model such nonlinear combinations of parts as you describe (e.g. compositional goals) with linear parameter compositions.
In particular, note that the linear hypernetwork performs very similarly to the nonlinear hypernetwork in these experiments, both clearly outperforming the nonlinear monolithic baselines.

Reviewer comment (Experiments):
It uses two two-layer neural networks.

Author response:
Our theoretical analysis of the modular teacher-student setup focuses on two-layer networks, as we are not aware of identification results in the standard teacher-student setup for arbitrary depth which our modular teacher-student setting generalizes.

Moreover, our experiments use neural networks of varying depth beyond two layers - both for the learned models and for the teacher.

Reviewer comment (Theory):
It assumes knowing the correct (teacher) architecture.

Author response:
We would like to clarify that in the teacher-student analysis, our theory explicitly considers the case where the number of modules and hidden units is unmatched, showing that it is theoretically possible to construct counterexamples in these cases.

Moreover, we run a number of experiments on the effect of over-parameterization in the student when verifying the theory in Figure 2C (formerly 2D), in the hyperteacher setting in Figure 3F, in the compositional preferences environment in Figure A8 and in the compositional goals environment in Figure A9.
They show that the sensitivity to over-parameterization is less crucial than suggested by the theory, especially in the more complex environments and compositional generalization is possible despite a mismatch in the number of modules (and hidden units in case a teacher is present).
We further clarify this point when presenting our experimental results.

The analysis does not assume that the teacher architecture is known given that the irreducibility assumption is not violated.
The teacher can be any architecture that can be implemented by a MLP, including a convolutional architecture.
Other parts that specify the architecture in terms of width of the hidden dimension $h$ and the number of modules $M$ are part of the analysis.

Reviewer comment (Theory):
The theory assumes knowing the number of modules (M) and hidden units (h).

Author response:
We would like to clarify that in the teacher-student analysis, our theory explicitly considers the case where the number of modules and hidden units is unmatched, showing that it is theoretically possible to construct counterexamples in these cases.

Moreover, we run a number of experiments on the effect of over-parameterization in the student when verifying the theory in Figure 2C (formerly 2D), in the hyperteacher setting in Figure 3F, in the compositional preferences environment in Figure A8 and in the compositional goals environment in Figure A9.
They show that the sensitivity to over-parameterization is less crucial than suggested by the theory, especially in the more complex environments and compositional generalization is possible despite a mismatch in the number of modules (and hidden units in case a teacher is present).
We further clarify this point when presenting our experimental results.

Reviewer comment (Writing):
The paper is not very well written and is hard to follow. There are no simple examples explaining the problems the authors are trying to solve.

Author response:
We have revised the presentation of the paper in general and the theoretical exposition in particular. We hope this makes the paper easier to follow. In particular, section 3.2 which contains our main theoretical result has been revised. It now starts with a concrete example accompanied by an illustration in Figure 2B and states definitions and theorems in a way we hope is easier to parse.

Reviewer comment (Presentation):
There is no section that explicitly discusses related work.

Author response:
Thank you for making us aware of this potential misunderstanding.  
We have treated related work as part of our discussion section. We clarify this now by changing the section title to "Related work and Discussion". In addition, we added an "Extended related work" section in Appendix~E.

Reviewer comment (Theory):
Could you explain why the network modules have to be hypernetworks in your setting?

Author response:
Thank you for pointing out that this was not clearly communicated.  
We clarified the setting we consider in Section 2:

To endow our data generating distributions with compositional structure, for each task we recombine sparse subsets of up to $K$ components from a pool of $M$ modules.  
Each task can be fully described by a  latent code $z$ that specifies the combination of module parameters $(\Theta^{(m)})_{1 \leq m \leq M}$ which parameterizes a task-shared function $g$ with parameters

$\omega(z) = \sum^{M}_{m=1} z^{(m)} \Theta^{(m)}$.

Given this vantage point, hypernetworks are one possible choice to capture this class of compositional problems using neural networks.  
We believe they are a natural choice, as they can be seen as implementing functions with composable parameters using neural networks both to implement the function itself as well as to implement the parameterization of the function using module composition.  
The strong performance of the hypernetworks in our experiments on the compositional preferences and compositional goals grid worlds provide evidence, that this module composition in weight space is indeed an expressive architecture to capture varied compositional structure.