Discovering modular solutions that generalize compositionally

Interpretability of linear hypernetworks: Other modular solutions (such as the neural module networks of [2]) have very intuitive explanations of what each module does and how that composes with other modules. Linear parameter combinations within a hypernet don't seem to have such a straightforward intuition. The authors' attempt at understanding what kinds of composition the model can learn is in the experiments of Sec 4.2, but since those are not clearly explained, it's hard to get any intuition from them. I encourage the authors to include an intuitive description of what the compositionality is in those tasks and how linear hypernets should capture that compositionality.

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.

Throughout the abstract and much of the text, it's unclear what types of problems the authors are tackling. There's a mention of "demonstrations" and "action-value functions", which somewhat hints at a solution geared toward RL. Much later, it becomes clear that a "demonstration" is the output of the teacher network, and that action-value prediction is just one application of the linear hypernets. Up to this point [Section 3], it was still unclear to me what the authors meant by identification of modules. Is it that the student can determine which modules to use, or that it can find the right set of modules? Is this in a setting where the student knows the latent codes? These are later clarified, but it might be worth doing so earlier on

Thank you for raising these issues. We have made substantial efforts to improve the presentation of the paper and attempted to clarify these points earlier in the text.

It's odd that the authors combined a related work section with the discussion, but I think it works okay.

Indeed as you point out, 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 to Appendix~E.

[1] Barreto, André, et al. "Fast reinforcement learning with generalized policy updates." Proceedings of the National Academy of Sciences 117.48 (2020): 30079-30087.

Thank you for your comprehensive review and for the useful feedback that we incorporated in the revised version of our manuscript. We are grateful that you appreciated the combination of theoretical results and empirical evidence despite your concerns with the clarity of the exposition of the experimental setting. Following your suggestions, we have made substantial efforts to improve the presentation of the latter and address your individual questions point by point below. We hope this strengthens your assessment for acceptance of our work.

The experimental setting is unclear, especially the sequential decision-making portion, which makes it hard to assess its impact

We have substantially reworked the experimental section to clarify the experimental setting in our updated version of the manuscript. In addition, we provide detailed descriptions of the experimental setting in Appendix C.

It is unclear whether the three assumptions (compositional support, connected support, and no over-parameterization) are necessary conditions or just sufficient

Thank you for raising this issue. We have now added a sentence immediately after Theorem 1 clarifying this point. In short, under the same conditions with no additional assumptions, the three conditions of connected support, compositional support as well as no over-parameterization are not only sufficient but necessary conditions for the implication in Theorem 1 to hold. We illustrate this by constructing counterexamples to the implication when some of the three conditions are violated, which we detail in Appendix A2.

The choice of linear hypernetworks as the base modular model somewhat limits the intepretability of the compositionality

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.

There are very few theoretical treatments of compositionality in the context of neural nets in the literature. In particular, I am only aware of [1] (which the authors actually fail to mention).

Thank you for pointing us to this reference which we indeed missed. We have added it to the updated manuscript.

Section 4.1

How is compositional generalization measured? Does the agent train the task-specific coefficients on the new task without training the basis vectors

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.

"certain overparameterization is necessary for meta-learning to succeed" -- what does this mean? Doesn't this contradict the theory? Why is it necessary? Haven't all previous results used the "right" number of modules? Also, all results in Fig. 3F seem to show that more module overparameterization is consistently better. It isn't clear how this connects to the theory or the remaining results and claims.

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.

Thank you for your encouraging review bringing to attention the importance of the question on how to learn modular structure for compositional generalization. Below we aim to contextualize your concerns on the constraints of our framework replying to your questions point-by-point. We hope these clarifications will convince you to vote for acceptance of the paper, and we remain available for any further questions.

My largest concern is that the framework is very constrained, and some constraints, e.g., linearity, may be essential to deriving the results. It indicates that there may be difficulty when generalizing the result to more complex situations.

(1) Linear assumption in hypernetwork.

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.

(2) It uses two two-layer neural networks.

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.

(3) It assumes knowing the correct (teacher) architecture.

(4) The theory assumes knowing the number of modules (M) and hidden units (h).

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.

Thank you for taking the time to consider our responses. We are replying to your remaining questions point-by-point below. We hope these clarifications will convince you to vote for acceptance of the paper, and we remain available for any further questions.

(1) Linear assumption in hypernetwork. As far as I understand, theoretical derivation is still important for non-linear combinations. Though the derivation may not be easily tractable, there are many ground-truth non-linear combinations of parts in the real world.

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.

(3) It assumes knowing the correct (teacher) architecture. Does the theory assume that the teacher architecture (e.g., fully connected or convolutional) is known? If so, why is it reasonable? How about in human-computer interaction settings, where the teacher is a human?

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.

Also, the experiments are not in real and complex settings, so a strong theoretical result is important.

We would like to emphasize that our teacher-student analysis extends state-of-the-art teacher-student analyses, like for example [1,2], to the full multi-task case where task latent variables must be inferred. We believe it is a strong theoretical result in the context of the literature.

[1] Tian, Yuandong. "Student specialization in deep rectified networks with finite width and input dimension." International Conference on Machine Learning. PMLR, 2020.

[2] Simsek, Berfin, et al. "Geometry of the loss landscape in overparameterized neural networks: Symmetries and invariances." International Conference on Machine Learning. PMLR, 2021.

Thank you for your comments and for taking the time to review our work. We have spend considerable effort to improve the presentation of the paper in order to address your main point about the difficulty of understanding our initial submission. Below we address your individual suggestions and questions point by point. We hope these clarifications will convince you to vote for acceptance of the paper, and we remain available for any further questions.

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.

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.

There is no section that explicitly discusses related work.

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.

1. Could you rearrange the paper to have a background section explaining things like MAML, hypernetworks, etc, for readers who are not familiar with these concepts?

We have revised the methods section to make the exposition to hypernetworks and MAML easier to digest for readers not familiar with these concepts and added a missing reference to the corresponding appendix section D, where we provide a detailed description of all models.

2. Could you explain why the network modules have to be hypernetworks in your setting?

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.

Thank you for expressing appreciation for the significance of our work and your very constructive feedback. Based on your input, we have significantly revised the presentation of the theory in the main text as well as the appendix, and are overall confident the quality of the paper improved as a result. We address your questions point by point below, and do our best to clarify important miscommunications. Please do not hesitate to let us know if anything remains unclear.

Notation is used but not properly introduced [...]

Thank you for raising this issue. We have revised our notation by simplifying it where possible and ensuring that symbols are properly introduced. We hope this makes it easier to parse our theory and addresses your questions concerning the notations.

Definitions 3.1,3.2 and 3.3 are not clear

Following the revised notation, we have tried to improve the clarity of the definitions in Section 3.2. We refrain from restating them here to avoid cluttering our response and would like to kindly ask you to revisit our updated manuscript.

For example, where defining irreducibility, the fact that all rows of are pairwise different means that each hidden neuron will be activated for a different feature - and extremely important point for a work concerned with whether meta-learners can extract ground-truth features from a "metateacher". This is not stated. Similarly, how this interacts with nonlinearities on the hidden neurons is ignored and what it means for no columns of to be is also not mentioned. Why would a full column of be ?

Thank you for raising this point. To understand the purpose of the irreducibility condition, consider the standard (i.e. single task) teacher-student setting with two-layer neural networks. Note that there are many different two-layer neural networks that are functionally equivalent. When aiming at characterizing the exact set of neural networks that are functionally equivalent to a given teacher network, we are free to pick the teacher network to be any member of that equivalence class. A convenient choice is to choose it to be minimal, i.e. with the smallest possible number of hidden units. The definition of irreducibility is a necessary (and sufficient in the case of the smooth nonlinearity originally used in the main text) condition for such minimality: whenever two input weights are equal, one can merge the associated hidden neurons without changing the function implemented by the network. Likewise, if an output weight is 0, one can remove the associated hidden neuron. We have added a paragraph that explains this point at the beginning of Appendix A.1, and revised the definition of irreducibility to more clearly communicate the above point.

Compositional Support seems to just be saying that $U$ is a basis $\mathbb{R}^m$ of and so I am not certain of how this new terminology is necessary. Is compositional support a weaker case of having a basis as $U$ does not need to be the minimal spanning set? Would Definition 3.1 and 3.2 together then imply that $U$ is a basis?

Compositional support indeed simply ensures that all modules have been encountered at least once during training. We adopt the notion of "compositional support" from related work [1]. It is a useful definition for our theory as violating it will pose an obvious but important problem for compositional generalization and we can use this property as an experimental intervention to inform to what extent a learned model has discovered underlying compositional structure (c.f. Figures 3C, 4C and 4F).

Indeed as you correctly state, compositional support is a weaker case than having a basis of $\mathbb{R}^m$ as $Z$ (formerly $U$) does not need to be the minimal spanning set.

Definition 3.3 is just extremely difficult to parse in general and so is Theorem 3.1.

Thank you for raising this point. We have revised the formalism used to state Definition 3.3 and Theorem 3.1 in an attempt to make them easier to parse.

The figures, while visually neat and well done, are vague and their captions unhelpful. For example, where it is said "See Figure 2B for an illustration of a connected support" and then the caption does not explain what a connected support is or how the connected vs disconnected task families connects to the actual Definition 3.3. For example, in Figure 1B, the tiling of the $x,y$ space is not connected at all to the rest of the work beyond the notation of $p(\tau)$.

Thank you for pointing out the issues with our figures and corresponding captions. According to your suggestions, we have tried to improve their clarity. In particular, we give a detailed explanation of the example visualized in Figure 2A (formerly 2B) in the beginning of section 3.2, and have reworked Figure 1B to better aid its purpose of illustrating the modular teacher-student setup and how we can use it to investigate compositional generalization. It now shows a concrete toy example that illustrates how during training only a subset of module compositions of the teacher hypernetwork is used to construct the tasks and how this allows us to test for compositional generalization by presenting a novel composition.

For example, assuming that $\mathcal{P}_x$ has full support over the input space. I see how the learned weights of the linear hyper-network are compositional, in the sense that they operate similar to a set of basis vectors, which in the case of a linear mapping is also features for the network. However, this is then more similar to feature learning rather than exact modules. Or is the idea here that the features being learned which align with ground truth is the same as identifying a composable module?

In our setting, the linear hypernetwork takes a low dimensional task embedding $z$ and generates the weights of a base network which is a two-layer MLP. Each of these generated networks will have their own feature extractors over the input $x$, and the linear hypernetwork thus parametrizes a complex family of feature extractors by linearly combining the parameter modules (meta-features) through $z$.

In our teacher-student result, we show a characterization of what these modules and the distribution over $z$ must be, in order for the identification of these meta-features. Crucially, while we assume the input distribution $\mathcal{P}_x$ provided for each of the generated networks to have full support, we only require the support of the (hidden) task latent variables $\mathcal{P}_z$ to be compositional and connected for Theorem 1 (compositional generalization) to hold. In practice, this means we can show tasks from a subset of module combinations during training (in fact, a finite number of tasks linear in the number of modules), holding out a large number of combinations to test for out-of-distribution generalization as a measure of compositional generalization during evaluation.

How would this then tie in with disentanglement, which implies compositionality. While I would be open to such claims - the most obvious on to me here being that feature learning and module learning in the limited setting you study are the same thing - I think that argument needs to be made explicit. This would be a different take on modularity in general though and systematic generalisation which tends to focus more on how separate pieces are learned to be separate in spite of covariances and this makes the problem easier [1]. In your case it seems more like a claim that if the input space is sufficiently explored then the network will learn the ground truth features but just because this is the only way to learn the task (to learn a full-rank basis set) and is not in fact learning to identify or solve a smaller problem which is then composable.

Each parameter module defines a number of features in the teacher network that are composable in the parameter space of the teacher. Our result on linear identification in the student intuitively says that the learned student modules are a linear combination of the teacher modules. This means for each module in the teacher, there is a linear combination of student modules that will recover the teacher features in the student (up to permutation) on each task. These results are therefore somewhat complementary to the typical study of disentangelment in feature learning in the sense that each teacher module defines a whole set of features and we tackle the question whether a student is able to discover the teacher modules when only exposed to tasks that contain a subset of the possible module combinations. We hope that this explanation helped clarify your point. Please let us know if this did not fully answer your question.

The experiments of Section 4 seem to be more in line with the standard notions of compositionality [2], however due to the above issues grounding the theory to larger scale models of compositionality I also struggle to see how Section 4 fits in with the rest of the work, beyond just demonstrating that hyper-networks are better in compositional domains than ANIL and MAML. Is there a greater connection beyond this (I do think this result is important in its own right though)?

We have tried to improve the introductory paragraphs in the experimental section to make the connection to the rest of the paper more comprehensible in an attempt to clarify that the experiments of Section 4 consider a similar setup as the theory in Section 3 with the difference that

1. We learn from finite data, i.e. we no longer assume the input distribution $\mathcal{P}_\vx$ has full support
2. As a result we cannot consider the simplified objective of equation (2) but instead employ the bilevel optimization objective of equation (1), better suited for the many-shot setting, and computationally friendlier as it does not require a full optimization in the inner loop.
3. We move beyond the modular teacher-student setting, introducing a mismatch in the data generating process and the model architecture. In particular, we consider tasks for which it is a priori not clear how well neural networks can capture the underlying compositional structure.
4. We contrast the performance of our modular architectures with monolithic architectures to highlight the benefits of the former despite their less frequent use in practice.

Within this practical and more complex setting, we test predictions made by our theory and show that the central assumptions of compositional and connected support play a crucial role.

Could the authors please explain what Figure 2A is aiming to show with the permutation invariance? It is only referenced in condition (ii) but then not explained in the figure. I think I would benefit from understanding this point better.

Thank you for pointing out the missing explanation of Figure 2A in our initial submission. We have extended our explanation of this and other failure cases in Appendix A2 and moved the figure to this section. The caption for what is now Figure A1 now reads as follows:

Toy example of how correct identification of individual modules can lead to inconsistent neuron permutations preventing compositional generalization to unseen module compositions: Consider the simplified setting of a teacher and student network with two input neurons and three hidden neurons. Both the teacher and the student have $M=3$ modules. The upper right defines the teacher weights for each module and each neuron. For instance the weights denoted by B correspond to the weights connecting neuron 2 to the input for module 1. We now assume the student during training encounters three tasks. For each task exactly one of the teacher modules is used. Since in MLPs the ordering of neurons is permutation invariant, the student can perfectly match the teacher outputs, even when it uses a different ordering of the neurons. As a result, the student modules can perfectly fit all three tasks, despite permuting the neuron indices. For instance, neuron 2 in module 3 of the student contains the weights H whereas the corresponding neuron in the teacher contains the weights G. When we now present a new task during out-of-distribution evaluation that mixes all three modules in the teacher, the student is required to mix the weights of each neuron across modules. Since the neuron permutations in the student are inconsistent with those of the teacher, the student is unable to create the correct neuron weights and ends up unable to generalize to the OOD task.

[1] Wiedemer, Thaddäus, et al. "Compositional generalization from first principles." arXiv preprint arXiv:2307.05596 (2023).