Discovering modular solutions that generalize compositionally
Authors: Simon Schug
Seijin Kobayashi
Yassir Akram
Maciej Wolczyk
Alexandra Maria Proca
Johannes von Oswald
Razvan Pascanu
Joao Sacramento
Angelika Steger
Venue: ICLR 2024 poster (2024)
Keywords: modularity
compositionality
compositional generalization
teacher-student
meta-learning
hypernetworks
Primary area: transfer learning, meta learning, and lifelong learning
TL;DR: Meta-learning can discover modular solutions that generalize compositionally in a multi-task setting.

Abstract: Many complex tasks can be decomposed into simpler, independent parts. Discovering such underlying compositional structure has the potential to enable compositional generalization. Despite progress, our most powerful systems struggle to compose flexibly. It therefore seems natural to make models more modular to help capture the compositional nature of many tasks. However, it is unclear under which circumstances modular systems can discover hidden compositional structure. To shed light on this question, we study a teacher-student setting with a modular teacher where we have full control over the composition of ground truth modules. This allows us to relate the problem of compositional generalization to that of identification of the underlying modules. In particular we study modularity in hypernetworks representing a general class of multiplicative interactions. We show theoretically that identification up to linear transformation purely from demonstrations is possible without having to learn an exponential number of module combinations. We further demonstrate empirically that under the theoretically identified conditions, meta-learning from finite data can discover modular policies that generalize compositionally in a number of complex environments.

Published as a conference paper at ICLR 2024
DISCOVERING MODULAR SOLUTIONS THAT GENERALIZE COMPOSITIONALLY
Simon Schug" ETH Zurich
Seijin Kobayashi ETH Zurich
Yassir Akram
ETH Zurich
Maciej Wolczyk IDEAS NCBR
Alexandra Proca Imperial College London
Johannes von Oswald
ETH Zurich
Google Research
Razvan Pascanu' Google DeepMind
João Sacramento'
Angelika Steger
ETH Zurich
ETH Zurich
ABSTRACT
Many complex tasks can be decomposed into simpler, independent parts. Dis- covering such underlying compositional structure has the potential to enable com- positional generalization. Despite progress, our most powerful systems struggle to compose flexibly. It therefore seems natural to make models more modular help capture the compositional nature of many tasks. However, it is unclear under which circumstances modular systems can discover hidden compositional structure. To shed light on this question, study teacher-student setting with modular teacher where we have full control over the composition of ground truth modules. This allows us to relate the problem of compositional generalization that of identification of the underlying modules. In particular we study modular- ity in hypernetworks representing general class of multiplicative interactions. We show theoretically that identification up to linear transformation purely from demonstrations possible without having learn an exponential number of mod- ule combinations. We further demonstrate empirically that under the theoretically identified conditions, meta-learning from finite data can discover modular policies that generalize compositionally number of complex environments,
INTRODUCTION
Modularity pervades the organization of artificial and biological systems. It allows building com- plex systems from simpler units: objects are constructed from parts, sentences are compositions of words and complex behaviors can be decomposed into simpler skills. Discovering this underlying modular structure promises great benefits for learning. Once a finite set of primitives has been ac- quired, it can be recomposed in novel ways to quickly adapt to new situations offering the potential for compositional generalization. For instance, having learned to solve tasks that require pushing green boxes and tasks that require jumping over red boxes, an agent that has decomposed its skills appropriately, can flexibly solve tasks that require jumping over green boxes and pushing red boxes.
While monolithic architectures, which the whole network engaged for all tasks, are powerful, they lack a built-in mechanism to decompose learned knowledge into modules. This puts them a potential disadvantage as they might have to learn each of the exponentially many combinations explicitly even if their constituent modules were all previously encountered. For instance, it is unclear to what extent large language models can compositionally generalize beyond the training distribution (Srivastava et al. 2023) Press et al. 2023 Dzin et al. (023).
This raises the question whether an architecture built from independent modules that can be flexibly combined can more easily extract modular structure from family of tasks with underlying composi- tional structure. Typically modular architectures are expressive enough to collapse to the monolithic solution of modeling all task combinations encountered during training and often do so in practice
"Equal contribution; correspondeece to (sschug, sei jink
1 Shared senior authors
"ooe available
com http://////.o..aa//////aaaaaaaa/ http:///////..ooooooo/oo/oaaaaaa

---

Published as a conference paper at ICLR 2024
A Modular architecture
B Modular teacher-student
Train only on tasks based a subset of medule compositions.
Embeediing
mod th
moods ee
mod Pt
mod modilae
Latents
Parameter composition
Composed task
0.4
8(F:WM)
WL ()
/0.5
Hypernetwork iid:
0.6 0
Monolithic architecture
Target network
0.5- 0.6 8
((x:Ws)
Test for compositional generalization on unseen compositions.
Task-specific
11.4
Nonlinear features
Task-shared
0
0.4 - 0.3 (
((::
Figure 1: Diagram contrasting hypernetwork as modular architecture linear combination of features as a monolithic architecture. B Toy example illustrating the modular teacher-student setting where tasks are drawn from parameter module compositions of teacher hypernetwork.
(Shazeer et al. 2017 Kirsch al. 2018 Rosenbaum al. 2018 Mittal al. 2022). Identifying
the correct modules is further complicated by the fact that the learning system only ever encounters data generated from mixtures of modules without access to the ground truth decomposition.
Here we aim to illuminate the question what properties of the data generating process allow to dis- cover modular solutions enabling compositional generalization. We do within the framework of multi-task learning where a learner encounters input/output pairs on set of tasks composed of few underlying modules. As modular architecture. we consider hypernetworks (Ha et al.] 2017). general form of multiplicative interaction (Jayakumar et al.) (2020) which allow to express composi- tions in weight space used by many multi-task models (Ruvolo & Faton 2013 Perez et al. 2018
Dimitriadis et al. 2023 Rame et al. 2023). Our main contributions are as follows:
We setup a multi-task tcacher-student setting with a modular teacher instantiated as a linear hy- pernetwork that gives us precise control over the underlying modules and their composition. It allows us to cast the problem of compositional generalization as an identification problem.
We provide a new identifiability result in this setting showing that in the infinite data limit the modules of a perfect student with zero training loss are linear transformation of the teacher modules. Crucially, observing linear number of sparse module combinations sufficient for the student generalize to any combination of modules.
an extended teacher-student setting, we demonstrate empirically how meta-learning from finite data of sparse module combinations can discover modular solutions that generalize composition- ally under the theoretically identified conditions in modular but not monolithic architectures.
We show that our results translate to complex settings of learning action-value functions for com- positional preferences and learning of behavior policies for compositional goals.
METHODS
Compositionality 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 E MRM with ||=||o K, that specifies combination of module parameters (0(m) I which parameterizes task-shared function g with parameters
(2) Em=1 ((m)) Being exposed only to demonstrations of such tasks with their underly- ing modular structure hidden, the naive strategy for a learner would be learn every combination separately. As the number of modules grows, the quantity of possible combinations increases ex- ponentially, requiring more and more capacity. contrast, identifying underlying compositional structure not only affords more compact representation, it also potentially allows the learner to generalize to unseen compositions.
Modular and monolithic architectures. For learning system to be able to leverage the compo- sitional structure of a family of tasks, its model architecture needs to be able to efficiently represent the corresponding class of compositions. Here contrast modular and monolithic (non-modular)
2

---

Published as a conference paper at ICLR 2024
architectures (see Figure TA) and investigate under which circumstances they admit functionally modular solutions that reflect the compositional structure of the data generating distribution. To test models for compositional generalization, we hold-out number of module combinations when we sample tasks for training. At evaluation time we can then test the ability of all models to solve the held-out out-o--distribution (OOD) tasks.
As our modular architecture of choice we consider hypernetworks (Ha et al. 2017) that allow to express a wide class of multiplicative interactions (Jayakumar et al. 2020). A hypernetwork h is a neural network that generates the parameters w for target neural network g( w) which itself processes the task data. The hypernetwork takes as input a task-specific embedding vector & and is parameterized by the task-shared parameters 0, i.e. w = h(Ø; 0). We can think of h as a function that maps 6 to selection of flattened module parameters (e (g M- In the case of a linear hypernetwork, is a linear function of & resulting in a linear combination of the module parameters. In the case of nonlinear hypernetworks, h is separate neural network whose penultimate layer can be thought of as selecting among the module parameters.
As an example for a monolithic architecture that only supports less expressive form of composition, we consider the commonly used motif of combining a nonlinear feature encoder shared across tasks with task-specific linear decoder (c.g. Lee al. 2019; Bertinetto et al., 2019; Raghu et al. (020). Here we focus on ANIL |Raghu et al. 2020). i.c. fxx Φ, 0) = & g(æ; 0) where & is a task sseccccc readout weight matrix. In addition,
consider the widespread MAML architecture (Finn et al.)
2017) with fæx Φ 0) = g (z; x(0)) where the task-shared parameters 0 are used to initialize the task-specific parameters Φ. For additional details on the architectures please consider Appendix D
Meta-learning Even if a model supports flexible compositional representations, discovering the hidden latent structure of task distribution simply from observing demonstrations is diffieult. A common paradigm to extract structure shared across tasks is meta-learning which we will consider in the following (Finn et al., 2017: Zhao al., 2020). Formally, we consider a distribution Pa over task latent codes z which has compositional structure. A given task z defines a joint distribution over input æ and output y. from which two datasets are sampled, Dau = {I,, bi No and Da = {a/;, v 1=1 Nuasy We allow the model to infer the task latent code from the demonstrations in the support set Duupput by optimizing the task-specific parameters Φ, obtaining (=))), before querying the model on the query set Druery to update the task-shared parameters 0. The meta- bilevel optimization problem
learning problem can be formalized
min Eg~P, (L)= A: )
s.t. (=()) E arg min L(ø, 8: D
(I)
where L is a loss function; here we using the mean-squared error and the cross-entropy loss. To derive our theory, we assume the infinite data limit which allows us to consider a simpler multi-task learning problem with D= = Dquery Diuppon resulting in a single objective
min EzwP, min L(dx, 0: Dz)
(2)
+z
We optimize it in practice by running gradient descent on the task-specific variable dx until con- vergence for each gradient update of the task-shared parameters 0. In contrast to the bilevel op- timization, no second order gradients need to be computed. At inference, we consider new dis- tribution Pless for which the support can be disjoint with that of Pz, and report the expected loss Egape (mins. L(Øz, (:)))
3 THEORY
We seek to understand what properties of the data generating distribution allow a modular architec- ture to discover a functionally modular solution and thereby enable compositional generalization. The core idea is to extend the classic teacher-student problem to multi-task setting, where the teacher has a compositional representation. This allows us to relate the problem of compositional generalization to that of parameter identification in the student as a necessary and sufficient condi- tion. For both the teacher and the student we choose linear hypemnetorrs (Ha et al., 2017) which afford a natural interpretation in terms of module composition. Each task can then be described concisely given a task latent variable that conditions the teacher hypernetwork to generate demon- strations given the selected sparse, linear combination of modules. We show that in the infinite data
3

---

Published as a conference paper at ICLR 2024
A
2--P2
B
connected support Disconnesed support
Contiraous
Teacher
Mod0 In OWNE OWNERMIA wwr.etttoooo
(1.100
(6.1 1.400)
Module sligmment
Module aligmment
0.9-
0.00
Tosk
0.5
anPo
1
21?
Tast Surte
0.7-
0.85
y
(000 O
(0.0 2..0)
0.0
Student
Continauus Dieret Task distribution
18
4
18
32
1AA famvikes
Hidden everparaeeeiaaiion
Hidden Everparametizateon
Figure 2: A Schematic of the multi-task teacher-student setup. B Visualization of connected and disconnected task distributions in a teacher with M = 6, K 2. Without the central node support would be disconnected. C Module alignment between the student and teacher is high for both continuous and discrete task distributions with connected but not disconnected support. D Module alignment sensitive to overparameterization. Numbers denote the factor by which the student dimension is larger than the teacher. Error bars in C,D denote standard error of the mean for 3 seeds.
limit the student can identify the teacher modules up to linear transformation without access to the task latent variable the following conditions are satisfied:
(i) The training distribution has compositional support, i.e. all modules wish compose in the OOD task have been encountered at least in one training task (c.f. WWeddemer et al. 2023).
(ii) The training distribution has connected support, i.c. no subset of modules appears solely in isolation from the rest, otherwise we cannot guarantee modules can be consistently composed in the student.
(iii No overparameterizatton of the student wrt. the teacher. For instance, if the student has more modules available than there are tasks in the training distribution, it could use one module per task, fitting all of them perfectly but remaining ignorant of their shared structure.
3.1 MULKI-TASK TEACHER STUDENT SETUP
In standard teacher-student setting, a student network is trained to match the output of teacher given inputs x E I drawn from some distribution Pa- Here, we consider a multi-task counterpart to this setting (Figure A In addition to the input, we sample hidden task latent variable from some distribution PN independently from æ. In the following, we denote all quantities relating to the student with hat, e.g. Ô, and those relating to the teacher without.
For both the teacher and student architecture we use lank-cond hypernetorrss (c.f. Fig- ure [[A). Specifically, we consider a 2-layer multi-layer perceptron (MLP) with h hidden units and one output unit for which the first layer weights are generated by a linear hypernetwork as [xx; a, O =) a. (W,
where vi is an activation function, a is the task- readout vector, and θ (0(m)) ISmSM are the parameter modules parameterizing the hypernetwork. The first layer weight W generated by linearly combining the (0(m)))mms) with the task latent variable z (M()) i.c.
( z) E ((m) The teacher then produces task-conditioned demonstrations as in Equation The task thus has compositional structure as defined in Section 2
The student is also implemented as a linear hypernetwork with the hidden dimension potentially different from the teacher. Crucially, the student only has access to the input but does not observe the task latent variable z. We allow the student however to infer task-specific parameter of dimension M for each task given contextual information which will then condition the forward pass. We estimate such that it minimizes the objective detailed in Section given demonstrations from the teacher. The student then optimizes the multi-task objective stated in Equation 2 where given the task shared parameters 0 (Ô, a), and the task-specific parameters Φx 2, the task loss is defined as
1
L(2, (6, D.) E. [1][(x à, é, 2) - f(x; a, 0,
(4)
3.2 IDENTIFICATION & COMPOSITIONAL GENER ALIZATION
We will now show under what conditions a student, that perfectly fits the teacher on the training distribution, recovers the teacher's modules and achieves compositional generalization.
4

---

Published as a conference paper at ICLR 2024
Let us first study a motivating example for teacher with M 6 modules illustrated in Figure 2B to build some intuition. For that purpose we consider task families associated with binary masks that indicate which sparse subset of K 2 modules is present. For instance, the binary mask bi = (1, 1, O, O, 0, 0) indicates that Z1 is a task family consisting of latent vectors for which exactly the first and second entry are non-zero, e.g. = (0.6, 0.7, 0, 0, 0,
We now consider the case where during training the student observes tasks from four different task families specified by the binary masks 6 (1, 0, 0, 0) b2 = (0, 1, 0, 0,00 0) (o, O, O, 1, by = (0, 0, 0, O, 1,) For each task, we assume that the student sees an infinite number of samples covering the whole input space and learns to perfectly fit the teacher. After train- ing we ask the student to solve a new task specified by the task latent code 2 = (1. O, O, O, 0) be able to solve
.b3
Given that it has seen all modules present this new task, we might expect it. It turns out that this is not the case since the neuron permutations between modules indexed by (1,2, 2.3) and (4, 5, 6} are likely inconsistent as it has never encountered a task that connects these two subsets of modules (also see Appendix [A.3 for more detailed explanation). We can avoid this pathology by adding a fifth task family to our set of training tasks given by the binary mask bs = (0, 0, 1, 1,0.0) As we will show now, this leads compositional generalization and the student will perfectly solve any task including 2;. In fact, the student modules will be a linear combination of the teacher modules which we refer to as linear identification.
Let us now consider a general set of task families = (71, .2..... where each task family Zi E Z is non-empty set of R.M with an associated binary mask by- In the following require that span Zx) = {v O by v E IR }, where span (A) denotes the set of linear combinations of elements of A. With this definition at hand we can formulate sufficient conditions on the training distribution P= over task latent variables whose support contains the union of all task families Z.
Definition 3.1 (Compositional support). P= has compositional support if all dimensions of RM are covered by at least one of the binary masks by. i.e. Σ/ by has no entry that is O.
Intuitively having compositional support simply ensures that all modules have been encountered at least once. As we have seen in the example above this is not enough, leading to our second condition.
Definition 3.2 (Connected support). P= has connected support if there exists a path between any two task families, where two task families Zk, Zi are said to be connected if their associated binary masks share a non-zero element, i.c. by A b #
With these two definitions we can now state a simplified version of our main result showing that a perfect student will generalize compositionally. Formal versions of the following theorems well as their proofs can be found in Appendix A.4 and shown for both ReLU as well as a general class of smooth activation functions.
Theorem (Compositional generalization, informal). Assuming that Pz has compositional and connected support, Pa has full support in the input space and the student dimensions match that of the teacher, i.e. M M I n, = h, then under additional smoothness condition on Pa and non-degeneracy of ( a),, we have that
EzaPa min L(2, (6,0): =0
min I(2, = e R
L.e. the student optimizes (Ô, à) to fit the teacher on Pz, it achieves compositional generalization.
Furthermore, if the student achieves zero loss on a finite number, N = dim(spann (zx) + 1, i.ddd samples of tasks for each 7& from Pa it will generalize compositionally almost surely.
This result is noteworthy since it implies that fitting number of tasks linear in the number of modules suffices to generalize to all exponentially many module combinations.
The assumptions of compositional, connected support and h h are not only sufficient, but also necessary conditions for the implication to hold (all other things being equal), since we can construct counterexamples, detailed in Appendix A.3 when one of the conditions is violated. Furthermore, achieving compositional generalization under the above assumptions is equivalent to linear identi- fication of the teacher modules in the student. For any hidden neuron index i, denote by O the slice of θ of dimension n X M which maps the task latent variable z to W; i.e. w; = O,2.
5

---

Published as a conference paper at ICLR 2024
Theorem 2 (Linear identification, informal). Assuming Pa has full support in the input space,
M M S n, h =b=1 h, non-degeneracy of ④ and a, then ming T(2, (Q, D2) 0 az is equivalent to the existence of an invertible matrix F, permutation a and a sign flip E [-1, such that
On) & =
[1,
i.e. the learned student modules are a linear combination of the teacher modules.
3.3 EMPIRICAL VERIFICATION IN THE THEORETICAL SETTING
We empirically verify our theoretical findings by randomly initializing a teacher model following Section [.1] and measuring the degree of identification achieved by the student network after training until convergence on various task distributions. See Appendix CopendixC..Tjeeee for details on the experiments. We measure module alignment as min, m |s(O,, (,F) where s is the cosine similarity function, and F obtained by regressing on z. A module alignment of implies linear identification.
Identification requires connected support We construct task distributions with both connected and disconnected support. Figure [C shows that the student parameter modules have high align- ment both for a continuous (z E RT)) and discrete (z E {0, 1)m) task distribution when the task support is connected. Disconnecting the task support leads to a noticeable degradation.
Identification is sensitive to overparameterization. Given a task distribution with compositional and connected support, we investigate the effect of overparameterization on identification. Specifi- cally, we vary the ratio of the number of hidden units h and modules M of the student with respect to the teacher. While Theorem holds only for î = h and M = M, optimizing the loss well enough typically requires slight overparameterization. Figure shows that a slight overparameterization is indeed beneficial for linear identification. Yet, beyond h 8h and A 2.M, the identification no- ticeably degrades for discrete training tasks and cannot be alleviated by adjusting the weight decay (c.f. Table AT in Appendix [I. For a continuous task distribution, overparameeeriizaioo appears unproblematic, indicating that the learning dynamics introduce a beneficial inductive bias.
EXPERIMENTS
We now empirically investigate if modular architectures realized as hypernetworks can composition- ally generalize when meta-learning from finite data, assess whether highly expressive monolithic architecture can match this performance, and determine abstract latent space compositionality can be discovered despite unknown functional mapping to the task data. We further verify that both theoretically identified conditions compositional and connected task support are important for compositional generalization beyond the teacher-studen setting.
4.1 HYPERTEACHER
We generalize the reacher-studens setup from Section 3.1 to a teacher hypernetwork that modularly parameterizes multiple layers of target neural network and only show a limited amount of data to the student per task. To more faithfully reflect the symbolic nature of many real-world compositional tasks, each task is sparse and discrete combination of modules as exemplified in Figure 3A. See Appendix
C.2 for more details. As introduced in Section we contrast monolithic architectures modular architectures: the nonlinear and linear hypernetwork, see Appendix D for a detailed description of all models. The nonlinear hypernetwork specifically allows us to investigate the effects of architectural mismatch with the teacher network.
namely ANIL and MAML
Modular but not monolithic architectures compositionally generalize. To test the ability for compositional generalization, we hold out a fraction of the possible teacher module combinations during training. Only if we ensure that all teacher modules part of the out-of-distribution (OOD) set have been encountered individually as part of another combination during training does our training distribution have compositional support. Figure 3 shows that both the nonlinear and linear hypernetwork achieve high OOD accuracy while ANIL and MAML fail to do so. When the training support is composoooiii the performance of the hypernetworks drops below that of ANIL and MAML. Figure| 3E demonstrates that this holds as long as the number modules chosen for task combinations is K > 1. In the case where exactly one module is chosen per task, K 1, the

---

Published as a conference paper at ICLR 2024
A
B
D
Composttinnll
Compositioal
Connected task
dittrbuioon accursyy
Non-sompoitiona support
modu (89)
composts surport
Linor modul (P9)
support
c.
teacher
teacher modules
2.6-
0.5-
0.5-
S
C4-
Choose 1 to x modules
Task
hepermswerk
hypernetook
Iypereeworkk
hypernetent
typeretactk
Noriinear hyperntteerk
- Linaer hypar natwork
--- AMI
MAML
stribuuton accurscy
16
teacher modules
teacher moduless
Mo Me wwtrpprerrreeeeeeo www.proreeeoooooooooo
6.6-
6.--
0.2
1-80-K
cimeessio
dimension
Figure 3: A In the hyperteacher we pick between 1 to K of the M teacher modules, adding them in parameter space to create a task. ANIL and MAMI. fail to generalize to OOD tasks regardless of the support of the training distribution, while hypernetworks achieve good OOD accuracy only when the task support is compositional. C+D When the task distribution is compositional and connected be linearly decoded from the student task embeddings. E Hypernetworks
the teacher modules have high OOD accuracy for > but low OOD accuracy for = 1. F OOD performance of hypernetworks is sensitive overparameterrizatoo in both the hidden dimension and module dimension. Error bars in B-F denote the standard error of the mean over seeds.
condition of connected support is not satisfied and as predicted by our theory, both the nonlinear and linear hypernetwork perform poorly terms of compositional generalization despite the task support being compositional.
Modules can be linearly decoded in compositional solutions given connected task support. Given our theoretical result, we expect that a linear hypernetwork that achieves high OOD accuracy to recover the module parameters of the teacher up to linear transformation. To test this we train linear decoder to predict the ground truth task latent variables of a task given the learned embeddings of the student on validation set and then test linear decoding performance on the OOD set. Shown in Figure 3C.D, close to one in the case where the task support is compositional and connected. In the case where the support is non-compositional, the decoding performance is reduced reflecting the fact that there was no opportunity to learn about one of the modules. Despite having seen all modules during training, performance further drops significantly when the task support disconnected predicted by our theory. Surprisingly, these results extend to the nonlinear hypernetwork where there is an architectural mismatch between the teacher and the student and there no prior reason for the relationship between learned embeddings and ground-truth task latent variables to be linear.
find that in line with the theory, get high decodability with an R2 score
Compositional generalization is sensitive overparameterization. The theory in the discrete setting requires the student to have the same hidden and module dimension as the teacher. We test how sensitive the ability for compositional generalization in terms of OOD performance is to overparameterizzaioon on these two axes. In Figure 3 we observe that while certain overparam- eterization is necessary for the optimization succeed, performance starts decrease for larger overparameterrizatoo..
Meta-learning overfits to the number of modules entering the composition. Despite linear iden- tification as detailed above, we further observe that learners overfit to the particular number K of modules combined within a task, with OOD accuracy decreasing for larger than what was en- countered during training, e.g. see Table A6 Preliminary experiments suggest that this be partially alleviated by allowing for more gradient steps to perform task inference during evaluation, but does not present full remedy of the issue
7

---

Published as a conference paper at ICLR 2024
A
B
Connected tack
311
Norimnaar hyperstserkk
H
Distomedded uas support
hyper ryperneteoie
Oucei-dlttrbutonn
sccrtiyy detrrbutttoonnn
NAML
2.08
Combine g out 2f preferencee vectors prafersnces
0.12-
0.1
Prefecnte
Uneer
Korlinear hypernctwetk
Compostional support
Nursompaniioonl
Colors
tspernstwor'
support
E
0.-
Target object
Out Out-of diseriuuttooonn
ANE
Interaction
0.6-
0.6
Lincor hypernetwork
network
Nordl ncar byparnetwerk
Out
quadirant
04-
0.-
50k
100%
15ck
200k
Non SUREE osmtonatonnl
Compositinnl
Traininn Steps
Figure 4: In the compositional preference grid world the agent has modular preferences over colors and gets a reward corresponding to the current preference for the color of an object. B Dis- connecting the task support increases the OOD loss of hypernetworks. C Hypernetworks achieve better OOD loss than ANIL and MAML when the task support is compositional. D In the composi- tional goal grid world an agent has to walk Goals are composition of the maze, target object, target action and goal quadrant. E Hypernet- works achieve better OOD accuracy wrt to the optimal policy than ANIL and MAML. When holding out one of the goal quadrants the OOD accuracy decreases for hypernetworks more strongly than for ANIL and MAML. Error bars in B,C,F denote the standard error of the mean over seeds. 4.2 COMPOSITIONAL PREFERENCES
target object and perform the correct target action.
In the teacher-student setup we ensure by construction that the compositional representation of the data generating process is contained in the class of compositions the student can efficiently express. We now consider a setting where an agent has compositional preferences shared across tasks but the functional mapping its action-value function for each specific task is unknown. study the compositional preference grid world inspired by Barreto et al. (2020) where an agent can obtain rewards by collecting objects with different colors. For each task, up to K = 3 out of M preference modules are sampled and summed to obtain hidden preference vector which specifies the reward agent obtains when moving to an object a specific color (see Fi 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 see Appendix
particular,
Modular architectures learn composable action-value functions. Figure demonstrates that the linear and nonlinear hypernetoor achieve a lower mean-squared error loss predicting action- values on OOD tasks compared ANIL and MAML when the task support is compositional. The OOD loss significantly increases when one of the preference vectors is held-out from the training set, resulting in non-compositional support.
Encountering preference vectors in disjoint clusters interferes with compositional generaliza- tion. We can emulate disconnected task support similar our definition in the teacher-student setting by constructing the training tasks to contain subsets of preference vectors in disconnected clusters. We observe that this intervention has noticeable effect on the OOD loss achieved by both the linear and nonlinear hypernetwork, as shown in Figure 4, despite the number of observed combinations being comparable for the connected and disconnected support.
4.3 COMPOSITIONAL GOALS
Many real-world tasks have compositional goals raising the question whether modular architecture can discover behavior policies that decompose according to the latent goal structure. To this end, we consider the compositional goal grid world depicted in Figure 4D where agent has to navigate to a target object to employ a specific target action. The goals have compositional structure and we can test for compositional generalization by presenting novel compositions of goals. A goal is specified
8

---

Published as a conference paper at ICLR 2024
by four factors: maze layout, target object, target interaction and goal quadrant. For example possible goal would be "In maze I, perform interaction 3 on the purple object in quadrant 4". We measure the accuracy of the agent with respect to the optimal policy. For a more detailed task description see Appendix C.4
Modular architectures learn composable behavior policies. Figure demonstrates that both the linear and nonlinear hypernetwork achieve higher OOD accuracy of predicting the optimal behavior policy on held-out goal compositions than ANIL and MAML.
Modular architectures outperform monolithic ones by leveraging compositional structure. To construct non-compositional support in this task, we hold-out one of the goal quadrants from the training distribution. Figure 4F shows that this intervention has a more pronounced effect on the linear and nonlinear hypernetwork than on ANIL and MAML, indicating that generalization OOD tasks is due to the hypernetworks leveraging the compositional structure of the task.
5 RELATED WORK & DISCUSSION
The findings we present here are grounded in theoretical teacher-student analysis (Gardner & Derrida 1989) building on recent results that characterize the global minima of the loss of ML.s in the infinite data limit (Tian] 2020 Simsek et al. 2021). We extend the classical teacher-student setting to the multi-task setting where in addition to matching the teacher, the student needs to infer a hidden task latent variable and we instantiate the teacher as linear hypernetwork create compositional task structure. We show that parameter identification is both necessary and sufficient condition for compositional generalization and we prove new identifiability result that requires the task support of the training distribution to be compositional and connected two conditions we empirically verify to be crucial for compositional generalization.
Our results complement Lachapelle et al. (2023) who demonstrate that LI ---reuullrizationn enables disentangled identification in multi-task learning of sparse, linear feature combinations. Generally, there has been widespread interest in assessing compositional generalization (Lake & Baroni 2018
Ruis et al. 2020; Huppes et al.] 2020) as a crucial shortcoming of deep learning (Loula et al.
2018 Dessi & Baroni 2019 Keysers et al. 2020) across various subfields (Battaglia et al. 2018
Atzmon et al.] 2020; Montero et al. 2021 Liu et al. 2023). While with the advent of pretrained
large language models, compositional abilities have seemingly come into reach (Zhou et al. 2023
Orhan 2022 Furrer et al., 2021 Csordás et al. 2021) and many specialized architectures RUssin
al. 2019 Li et al. 2019 Gordon et al.) 2020 Andreas 2020 Liu et al.) 2020; Lake 2019]
Nye et al./ 22200] Kaiser & Sutskever (016) are outperformed by pretrained large language models
(Furrer et al.) 2021), I), remains an open question whether the ability for compositional generalization
extends beyond the training distribution (Srivastava et al. 2023 Press et al. 2023 Dziri et al., 2023)
Different from prior work we put particular emphasis on the multi-task setting and compositional-
ity in weight space a common and powerful motif in deep learning (Kumar & Daume III 2012) Ruvolo & Eaton, 2013] Perez et al., 2018 Jayakumar et al., 2020 Dimitriadis et al.) 2023). We show that linearly combining parameters can capture nonlinear latent structure raising an important question for future research on the interpretability of the resulting modules. We derive our theory in the infinite data regime while in practice resort meta-learning from finite samples. Indeed the sample complexity appears to scale unfavorably as the number of teacher modules increases (c.f. F..) Running meta-learning on this scale is computationally expensive and future work would ideally amortize the inference procedure when scaling to larger problem instances (Zhmogi- nov al. 2022). Moreover, while our identification result in the multi-task teacher-student setting indicates that overparameterization of the student can be detrimental for compositional generaliza- tion. the situation is less clear outside the teacher-student setting necessitating further investigation. Finally, we note that we characterized the solutions of the modular teacher-studen setting at equi- librium in the infinite data limit. Even when the modular solution constitutes a global optimum, we cannot provide any guarantees that it can be reached with gradient-based learning. Darvis et al. (2023) have made an important first step towards analyzing the learning dynamics of modular and monolithic architectures using deep linear networks in synthetic, linearly-separable datasets. Un- derstanding and analyzing the learning dynamics of the full modular sancher-student setting, as has previously been done for more restricted architectures (Seung et al. 1992; Saad & Solla 1995] Goldt et al., 2019). remains an important open question, the answer to which might help elucidate why even modular architectures often collapse monolithic solutions.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
We would like to thank Flavio Martinelli and Wulfram Gerstner for fruitful discussions and Nino Scherrer for valuable feedback on the manuscript. Simon Schug would like to kindly thank the TPU Research Cloud (TRC) program for providing access Cloud TPUs from Google. This research was supported by an Ambizione grant (PZ00P3 186027) from the Swiss National Science Founda- tion and an ETH Research Grant (ETH-23 21-1). Alexandra Proca is funded by the Imperial College London President's PhD Scholarship.
AUTHOR CONTRIBUTIONS
This paper was a collaborative effort. To do this fact better justice we give an idea of individual contributions in the following.
Simon Schug" Original idea, conceptualization of the project, development of theory, experimental design, implementation, running experiments, writing of the main text, writing of the appendix, literature review, creation of figures.
Seijin Kobayashi Original idea, conceptualization of the project, development theory, proving main theorems, experimental design, implementation, running experiments, helped writing of the main text, writing of the appendix, helped creation of figures.
Yassir Akram Helped develop theory, helped proving main theorems, provided feedback on the manuscript, helped writing of the appendix.
Maciej Wolezyk Worked on precursor to this project whose negative results formed the basis for the current project, provided feedback in regular project meetings, provided feedback on the manuscript.
Alexandra Proca Worked on a precursor to this project whose negative results formed the basis for the current project, provided feedback on the manuscript.
Johannes von Oswald Worked on precursor to this project whose negative results formed the basis for the current project, conceptualization of the project, provided feedback in regular project meetings, provided feedback on the manuscript.
Razvan Pascanu Senior project supervision, conceptualization of the project, provided feedback on the theory. provided feedback on experimental design, provided feedback in regular project meet- ings, provided feedback the manuscript.
João Sacramento Senior project supervision, original idea, conceptualization of the project, helped develop theory, experimental design, provided feedback in regular project meetings, provided feed- back on the manuscript.
Angelika Steger* Senior project supervision, conceptualization of the project, development of the- ory, helped proving main theorems, provided feedback in regular project meetings, provided detailed feedback the manuscript.
REPRODUCIBILITY STATEMENT
To ensure the reproducibility of our work, providing the code for all our experiments online https: / com/ smonsays
and complement the main complementing
text with experimental details in Appendix C For the theoretical results, our intuitive exposition the beginning of Section and empirical verification in Section (3. with detailed proofs in Section A and example failure cases build intuition in Section A.3
10

---

Published as a conference paper at ICLR 2024
REFERENCES
Ferran Alet, Tomas Lozano- Perez, and Leslie P. Kaelbling, Modular meta-lcarning, In Aude Bil- lard, Anca Dragan, Jan Peters, and Jun Morimoto (eds.). Proceedings of The 2nd Conference on Robot Learning. volume 87 of Proceedings of Machine Learning Research, pp. 558.868. PMLR, m www.p/ssssss11pmm attpaaleellll....... html
October 2018. http:tttst/oroceeediiiii htt.:/eddi...........ee...sseeee
Jacob Andreas. Good-Enough Compositional Data Augmentation. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 7556-7566, Online, 2020. Association for Computational Linguistics. doi: 10. (0.18663//1/2222.acll-a acl-main.676. URL
https:/ www. http:/wwwwwwww...o///// org/ hat///atoooooo/////00
main. 676
Yuval Atzmon, Felix Kreuk, Uri Shalit, and Gal Chechik. A causal view of compositional zero-shot recognition. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1462-1473, Cur- Associates, Inc., 2020. DBLAttpss//proceddng.coo neurigg........ss. cc/paper files/
Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacei, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzie, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Wojeiech Stokowiee, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL http: / com dw/pmindpiiinn
Dzmitry Bahdanau, Harm de Vries, Timothy J. O'Donnell, Shikhar Murty, Philippe Beau- doin, Yoshua Bengio, and Aaron Courville. CLOSURE: Assessing Systematic Generaliza- tion of CLEVR Models, October 2020. URL http: arXiv:1912.05783 [cs].
André Barreto, Shaobo Hou, Diana Borsa, David Silver, and Doina Precup. Fast reinforcement learning with generalized policy updates. Proceedings of the National Academy of Sciences, 117 300799 300709300797777777 1907370117. URL hottp::ts:///mm
December 2020. ISSN 0027 -424,, 1091-6490. doi: 10. 10.103//pms.. wwwpo///////oooo/// 10730000.3.pp./ 1907370117
Joost Bastings, Marco Baroni, Jason Weston, Kyunghyun Cho, and Douwe Kiela. Jump to bet- ter conclusions: SCAN both left and right. In Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 47-55, Brussels, Bel- gium, 2018. Association for Computational Linguistics. doi: 10.18653/
URL
http: //cc///cllwe web. org/ htth//nnhhooio logy/ / WW1 5407
Peter W. Battaglia, Jessica B. Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, Caglar Gulcehre, Francis Song. Andrew Ballard, Justin Gilmer, George Dahl, Ashish Vaswani, Kelsey Allen, Charles Nash, Victoria Langston, Chris Dyer, Nicolas Heess, Daan Wierstra, Pushmeet Kohli, Matt Botvinick, Oriol Vinyals, Yujia Li, and Razvan Pascanu. Relational inductive biases, deep learning, and graph networks, October 2018. URL 01261 arXiv:1806.0 806.0126.1111 [cs, stat].
Luca Bertinetto, Joao Henriques, Philip Torr, and Andrea Vedaldi. Meta-learning with differen- tiable flomed-fomm solvers. In International Conference on Learning Representations, 2019. URL
Lukas Biewald. Experiment Tracking with Weights and Biases, 2020. URL https: www.
wandb. com/
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal
Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-No Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
7 Ligithubbb. com cooooo////ag
11

---

Published as a conference paper at ICLR 2024
Róbert Csordás, Kazuki Irie, and Juergen Schmidhuber. The Devil is in the Detail: Simple Tricks Improve Systematic Generalization of Transformers. In Proceedings of the 2021 Conference Empirical Methods in Natural Language Processing, pp. 619-634, Online and Punta Cana, Dominican Republic, 2021, Association for Computational Linguistics. doi: 10. 22.18653/v12022/ emnlp-main.49, URL htt.:: /tt////alaaatthocoooo ogy wtg. 2021 emn 49
Verna Dankers, Elia Bruni, and Dieuwke Hupkes. The Paradox of the Compositionality of Natu- ral Language: A Neural Machine Translation Case Study. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 4154- 4175, Dublin, Ireland, 2022. Association for Computational Linguistics. doi: 10.18653/v1/2022. acl-long URI.tttp.a /ac Lanthl Sonthology.oo org//022 aal-locclllloooo long-
286
Roberto Dessi and Marco Baroni. CNNs found to jump around more skillfully than RNNs: Com-
positional Generalization in Seq2seq Convolutional Networks. In Proceedings of the 57th An- nual Meeting of the Association for Computational Linguistics, pp. 3919-3923, Florence, Italy,
URL July 2019. Association for Computational Linguistics. doi: 10.1663//1//1991133
Coline Devin, Daniel Geng, Pieter Abbeel, Trevor Darrell, and Sergey Levine. Plan Arithmetic: Compositional Plan Vectors for Multi-Task Control, August 2020. http:/ttt:///iiii..///
abs/1910. 14033 1rXiv:1910.14000 [cs, stat].
Nikolaos Dimitriadis, Pascal Frossard, and François Fleuret. Pareto Manifold Learning: Tack- ling multiple tasks via ensembles of single-task models. In Andreas Krause, Emma Brunskill, Kyunghyun Cho. Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.). Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 8015-8052. PMLR, July 2023. TRI.https:: /tt::/prroccc proceedoceeedeeeee ngs.m... mmrr. press/v202 / wwwiitiiiiiiiiiooo 223tmmmmllll.ll....
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D. Hwang. Soumya Sanyal, Sean Welleck, Xiang Ren, Allyson Ettinger, Zaid Harchaoui, and Yejin Choi. Faith and Fate: Limits of Trans- formers on Compositionality, June 2023. URL arXiv:2305. 18654 [es].
18654
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-Agnostic Meta-Learning for Fast Adap- tation of Deep Networks. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning. volume 70 of Proceedings of Machine Learn- ing Research, pp. 1126- 1135. PMLR, August 2017. URL bttps://proceed// press/ 21555/7000 / innl 7a html
.mlr.
Jerry A. Fodor and Zenon W. Pylyshyn. Connectionism and cognitive architecture: A crit- ical analysis. Cognition, 28(1-2)33-1, March 1988. ISSN 00100277. doi: 10.1016/
0010-027766399003--55 URL htttp://linkingghtn www.hnhhhh..eeeeiiieeeee elsevierr wwt///ttreeee////////eecc p11/001002718890033
E. Paxon Frady, Spencer Kent, Quinn Tran, Pentti Kanerva, Bruno A. Olshausen, and Friedrich T. Sommer. Learning and generalization of compositional representations of visual scenes, March
2023. URI/nttp:://rrwww.or/////a/ www.oorg//a......... /2300 13691 arXiv:2303. 13661 [cs].
Daniel Furrer, Mare van Zee, Nathan Scales, and Nathanael Schärli. Compositional Generalization in Semantic Parsing: Pre-training vs. Specialized Architectures, September 2021. URL http: arXiv:2007. .08970 [cs].
E. Gardner and B. Derrida. Three unfinished works on the optimal storage capacity of net- works. Journal of Physics A: Mathematical and General, 22(12):1983, June 1989. doi: 10. 0088/0305-444002211118 12/004
URL
doi
44700/
Badih Ghazi, Rina Panigrahy, and Joshua Wang. Recursive Sketches for Modular Deep Learn- ing. In Kamalika Chaudhuri and Ruslan Salakhutdinov l, Proceedings of the 36th Interna- tional Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Re- search, pp. 2211 -2220. PMLR, June 2019. URL ht https: http:////cccceeeeeeeeeee .mlr, ...ss/ v97 9a html
12

---

Published as a conference paper at ICLR 2024
Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dy- namics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché-Buc, E. Fox. and R. Gameet (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL htttps::/ppoceeedi wwwrrnurrrssee www.:eeww.........ee.eeee..eeee. files/ Papee
Jonathan Gordon, David Lapez-Paz, Marco Baroni, and Diane Bouchacourt. Permutation Equivari- ant Models for Compositional Generalization in Language. In International Conference on Learn- ing Representations, 2020. URRlshtttsss/speeeeeeet wet.eeeeweeeeeormmd///sssyeoee
David Ha, Andrew M. Dai, and Quoe Le HyperNetworks. In International Conference on Learn- ing Representations, 2017. URL /tttppss///ppeeeeeiee net/ httpa///eeeaeeeeeeeeeeeeeeeeerrr ftttt//ooommmiirkppceeeeee
Robert F Hadley. Systematicity in connectionist language learning. Mind & Language, 9(3):247- 272, 1994. Publisher: Wiley Online Library.
Dieuwke Hupkes, Verna Dankers, Mathijs Mul, and Elia Bruni. Compositionality decomppsee. how do neural networks generalise?, February 2020. Utl.hcpp:./aaxxiv.orr///b///// [cs, stat].
08351 arXiv:1908.0m 08351
Plotly Technologies Inc. Collaborative data science, 2015. URL https::/plot..2 Place: Montreal, QC Publisher: Plotly Technologies Inc.
Devon Jarvis, Richard Klein, Benjamin Rosman, and Andrew M. Saxe. On The Specialization of Neural Modules. In The Eleventh International Conference on Learning Representations, 2023.
URL htttps:://ppeeeeiiiiit
net wtt://orrmmmmmmhm1mm
Siddhant M. Jayakumar, Wojciech M. Czarnecki, Jacob Menick, Jonathan Schwarz, Jack Rae, Si- mon Osindero, Yee Whye Teh, Tim Harley, and Razvan Pascanu. Multiplicative Interactions and Where to Find Them. In International Conference on Learning Representations, 2020. URL
https::/cpenzeviie net/ httr//oormmm//mmm daa.umiid=ry1n66vv...
Justin Johnson, Bharath Hariharan, Laurens Van Der Maaten, Li Fei-Fei, C. Lawrence Zitnick, and Ross Girshick. CLEVR: A Diagnostic Dataset for Compositional Language and Elementary Vi- sual Reasoning. In 2017 JEEE Conference on Computer Vision and Patlern Recognition (CVPR), pp. 1988-1997, Honolulu, HI, July 2017. IEEE, ISBN 17788---336-0457-400
doi: 109VVR..
2017.215. URL htttps:://eeeeeee
leee org htt/ www.d/nnww/tooocccooomoo /00.c//mmmmet/9999988 8099698/
Lukasz Kaiser and Ilya Sutskever. Neural GPUs Learn Algorithms, March 2016. URL http:
/ wlarxivoorg/aee
08228 arXiv:1511 1.08228 [cs].
Daniel Keysers, Nathanael Schärli, Nathan Scales, Hylke Buisman, Daniel Furrer, Sergii Kashubin, Nikola Momchev, Danila Sinopalnikov, Lukasz Stafiniak, Tibor Tihon, Dmitry Tsarkov, Xiao Wang, Mare van Zee, and Olivier Bousquet. Measuring Compositional Generalization: A Com- prehensive Method on Realistic Data. In International Conference on Learning Representationss 2020. URL hatpp::pp///.......t...ttt..tt.t
Najoung Kim and Tal Linzen. COGS: A Compositional Generalization Challenge Based on Se- mantic Interpretation. In Proceedings of the 2020 Conference on Empirical Methods in Natural
Language Processing (EMNLP), pp. 9087-9105, Online, 2020. Association for Computational Linguistics. doi:
1/0663//12020.mmmmppp. anthology/2000 2020. emnlp www.pmmmaini.........m... 731
URL https:/ www. aclweb. org/
Hiroaki Kingetsu, Kenichi Kobayashi, and Taiji Suzuki. Neural Network Module Decomposi- tion and Recomposition, December 2021. URL http://awwii...gg//////./2....// 1.2. 13208 [cs].
Louis Kirsch, Julius Kunze, and David Barber. Modular Networks: Learning to Decompose Neu- ral Computation. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. and R. Garnett (eds.). Advances in Neural Information Processing Systems, volume 31. Cur- ran Associates, Inc., 2018 URL ht ht http:///ww....ee//oooeeeeeenoon.
htpp:///.....riiroooocpacaeeeere files/ Teaerr/0......iiaeeeeeeeeeeaaeea
13

---

Published as a conference paper at ICLR 2024
Abhishek Kumar and Hal Daume III. Learning Task Grouping and Overlap in Multi-task Learning, June 2012. URL http: http://axvvvv.r.oo/. awwaaaaa1o00 6417 arXiv:1206.6411 [cs, stat].
Sebastien Lachapelle, Tristan Deleu, Divyat Mahajan, loannis Mitliagkas, Yoshua Bengio, Si- mon Lacoste-Julien, and Quentin Bertrand. Synergies between Disentanglement and Sparsity: Generalization and Identifiability in Multi- Task Learning. In Andreas Krause, Emma Brun- skill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Pro- ceedings of the 40th International Conference on Machine Learning, volume 202 of Proceed- ings of Machine Learning Research, pp. 18171-18206. PMIR, July 2023. URL https: /pproceddings. mmr wte.. /tlpsss/v222/2 httpa////2222aacceeeeeeee22 html
Brenden M Lake. Compositional generalization through meta sequence-to-sequeeee learning. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates. Inc., 2019. URL https://proceedinss.eeooiii. ccppppppppppppppprrrrppppp files/paper/201s/ file/ flec///00e4t00055557766311100000
Brenden M. Lake and Marco Baroni. Generalization without systematicity; On the compositional skills of sequemce-to-ssqqeeec recurrent networks, June 2018. URL attp://arxiv.org/
abs/1711. 00350
arXiv: 1,1000030 [cs].
Kwonjoon Lee, Subhransu Maji, Avinash Ravichandran, and Stefano Soatto. Meta-Leaming with Differentiable Convex Optimization, April 2019. URL http::://aaa..io.oggaaa////////e 03758 arXiv:1904.037. [cs].
Yuampeng Li, Liang Zhao, Jianyu Wang, and Joel Hestness. Compositional Generalization for Prim- itive Substitutions. In Proceedings of the 2019 Conference on Empirical Methods in Natural Lan- guage Processing and the 9th International Joint Conference Natural Language Processing (EMNLP-JJCNLP), pp. 4293-4302 Hong Kong, China, November 2019. Association for Com- putational Linguistics. doi: 10. 10.18653/v1/D19-1 1438. URL httpss://cclaatholooy D19-1438 1438
Jiashuo Liu, Zheyan Shen, Yue He, Xingxuan Zhang, Renzhe Xu, Han Yu, and Peng Cui. Towards Of--st--istiibttionn Generalization: A Survey, July 2023. URL
org. abs/
2108.1362. arXiv:2108. 13624 [cs].
Qian Liu, Shengnan An, Jian-Guang Lou, Bei Chen, Zeqi Lin, Yan Gao, Bin Zhou, Nan- ning Zheng, and Dongmei Zhang. Compositional Generalization by Learning Analytical Ex- pressions. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 11416- 11427. Curran Associates, Inc., 2020. URL 2aper/2020/E5 fil /ttp:///e//aacceeeeeeeeeeee////.
wwwrrps....... cappr files/
João Loula, Marco Baroni, and Brenden Lake. Rearranging the Familiar: Testing Compositional Generalization in Recurrent Networks. In Proceedings of the 2018 EMNLP Workshop Black-
boxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 108-1 Brussels, Bel-
gium, 2018. Association for Computational Linguistics. doi: 10. 12.18653/11/W18-5442 URL
http::/aclwest org/ http///nthhtoootlaaeeee /W
Sarthak Mittal, Yoshua Bengio, and Guillaume Lajoie. Modular Architecture Enough? In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neu- ral Information Processing Systems, 2022. URL http:::ttt//t..t htttp::://ppeeeevieee wwt.:/eef//o/ioooooo
3- rx
Milton Llera Montero, Casimir JH Ludwig, Rui Ponte Costa, Gaurav Malhotra, and Jeffrey Bow- ers. The role of Disentanglement in Generalisation. In International Conference on Learning Representations, 2021. URL https: http:///pppeee.e/e////ppee www.://wwweee///oomreee//-aaaaae
3KUVy
Maxwell Nye, Armando Solar-Lezama, Josh Tenenbaum, and Brenden M Lake. Learn- ing Compositional Rules via Neural Program Synthesis. In H. Larochelle, M. Ran- zato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Infor- mation Processing Systems, volume 33, pp. 10832-10842. Curran Associates, Inc.,
14

---

Published as a conference paper at ICLR 2024
2020. URL https///acceed/....arrii....p.pp http:/ppperriiieeeaaaea/222/2/2 File/ TLLL:7/55aas5001355fffffeeeeeae2
A. Emin Orhan. Compositional generalization in semantic parsing with pretrained transformers, December 2022.
http://aaaxx...oooa////225555 att.abbs...... 211121119111111111 15101 arXiv:2109.15101 [cs].
Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. FiLM: Visual
Reasoning with General Conditioning Layer. Proceedings of the AAAI Conference Artificial Intelligence, 32(1), April 2018. ISSN 2374-3468, 2159-5399, doi: 10. 10.1609/auai.v32... h3tp://....i.....11.11 1661. wtw.tticee/viewwel1l1e
URL httpt:://o/off aaai. www//iinneee
Steven Phillips. Connectionism and the problem of systematicity. PhD Thesis, University of Queens- land, 1995.
Edoardo M. Ponti, Alessandro Sordoni, Yoshua Bengio, and Siva Reddy. Combining Modular Skills in Multitask Learning, March 2022. http:////aarxx wttp.//.w...//b//2.2233233 arXiv:2202.1391 [es].
Ofir Press, Muru Zhang, Sewon Min, Ludwig Schmidt, Noah A. Smith, and Mike Lewis. Mea- suring and Narrowing the Compositionality Gap in Language Models, May 2023. htthtp:
//arxiv. prexiv.org/abaee wwtp://rraabba//..oooo 03350 arXiv:2210.03350 [cs].
Aniruddh Raghu, Maithra Raghu, Samy Bengio, and Oriol Vinyals. Rapid Learning or Fea- ture Reuse? Towards Understanding the Effectiveness of MAML. In International Confer- ence on Learning Representations, 2020. URL hnttps://openreye: net / forum?id=
TkgMkCEtPR
Alexandre Rame, Guillaume Couairon, Corentin Dancette, Jean-Baptiste Gaya, Mustafa Shukor, Laure Soulier, and Matthieu Cord. Rewarded soups: towards Pareto-optimal alignment by interpolating weights fine-tuned on diverse rewards. In A. Oh, T. Neu- mann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 71095-71134. Curran Associates, Inc., 2023. URL https::/proceedd/nngg//// ngs neur. www..ps.ee...ee
https/erre..sssceppppeeeeeieeeee fatppppprrr.llssppppr//2233333 Confereoneeeeee BLcc42225BBB85333335555882222223 Paneeee--eeeee--ooeeeeeeeeee pdf
Clemens Rosenbaum, Tim Klinger, and Matthew Riemer. Routing Networks: Adaptive Selection of Non-Linear Functions for Multi-Task Learning. In International Conference Learning Repre- sentations, 2018. htthttss: http::/:////peeeeeee/ee hww.:eewneett/oomniiiiyyyyeooo
Laura Ruis, Jacob Andreas, Marco Baroni, Diane Bouchacourt, and Brenden M Lake. A Benchmark for Systematic Generalization in Grounded Language Understanding. In H. Larochelle, M. Ranzato, R. Hadsell, M. F, Balcan, and H. Lin (eds.), Advances Neu- ral Information Processing Systems, volume 33, pp. 19861-19872. Curran Associates, Inc., 2020. URL
htwp://www..............o..ooo.. http::sss.....iii...ccccoaaa.aaa http:/pppprrreeeppeeeeeeeee filesspapee//022/ File 2111:/5500B222c00112222222666663
David E Rumelhart and James L. McClelland. On learning the past tenses of English verbs. 1986. Publisher: Cambridge, MA: MIT Press.
Jake Russin, Jason Jo. Randall C. O'Reilly, and Yoshua Bengio. Compositional generalization in deep seq2seg model by separating syntax and semantics, May 2019, URL htt.:ttp//// //arxiv. org/abs/ 1944. 09708 MtXiv:1904.09700 [cs, stat].
Paul Ruvolo and Eric Eaton. ELLA: An Efficient Lifelong Learning Algorithm. In Sanjoy Das- gupta and David McAllester (eds.), Proceedings of the 30th International Conference on Ma- chine Learning, volume 28 of Proceedings of Machine Learning Research, pp. 507-515, At- lanta, Georgia, USA, June 2013, PMLR. URI /tRttttt::// ruvolo13. 1.1.... html
pecttps://procee https////ccceeeeegaaa m r.l...pressv2)
Issue: 1.
David Saad and Sara Solla. Multilayer Neural Networks. Dynamics of On-Line Gradient Descent Learning for D. Touretzky, M. C. Mozer, and M. Hasselmo (eds.), Advances in Neural Information Processing Systems, volume 8. MIT Press, 1995. URL /ttpss:/pprceeednne wtwpediinnss.siiiiieeeee 10ll///5SSScebba4444000011001110 h10.
15

---

Published as a conference paper at ICLR 2024
H. S. Seung, H. Sompolinsky, and N. Tishby. Statistical mechanics of learning from examples. Physical Review A, 45(8):6056-600) April 1992. ISSN 1050-2947, 1694-1622. doi: 10.
1103/hhysRev..55.605 URL https: /tttpa://iik.../ wttp://////...111//./../ 6056
Noam Shazeer, *Azalia Mirhoseini, *Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hin- ton, and Jeff Dean. Outrageously Large Neural Networks: The Sparsely- -aate Mixture-of- Experts Layer In International Conference on Learning Representations, 2017. URL https: //openreview.. net for www.ffoommmiddido BIckMDg11
Berfin Simsek, François Ged, Arthur Jacot, Francesco Spadaro, Clement Hongler, Wulfram Gerst- ner, and Johanni Brea. Geometry of the Loss Landscape in Overparameterized Neural Networks: Symmetries and Invariances. In Marina Meila and Tong Zhang (eds.). Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 9722-9732. PMLR, July 2021. URL http::///grceeeeninggss/ http::eiiiiinnmm........e. v139 1tt.mmsekkaaaammll..
Paul Smolensky. Connectionism, Constituency and the Language of Thought. In Barry M. Loewer and Georges Rey (eds.), Meaning in Mind: Fodor and His Critics. Blackwell, 1991.
Aarohi Srivastava, Abhinav Rastogi, Abhishek Rao, Abu Awal Md Shoeb, Abubakar Abid, Adam Fisch, Adam R. Brown, Adam Santoro, Aditya Gupta, Adrià Garriga-Alonoo, Agnieszka Kluska, Aitor Lewkowycz, Akshat Agarwal, Alethea Power, Alex Ray, Alex Warstadt, Alexander W. Ko- curek, Ali Safaya, Ali Tazarv, Alice Xiang, Alicia Parrish, Allen Nie, Aman Hussain, Amanda Askell, Amanda Dsouza, Ambrose Slone, Ameee Rahane, Anantharaman S. Iyer, Anders Johan Andreassen, Andrea Madotto, Andrea Santilli, Andreas Stuhlmüller,, Andrew M. Dai, Andrew La, Andrew Lampinen, Andy Zou, Angela Jiang, Angelica Chen, Anh Vuong, Animesh Gupta, Anna Gottardi, Antonio Norelli, Anu Venkatesh, Arash Gholamidavoodi, Arfa Tabassum, Arul Menezes, Arun Kirubarajan, Asher Mullokandov, Ashish Sabharwal, Austin Herrick, Avia Efrat, Aykut Erdem, Ayla Karakas, B. Ryan Roberts, Bao Sheng Loe, Barret Zoph, Bart lomiej Bo- janowski, Batuhan Özyurt, Behnam Hedayatnia, Behnam Neyshabur, Benjamin Inden, Benno Stein, Berk Ekmekci, Bill Yuchen Lin, Blake Howald, Bryan Orinion, Cameron Diao, Cameron Dour, Catherine Stinson, Cedrick Argueta, Cesar Ferri, Chandan Singh, Charles Rathkopf, Chen- lin Meng. Chitta Baral, Chiyu Wu, Chris Callison-Burrs Christopher Waites, Christian Voigt, Christopher D. Manning, Christopher Potts, Cindy Ramirez, Clara F. Rivera, Clemencia Siro, Colin Raffel, Courtney Ashcraft, Cristina Garbacea, Damien Sileo, Dan Garrette, Dan Hendryeks, Dan Kilman, Dan Roth, C. Daniel Freeman, Daniel Khashabi, Daniel Levy, Daniel Moseguí González, Danielle Perszyk, Danny Hernandez, Danqi Chen, Daphne Ippolito. Dar Gilboa, David Dohan, David Drakard, David Jurgens. Debajyoti Datta, Deep Ganguli, Denis Emelin, Denis Kleyko, Deniz Yuret, Derek Chen, Derek Tam. Dieuwke Hupkes, Diganta Misra, Dilyar Buzan, Dimitri Coelho Mollo, Diyi Yang. Dong-Ho Lee, Dylan Schrader, Ekaterina Shutova, Ekin Do- gus Cubuk, Elad Segal, Eleanor Hagerman, Elizabeth Barnes, Elizabeth Donoway, Ellie Pavlick, Emanuele Rodolà, Emma Lam, Eric Chu, Eric Tang. Erkut Erdem, Ernie Chang, Fthan A. Chi, Ethan Dyer. Ethan Jerzak, Ethan Kim, Eunice Engefu Manyasi, Evgenii Zheltonozhsküi, Fanyue Xia, Fatemeh Siar, Fernando Martinez-Plumed Francesca Happé, Francois Chollet, Frieda Rong, Gaurav Mishra, Genta Indra Winata, Gerard de Melo, Germán Kruszewski, Gi- ambattista Parascandolo, Giorgio Mariani, Gloria Xinyue Wang. Gonzalo Jaimovitch-Lopez, Gregor Betz, Guy Gur-Ari, Hana Galijasevic. Hannah Kim, Hannah Rashkin, Hannanch Ha- jishirzi, Harsh Mehta, Hayden Bogar, Henry Francis Anthony Shevlin, Hinrich Schuetze, Hi- romu Yakura, Hongming Zhang, Hugh Mee Wong, Ian Ng. Isaac Noble, Jaap Jumelet, Jack Geissinger, Jackson Kernion. Jacob Hilton, Jachoon Lee, Jaime Fernández Fisac, James B. Si- mon, James Koppel, James Zheng, James Zou, Jan Kocon, Jana Thompson, Janelle Wingfield, Jared Kaplan, Jarema Radom, Jascha Sohl-Dickstein, Jason Phang. Jason Wei, Jason Yosinski, Jekaterina Novikova, Jelle Bosscher, Jennifer Marsh, Jeremy Kim, Jeroen Taal, Jesse Engel, Je- sujoba Alabi, Jiacheng Xu, Jiaming Song. Jillian Tang, Joan Waweru, John Burden, John Miller, John U. Balis, Jonathan Batchelder, Jonathan Berant, Jörg Frohberg, Jos Rozen, Jose Hernandez- Orallo, Joseph Boudeman, Joseph Guerr, Joseph Jones, Joshua B. Tenenbaum, Joshua S. Rule, Joyce Chua, Kamil Kanclerz, Karen Livescu, Karl Krauth, Karthik Gopalakrishnan, Katerina Ignatyeva, Katja Markert, Kaustubh Dhole, Kevin Gimpel, Kevin Omondi, Kory Wallace Math- ewson, Kristen Chiafullo, Ksenia Shkaruta, Kumar Shridhar, Kyle McDonell, Kyle Richardson,
16

---

Published as a conference paper at ICLR 2024
Laria Reynolds, Leo Gao, Li Zhang, Liam Dugan, Lianhui Qin, Lidia Contreras-Ochande Louis- Philippe Morency, Luca Moschella, Lucas Lam, Lucy Noble, Ludwig Schmidt, Luheng He, Luis Oliveros-Colón, Luke Metz, Lütfi Kerem Senel, Maarten Bosma, Maarten Sap, Maartje Ter Ho- eve, Maheen Farooqi, Manaal Faruqui, Mantas Mazeika, Marco Baturan, Marco Marelli, Marco Maru, Maria Jose Ramirez- Quiitaaa,, Marie Tolkiehn, Mario Giulianelli, Martha Lewis, Martin Potthast, Matthew L. Leavitt, Matthias Hagen, Mátýás Schubert, Medina Orduna Baitemirova, Melody Arnaud, Melvin McElrath, Michael Andrew Yee, Michael Cohen, Michael Gu, Michael Ivanitskiy, Michael Starritt, Michael Strube, Michal Swedrowski, Michele Bevilacqua, Michi- hiro Yasunaga, Mihir Kale, Mike Cain, Mimee Xu, Mirac Suzgun, Mitch Walker, Mo Ti- wari, Mohit Bansal, Moin Aminnaseri, Mor Geva, Mozhdeh Gheini, Mukund Varma T. Nanyun Peng. Nathan Andrew Chi, Nayeon Lee, Neta Gur-Ari Krakover, Nicholas Cameron, Nicholas Roberts, Nick Doiron, Nicole Martinez, Nikita Nangia, Niklas Deckers, Niklas Muennighoff, Ni- tish Shirish Keskar, Niveditha S. Iyer, Noah Constant, Noah Fiedel, Nuan Wen, Oliver Zhang, Omar Agha, Omar Elbaghdadi, Omer Levy, Owain Evans, Pablo Antonio Moreno Casares, Parth Doshi, Pascale Fung, Paul Pu Liang, Paul Vicol, Pegah Alipoormolabshhi, Peiyuan Liao, Percy Liang, Peter W. Chang, Peter Eckersley, Phu Mon Htut, Pinyu Hwang, Piotr Milkowski, Piyush Patil, Pouya Pezeshkpour, Priti Oli, Qiaozhu Mei, Qing Lyu, Qinlang Chen, Rabin Banjade, Rachel Etta Rudolph, Raefer Gabriel, Rahel Habacker, Ramon Risco, Raphaël Millière, Rhythm Garg, Richard Barnes, Rif A. Saurous, Riku Arakawa, Robbe Raymaekers, Robert Frank, Rohan Sikand, Roman Novak, Roman Sitelew, Ronan Le Bras, Rosanne Liu. Rowan Jacobs, Rui Zhang, Russ Salakhutdinov, Ryan Andrew Chi, Seungjae Ryan Lee, Ryan Stovall, Ryan Techan, Rylan Yang. Sahib Singh, Saif M. Mohammad, Sajant Anand, Sam Dillavou, Sam Shleifer, Sam Wise- man, Samuel Gruetter, Samuel R. Bowman, Samuel Stern Schoenholz, Sanghyun Han, Sanjeev Kwatra, Sarah A. Rous, Sarik Ghazarian, Sayan Ghosh, Sean Casey, Sebastian Bischoff, Sebas- tian Gehrmann, Sebastian Schuster, Sepideh Sadeghi, Shadi Hamdan, Sharon Zhou, Shashank Srivastava, Sherry Shi, Shikhar Singh, Shima Asaadi, Shixiang Shane Gu, Shubh Pachchigar, Shubham Toshniwal, Shyam Upadhyay, Shyamolima Shammie Debnath, Siamak Shakeri, Simon Thormeyer, Simone Melzi, Siva Reddy, Sneha Priscilla Makini, Soo-Hwan Lee, Spencer Torene, Sriharsha Hatwar, Stanislas Dehaene, Stefan Divie, Stefano Ermon, Stella Biderman, Stephanie Lin, Stephen Prasad, Steven Piantadosi, Stuart Shieher, Summer Misherghi, Svetlana Kiritchenko, Swaroop Mishra, Tal Linzen, Tal Sehuster, Tao Li. Tao Yu, Tariq Ali, Tatsunori Hashimoto, Te-Lin Wu, Théo Desbordes, Theodore Rothschild, Thomas Phan, Tianle Wang, Tiberius Nkinyili, Timo Schick, Timofei Kornev, Titus Tunduny, Tobias Gerstenberg, Trenton Chang, Trishala Neeraj, Tushar Khot, Tyler Shultz, Uri Shaham. Vedant Misra, Vera Demberg, Victoria Nyamai, Vikas Raunak, Vinay Venkatesh Ramasesh, vinay uday prabhu, Vishakh Padmakumar, Vivek Sriku- mar, William Fedus, William Saunders, William Zhang. Wout Vossen, Xiang Ren, Xiaoyu Tong, Xinran Zhao, Xinyi Wu, Xudong Shen, Yadollah Yaghoobzadeh, Yair Lakretz, Yangqiu Song, Yasaman Bahri, Yejin Choi. Yichi Yang. Yiding Hao, Yifu Chen, Yonatan Belinkov, Yu Hou, Yufang Hou, Yuntao Bai, Zachary Seid, Zhuoye Zhao, Zijian Wang. Zijie J. Wang. Zirui Wang, and Ziyi Wu. Beyond the Imitation Game: Quantifying and extrapolating the capabilities of language models. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL https: httpss///peeeeeeeeeeemm .nt/ www.forrdddoooooo dirrumiidduyTLSBBee.
Yuandong Tian. Student Specialization in Deep Rectified Networks With Finite Width and Input Di- mension. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Confer- ence on Machine Learning. volume 119 of Proceedings of Machine Learning Research, pp. 9470- 9480. PMLR, July 2020. URL ht.ml
mml. wttp::m.w.......1./eeeee.... wwwppp//v1l/tiaan.......
Ivan I. Vankov and Jeffrey S. Bowers. Training neural networks to encode symbols enables com- binatorial generalization. Philosophical Transactions of the Royal Society B: Biologicul Sci- ences, 375(1791) February 2020. ISSN 0962-8436, 1471-2970. doi: 10 socclssooiiie wow.ooasssstyppollllll.ol.. weepubbisiinngg Sasing.nngg......... owg. /tg.g//oooii.g.ooooooooooo/oooi. /10. 1at.10098/rtbbb
2019.0309. URL
2019.0300
Thaddaus Wiedemer, Prasanna Mayilvahanan, Matthias Bethge, and Wieland Brendel. Compo- sitional Generalization from First Principles, July 2023. URL 2307. 05596 3Xiv:2307.05500 [cs, stat].
abs/
Zhenlin Xu, Mare Niethammer, and Colin A Raffel. Compositional Generalization in Unsupervised Compositional Representation Leaming: A Study on Disentanglement and Emergent Language.
17

---

Published as a conference paper at ICLR 2024
In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 25074-25087. Curran Associates, Inc., 2022. URL https::/pptoeeeeeee http://wwweeddd.....iiieiisseeee cc cc paper files/paper/2022/ fi 2lee/f
Dominic Zhao, Seijin Kobayashi, João Sacramento, and Johannes von Oswald. Meta-learning via hypemetworks. 4th Workshop on Meta- Meta-Learning at NeurIPS 2020 (MetaLearn 2020). NeurIPS, 2020.
Linfeng Zhao, Lingzhi Kong, Robin Walters, and Lawson L.S. Wong. Toward Compositional Gener- alization Object-Oriented World Modeling. Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.). Proceedings of the 39th International Conference on Machine Learning. volume 162 of Proceedings of Machine Learning Research, pp. 26841-26864. PMLR, July 2022. URL https: /ttpase//prcceedd www.nng....m.nn..m..mm..mm.m.mm. zhao22b. l
mm0.mmmrr.. www.r..pee//vlll. 62/
Andrey Zhmoginov, Mark Sandler, and Max Vladymyrov. HyperTransformer: Model Generation for Supervised and Semi-Suuprvvsed Few-Shot Learning, July 2022. URL httpp:/arxxxv
org/abs/ 2201. 04182 82 arXiv:2201 .04882 [es].
Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schu- urmans, Claire Cui, Olivier Bousquet, Quoe V. Le, and Ed H. Chi. Least-to-Most Prompting Enables Complex Reasoning in Large Language Models. In The Eleventh International Confer- ence on Learning Representations, 2023. URL https::tp///........ / ew. net. fort/ffrum?i 1= WZH 7999g
18

---

Published as a conference paper at ICLR 2024
Appendix
Table of Contents
Theoretical result
20
A.1
Single task parameter identification for ReLU MLP
20
A.2
Multi task parameter identification
25
A.3
Example failure cases to build intuition
28
A.4
Compositional generalization
30
Additional experiments
36
Multi-task teacher-student
36
B.2
Hypertcacher
36
B.3
Compositional preferences
38
B.4
Compositional goals
41
C
Experimental details
42
Multi-task seacher-student setup
42
Hyperteacher
44
Compositional preferences
46
C.4
Compositional goals
48
D Meta-learning models
49
D. MAML
49
D.2 ANIL
50
D.3 Linear hypemetwork)
50
D.4 Nonlinear hypernetwork
E Extended related work
51
F Addtional details
52
F Compute resources
52
Software and libraries
52
19

---

Published as a conference paper at ICLR 2024
A THEORETICAL RESULT
Before presenting the theoretical result for the multi-task teacher-student setting, first focus on the standard teacher student setting, which can seen as a special case of our multi-task setting where there exists only single task.
We will use the notation introduced Section 3.1 to describe network, in particular for any hidden neuron index i, we denote by w, the i-th row of the generated weight W (O,,2 x). θ the slice of θ of dimension n M which maps the task latent variable x to w,, i.c. w; = θ, and a, the i-th column of a.
A.1 SINGLE TASK PARAMETER IDENTIFICATION FOR RELU MLP
This section presents rigidity results that allows us, under the right conditions, to identify the learned parameters of student MLP learning imitate the output of teacher. Here the student may have the same or even a wider architecture. To our knowledge this the first such results for ReLU two-layer MLP.
There are clearly many functionally equivalent ways to implement a neural network. In a first step thus aim at characterizing the exact set of neural networks that functionally equivalent to to given teacher network. In order to get a nice characterization we require the teacher to be minimal. E.g., if the read-out value of hidden unit zero, this unit has influence on the output and can delete it. Formally, we assume that the teacher is irreducible:
Definition A.1. Irreducibility: The weights (W, a) of two-layer MLP f are called irreducible when there exists no two-layer MLP f with strictly less hidden units such that they are functionally equivalent, i.e. Vz E R2, f(x) f(x).
Clearly, we have the following necessary conditions for irreducibility:
output weight a is
no input weight vector ww is
no two input weight vectors wk; wi are identical
For a special class of smooth activation functions, vé E C and (0())(0) & 0 for infinitely many even and odd values of n, (imsee et al. (2021) show that the above conditions are sufficient for network to be irreducible.
For the ReLU nonlinearity, the above conditions are no longer sufficient for a teacher to be irre- ducible. As an illustration, if pair of weights wx, wy are positively colinear, then the 2 corre- sponding hidden neurons can be fused into a single hidden neuron without funetionally changing the network. See section A.I.2 for more details.
simple sufficient condition for the irreducibility input weights we that are colinear, and no output weight that is 0. While this is not neces- sary condition, we will restrict our identification results to such teacher networks for simplicity, ef. Section A.1.2 for the general case.
ReLU MLP is that the teacher has no two
A.1.1 IDENTIFICATION WITH RELU MIP
We now state the main theorem of this subsection. We will first treat the theorem for the special case that teacher and student network have the same number of hidden units, i.e. h = h. Subsequently, we will then state the theorem for 2 h and illustrate how to extend the proof. From here on, whenever we talk about measure, we refer to the Lebesgue measure.
Theorem 3. Assume teacher and student two-layer network parameterized by resp. (W, a) and (W,6 a). of hidden dimensions h h and single output unit such that
EnpP, [T(xx)] O.
(5)
20

---

Published as a conference paper at ICLR 2024
where
L(x) - - la (WW)) - -à aaT
(6)
and up is the ReLU nonlinearity.
Pz has full support in the input space and
(W, a) is such that
no output weight ag is O
no two input weights Wk, are colinear
then there exists a partition S1, S2 of (1..h] and permutation a of the neurons in the hidden layer such that we have
Vi, (1) > O,
Vi S1, (1) (ai a,
(8)
Vi S2, (i) ((i)) a,
Σ a,W; Σ 0;
(10)
Moreover, any two-layer teacher and student networks that satisfy equation equation satisfies equation 5
T0 also
The idea of the proof consists in noticing that 0 loss means that a ReLU MLP constructed by concate- nating the teacher and student hidden neurons need be functionally 0 everywhere. By grouping together the teacher and student weight vectors that are colinear with each other, we then show that the sum of each of hidden neurons belonging to the same group need to be a. Finally, because no two teacher weight vectors are colinear, we show that necessarily student weight needs be colinear with each teacher weight. The permutation a consists this mapping. and the partition S1, S2 in respectively student weights that are positively, resp. negatively colinear with the teacher.
linear function in
We now formally prove the claim.
Proof. We collect in W all the row vectors w of the matrices W from the teacher and student network. By using cu to denote the corresponding weight from a in the teacher network resp. minus one times the corresponding weight in the student network, we can rewrite the fact that are at zero loss as
Σ www x) =
for all x E R"
(11)
WEW
To deduce some consequences from this equation introduce
notation. For every W we denote by C., the set of all vectors in W that are collinear with w (including w itself) and with Nu : W C., all vectors in W that are not collinear with Furthermore to denote those vectors that are positively resp. negatively collinear with w.
C.p C+ U C
Consider some arbitrary w E W. By definition of Cww and New that
find a vector an R" such
w xg = O for all w E Cwo' and w xo + for all w E Now
and an open ball B around No such that we have for all w E N wn T x + 0 for all æ E B. Using 1s denote the indicator for the event E, i.c. ig 1 if & holds and 1g 0 otherwise, we thus have
(ww x) 1w #o>0
="
for all x E B, E Now
As w 2o for E Cww and B is ball around ap we know that both B+ = {a E B u x > 0} and B {x B w 0} open in R By definition of the sets City and CT. have
(wTe T if w C+
( w Tx)
for all æ E B+
if w E C
21

---

Published as a conference paper at ICLR 2024
resp.
w x
if w C
4 w x)
Tx)
for all x E B
if w C+
From equation
11 we thus get that
Σ
Cw ll Tooo
Σ cw
x
for B
NEN
weet
and
0
Cuo -o 22>0
w
ww
for all x E B
wEN.
WEC
As both B+ and B- are open sets and we know that for any open set O we have that a. ax = for all E O implies v we thus deduce that
Cw w
Σ
Cw
for all w' W.
(12)
WEC
We now interpret this equation in terms of our teacher-student setting, Recall that we assumed that no two rows in the teacher network are collinear and that in the teacher network no weight in zero. That is, for any w' E W from the teacher network - C needs to contain at least one vector from the student network. As teacher and student network have the same size, this thus implies the existence of a permutation a G₁, such that for all weight vectors w; E W from the teacher network there needs to exist weight vector (a(()) from the student network so that wo and
(() collinear. Denote by S1 the indices i such that t(() E C and by S [1..h] S1 the such that (() Cw.-- Observe that this implies that we rewrite equation 12 as w, = (n(1) Vi E S and a;W; = (a())
Vi S2.
Here we used that we defined Ow for weights from the student network as minus one times the corresponding weight from a. This shows equation [ and equation g equation 7 also follows immediately from the definition of S₁ and S2.
To prove equation 10 and the moreover-part of the theorem, show that if equation equation hold, then equation 5 and equation aree equivalent. As a is permutation we can rewrite equation
Σ
a, 0; 1, = Σ ( (1) 1 #ac0000
Vo E m
(13)
and equation as
Σ a,w; (1w.aso w.o = Σ Cott) wn(() (1.. woo 1w)
Vx E R". (14)
For indices i E S we have w.æ > 0 if and only if (1) > O. thus equation g implies that the terms for i and a(i) contribute identical values to the sums both in equation 13 as well as in
equation [4] We can thus concentrate on the values E Sz. For these we have a, w. > implies
(o) < 0 and vice versa. Using this and equation 9 both sides we see that
Σ a,W;a 1,
Σ 1E8
SES2
is equivalent to
Σ ( ieS.
10sn Σ
a,u. x wwaco-
iES.
The equivalence of equation and equation follows.
If the student is larger than the teacher, ie. h. then there is more flexibility: instead of having mapping from teacher hidden neuron i to single student hidden neuron (i)), each teacher neuron i now maps to of student neurons. In addition, this set partitions into positively colinear and negatively colinear subsets, resp. (1), S2))). Furthermore, we can have additional student neurons that do not correspond to any teacher neurons, provided they cancel out eventually. The sets Sp as well as S 5% in the following theorem correspond to such neurons.
From here on, refer to the operation that concatenates lists vectors.
22

---

Published as a conference paper at ICLR 2024
Theorem 4. Under the assumptions from Theorem
4 but with h 2 h, there exists a partition
S (2(())))((())all)) (
(So) of [1..h) hi) for some h'' Ruch that have
Vi E [I..].,, Vj E (i), 3o > 0 - anw;
(15)
Vj (2)1), 3a < 0 w0 arw;
(16)
0;W;; = Σ ر
Σ
رزشرةة
(17)
1 WVll (W.. MM V( E (1) x (1) U (2 x S2)) 3a > 0: 06 = atik
SES:()
(E52())
(18)
v(J E (1) x (1), 3a <0: wJ = 0'0k,
(19)
0 E رري Σ رزWر
(20)
JES(()
JE8,()
fj Sp. 0 C
(21)
Σ a:00;
(22)
Moreover, any two-layer teacher and student networks that satisfy equation 73 equation satisfies equation 5
22 also
Proof. The proof is identical up to equation 12
Firstly, we denote by So the indices of student vectors such that âzö 0. They correspond to student neurons that have 0 activities everywhere. This gives equation 21 Without loss of generality, we will exclude all student vectors in So from W, and redefine C, etc accordingly.
Secondly, recall that we assumed that no two rows in the teacher network are colinear and that in the teacher network no weight in is zero. That is, for any w' E W from the teacher network Cit, UC needs to contain at least one vector from the student network. For all teacher vector w;
where i E [1..h] we denote by resp. Si)) the subsets of [1..h] corresponding indices of student vectors that belong to C.+,, resp. C Observe that we can then rewrite equation 12 as
a,W, Σ å
Σ
ر Vi [1.h]]
JES,(0)
JES,(()
Here used that defined Cw for weights from the student network minus times the corresponding weight from a. This shows equation I7 equation 13 and equation 16 also follows
immediately from the definition of S and S2.
Finally, consider the student vectors W that have not been assigned to one of the sets S (1), S2(1) or Sg- We fix some (#0)))E(LL. such that [Cw, heeL.. forms partition of this set. We then partition each Cus; into "Ø; C,,, and denote by (i),, (i) the set containing the indices of student vectors belonging to resp. C+ C.. Then we can rewrite equation 12 as
Σ àBØ
Σ قههة E [1..V]] (ES
(ES(()
This shows equation 20 equation of St and S.g
18 and equation
also follows immediately from the definition
@ To prove equation 22 and the moreover- part of the theorem, we show that if equation 15 equation 21 hold, then equation and equation 22 are equivalent. We can rewrite equation as
Clearly,
@ (S,) form a partition of [1..A..
a, W;
1 = Σ
R
(23)
xE(
(E(1..k
23

---

Published as a conference paper at ICLR 2024
which gives
[
Σ
رررة
Σ
رzt0,] x
Lwyæ>0
(24)
ie[1..]]
JES, (6)
(ESx(1)
Σ
رر 1 100=>0 +
Σ
dgشرyy 10
(25)
sejt..M]
JES, (1)
(ES((1)
Σ
Σ
مstر 1 1 +
Σ مstرyy 1
Ve E IR...
(26)
kE (ES(()
)ES(())
which simpifies into
0 ==> Σ
" 1 --
Σ
" 1 x
(27)
iE[I..h] (ES2(1)
(ES
Σ م 1 m +
رققرة
L
Vo E R".
(28)
RE[1..]]] (ES(())
(ES
Noticing that twos + L-w2>0 for almost any and for any o O, we get
ā +
Σ
(29)
xE[I...] jE
((((,)))EE(())
for almost any æ, and thus
0 Σ
a ش +
Σ
Σ
(30)
iE (C
JES
<E(1.M))EE(00
the equivalence of equation 5 and equation 22 follows by adding equation 17
,qquatinn 20
eqqua-
tion 21 over all indices and using equation 30
A.1.2 IRREDUCIBILITY CONDITION FOR RELU NONLINEARITY
Here, comment the formal sufficient and necessary conditions for a ReLU two-layer ML.P to be irreducible. As it is not relevant for the remainder of the appendix, can be skipped.
Consider two-layer ReLU MLP, fx) a 4 (Wz) Σx (w æ:).
Firstly, a pair of weights Wk, wi are positively colinear, then they can be fused into one hidden unit. Thus, a necessary condition for irreducibility is that no two weights are positively colinear.
Secondly, if some pairs of weights are negatively colinear, we can decompose the corresponding function into the sum of hidden neuron and linear function, i.e. for any (w,, un, agg a1)
such that there exists a >> 0: wi -aw,, for any
a aptb(un ar) + amp æ:) = an vl un ω)) + arw x²) + warwo x camw x
(31)
= as (w) x) - arb(aww x:) caww x
(32)
(ax (ww) www x
(33)
(34)
where we obtain the second line by noticing that for any w, x. v(w" x) w w x).
Thus, any two-layer ReLU MLP can be transformed into a the sum a ReLU MLP with no pairs of weight being colinear, and linear function. Since linear function can be implemented by two hidden units with negatively colinear weights, any irreducible network should have at most a single pair of negatively colinear input weight.
Thirdly, the obtained linear function can further be absorbed by the ReLU hidden neurons if, by denoting w the weight of the linear function, there exists a subset of hidden neurons I such that
24

---

Published as a conference paper at ICLR 2024
1 a;
(35)
iE
since in that case, we have, for any Σ (w, x) - E
Σ
aavi (w x) w x
0: x
(36)
RET
iET
iEz
aa (ww x:)
(37)
EE
Thus, for f to be irreducible, the following necessary conditions need to be fulfilled:
There exists no output weight og that is
There exists no pair of input weights that are positively colinear
There exists not more than one pair of input weights that are negatively colinear
If pair of negatively colinear input weight exists, by denoting k. I the corresponding hid- den neuron indices, there are no subset of [1..h] (k, suchh that a = -. a or dyw) = Eiez 0:
The above conditions constitute in fact sufficient conditions for irreducibility of two-layer ReLU MLP. The proof follows a similar structure as that of Theorem and notice that necessarily 2 h.
A.2 MULTI-TASK PARAMETER IDENTIFICATION
We now come back the multi-task setting, where the teacher generates task-specific network using a hypernetwork conditioned on a task latent variable. Specifically, the weights of the second layer fixed across the tasks, but the weights of the first layer are obtained as linear combination of the modules: WOO,)) = ΣM 2(m) (mm) See section 3.1 for the full description.
Given a task latent variable z, clearly the theoretical results from the last section can be applied to the generated teacher and student network in order to establish relation between the teacher and student parameters. We will first extend the theoretical result establishing the sufficient and necessary conditions on the student parameters for the student to be able to fit the teacher on any task.
We provide results for MLPs with ReLU activation function as well as MLPs where the activation function belongs to general class of smooth functions, for student networks that can be wider than the teacher network. Theorem in the main text a special case of the following results.
A.2.1 RELU MLP
Theorem 5. Assume Par has full support the input space, M S n, M = M, h 2 h, and that there exists a full rank matrix in the linear spon of (0i)siigh and (a:)ı<igh are all non-zero. Assume furthermore that there exists at least one z such that no two rows of W (0, z) are colinear. Finally, assume vp is the ReLU nonlinearity. Then, ming L(ž, ((,,(,,0) a):)):: D2) 0 for almost any z is equivalent the existence of invertible matrix F, partition S (S (i), ()),)(())))))))))))))))))))))( (2))))e(().))))) @
25

---

Published as a conference paper at ICLR 2024
((()))) @ (So) (1.h) such that we have
Σ 0, = IF 1A WE [1..). E ( Vj 3a 0: 6, o0,
(38)
(a)
(39)
Vj E (2)1), 3a <0: <0 aO;;
(40)
a. Σ 6, Σ a,
(41)
jES(()
jESa())
iA E [1.A] V(, E S((1) x ((())S((() 82(1), 3a 0: 6, ab,
(42)
V(i k) E Si(1) x (2)1) <0:0, = =
(43)
0 = Σ ā,6, Σ ā,O,.
(44)
jES(() Vj So 0 0 6
jES₂(()
(45)
Before proving the theorem, we provide Lemma which will be useful.
Lemma 6. Assume there exists one point za such that W(O, (9, zo) has no two rows that are colinear. Then, for any linear subspace U containing zo. ( 2) has no two rows that are colinear for E U almost everywhere In particular, the set of points z in U such that W(O, z) has at least two colinear rows form a finite union of strict linear subspace of U.
Proof. Recall that the ith row of O z0) is given by 0:20. Consider pair i, with * j. As no two rows in ( 20) are colinear, there exists only a finite number of a E R such that the equation yields solutions of that are not in Ker θ n Ker θj. Given each of such a, the set of = that verifies the equation form a linear subspace. Taking everything together, the set of z such that ( 2) has least two colinear elements is a finite union of linear subspaces. Due the existence of za we know that these linear subspaces must not include any linear subspace containing zo and the claim follows.
(O, a0,))
We now prove the Theorem.
Proof. By assumption, there exists at least one point 2g such that ( =0) has no two rows that are colinear. Following Lemma for almost every z E. M the first layer weights have no two rows that are colinear. For all z, we & E arg ming L(2, (Ô, à); Dz). By assumption, L(É, L(,,0;
for almost every z. Considering the intersection of such and those that generate first layer weights with no two colinear rows, we can apply Theorem 4 for almost every z. For each such z the theorem guarantees us partition S = (S) (i), (2))))))))))))) (2)))ee + (S) S1(1))ep @ (So) such that equation 15 equation 22 hold for the generated weights.
A priori these partitions need not be the same. We prove this next. Denote for fixed S, US the set of all points in R for which Theorem 4 can be applied, and holds for this S. As Us U is dense in MR and the set of possible partitions S is finite, there thus has to exist at least one so that RM Sp) This means in particular that can write any point in R as linear combination of points in U.S. Thus for any point in RM there exists a & such that equation 15 equation 22 holds for S.
It remains to show the existence of an invertible matrix F, as well as the constancy of the various parameters w.r.t z. Using equation 17 in particular, we get that
Vz E RM 38 Vi [1..h], 0,,2= =
Σ a,,, (ES
Σ
a, 2
(46)
SE82())
Because there exists a full rank matrix in the linear span of the (0,) then clearly we can find some scalar coefficients (A)sen...) and linearly combine the equations 46 over i such that we
By "z U almost everywhere" we imply using the Lehesgue measure
26

---

Published as a conference paper at ICLR 2024
obtain equation of the form A Bi for all RW and corresponding Σ, where A is full rank. Necessarily then, B is full rank, and multiplying by its pseudoinverse yields Fz = where F is invertible matrix. for which the equations hold is thus linear function of z, which also implies the constancy of the various parameters a w.r.t. 2.
The proof is concluded by applying the fact that Pz Q= for all RMM implies P Q all equation 15 equation 22
In particular, if we restrict the student width to be equal to that of the teacher, we get the following corollary.
Corollary 6.1. Assume Ps has full support in the input space, M S n. M = M, = h, and that there exists a full rank matrix in the linear span of ( and (a:) kishh are all non-zero. Assume Surthermore that there exists at least one such that no two rows of W(O, z) are colinear. Finally, assume is the ReLU nonlinearity. Then, ming L(2, (+, a); D=) 0 for almost any is equivalent to the existence of an invertible matrix F, a partition = (S), S2) and a permutation o of [1..h) such that we have
E00 -a.
(47)
Ví, (o( >0
(48)
Vi E S1, (o(i)
a,
(49)
E S2, an(i) On() F 0,
(50)
A.2.2 MLP WITH INFINITELY DIFFERENTIABLE ACTIVATION FUNCTION
We first provide a reformulated version of the main Theorem 4.2. from Simsek et al. (202) which is the counterpart to Theorem 4
Theorem 7. Assume teacher and student two-layer network parameterized by resp. (W, a) and (W, à),, of hidden dimension resp. h, h with h and single output unit such that
Ra [T(xx)] =0
(51)
where
L(x) "1aa" (WW)) aa
(52)
Assuming
Pe has full support in the input space
(W, a) is irreducible, i.e.
no output weight is 0
no input weight wi 0
no two input weights wy, are identical
1/ _ and (/(m) (0) * for infinitely many even and odd values of
Then there exists partition (8.) (8)) of (1..h) for some h! such that have
Vi E [1.M],
Vj E St, v =
(53)
Σ à = a:
(54)
jES,
27

---

Published as a conference paper at ICLR 2024
and Vi E [1..k']
V(k, S x S Wi wt
(55)
Σ
aj 0
(56)
jES
In essence, if the teacher parameters irreducible, we learn the same parameters up to permutation when the hidden layer is of the same width. If we allow the student to have a wider width, the overparameterizattoon can either:
Allow hidden teacher neuron to subdivide into multiple student copies whose contribution the last layer sums up to that of a teacher neuron.
Allow for null neurons to appear, i.e. neurons who have no contribution on the final layer.
We now provide the identification statement in multi-task teacher-student setting.
Theorem 8. Assume Pa has full support in the input space, M S n, M = M, h h, and that there exists a full rank matrix in the linear span of (0.) and (ai)isish are all non-zero. Assume furthermore that there exists at least one z such that no two rows of W(O, z) are identical. Finally, assume the activation function vp satisfies &b E C and (0) + 0 for infinitely many even and odd values of Then, ming L(ž, (O, (); D2) 0 for all is equivalent the existence of invertible matrix F and partition (Si) (sigs for some h' of [1, [1,...h such that
Wi b E S;, =,F=O,
(57)
Σ a =
(58)
JES:
vi
"L......'NV. E S x S 6(0) -00)
(59)
Σ ā =0
(60)
JES
The proof follows the same steps as Cheoremm3excepp for the use of Theorem 7 instead of to prove the desired equations.
A.3 EXAMPLE FAILURE CASES BUILD INTUITION
We wish to uncover in the following the conditions under which student might achieve composi- tional generalization when trained on a training set of tasks. This subsection provides some failure cases where perfectly learning the training tasks does transfer to unseen tasks. It does not provide proofs, but rather an intuition on the later stated theorems and the assumptions they require.
In what follows, (em) denotes the canonical basis of R.M.
Wider student Let us consider teacher with one hidden neuron and three modules: O (A) for three vectors A, B, that are independent. The output weight of the unique hid- den neuron is a A. We allow the student to have more than one hidden neuron. We consider training task distribution defined by the binary masks {er feites +2, e2 +e) The task distribution is then trivially compositional and connected. Yet, we will show in the following that when -= = h 2, student can be constructed which perfectly fits the teacher on all training tasks, while being unable to generalize compositionally. Let a student network with the following parameter specifications:
M
m
m
A
B
0
0
B
C
Ôg
0
B
0
28

---

Published as a conference paper at ICLR 2024
identification
B Linear identification
Given teacher with M-3 modules and h-3 hidden neurons
Given teacher with M-3 modules and h hidden neurons
Wi
Madile
Module
Modkaee
Mustuee
Module
Module
Neuron
WW
Wa
W
Neuron
T
Wa
W.
PerwOrre
Neuran
W
W.
WA
POWISSSON
Neuron
W3
W.
WA
Nauron
W.
W
W.
Neuron
W
W
W;
W
Teacher modules
W
Teacher modules
Veuumin
h
training the student an distribution with disconnetted support
training the student on distribution with connected support
Task 6
Neuran
TWa
W,
W;
Task family 1: by
(
I
O)
Neuron
[Wg
Wy W.
Task farmly 2: 65
Neuron
W
W.
W
Task family 2: by 10
Neuron
Wa
W
W
Task farilyy 3:
Neuron
W.
Wa
Wi.
Neuron
W,
W.
WA_
Student modules
Student modules
lead inconsistent neuron permutations
modules Wi W. W, * ii; W.
leads consistent neuron permutations modules
OOD task family 600D (1 1)
Neuron
OOD task family boon (1 u 1)
Neuron 1-2: Wi-W, W, W.-W, W,
Neuron
W2 Wt Ws * Wa- W/ W,
Neuron 2-4h:
W2 WY Wg-Wi ML
Neuron
W/3 W. W. 4 Wg- W. We
Nauron Sints
Wy=W; W;
Wj=W,
Teather
Stunent
Teacher
Student
Figure A1: A No identification. Toy example of how correct identification of individual mod- ules can lead to inconsistent neuron permutations preventing compositional generalization to un- seen 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. For instance the weights denoted by W, correspond 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 used, Since in MLPs the ordering of neurons is permutation invariant, the student can perfectly match the teacher weights, even when it uses different ordering of the neurons. As result, the student modules can perfectly fit all three tasks, despite permuting the neuron indices. For instance, neuron 2 in module of the student contains the weights W, whereas the corresponding neuron in the teacher contains the weights Wn, When we now present a new task during out-t--distribution evaluation that mixes two of the modules in the teacher, the student is required to mix the weights of each neuron across these two modules as well. 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. B Linear identification. Toy example of how having connected support helps ensure that neuron permutations across modules are consistent allowing for compositional generalization to unseen module compositions: The teacher is setup identically A. Different before, the training distribution now has connected support, i.c. the binary masks defining the task families share non-zero entry. After learning the neurons of the student are still permuted compared the neurons of the teacher but this time the permutation consistent across modules, i.e. compared to the teacher only rows are permuted. As result, when presenting a novel task from a task family that mixes modules, the student is able match each of the teacher neurons and therefore compositionally generalizes. While this example only shows permutation of the learned student neurons consistent across modules, in general the student modules will be a linear transformation of the teacher modules, hence the naming linear identification. Importantly this lin- ear transformation is consistent across modules given the conditions of Theorem 2 are satisfied.
29

---

Published as a conference paper at ICLR 2024
output weight
(()
X
(2)
A
ú (3)
--
One can verify that this behaves as expected on the training tasks. However, if we now take task e1 e2 + eg, one can also verify that it does not behave as the teacher. The same argument can be made to break the compositional generalization of any teacher network width given this training task distribution.
Disconnected tasks We consider student with two neurons and two modules. Let θ1 (A(B) and 02 = D) for four linearly independent vectors A, B, C, D, and output weights all equal (1 02 A. If only consider tasks defined by masks {ei, e2 a valid student is given by θ₁ () and &₂ = B) and the same output weights. This student does however not generalize to the task described by mask ej + e2. Also see Figure A1 for more detailed example of this failure case.
Non-compositional support This is trivial failure case: if there is teacher module which does not appear in the training task distribution, then it is hopeless for the student to generalize to new tasks involving this unseen module.
A.4 COMPOSITIONAL GENERALIZATION
We will now provide the formal theorem for the ReLU nonlinearity the activation function. The- orem in the main text is an informal presentation of this theorem. To show compositional general- ization, similarly to Theorem which we obtain by applying Theorem i in fact common to all task latent variables in the sup- port. Coupled with the assumption that Pg has compositional support then suffices show that the sufficient conditions of Corollary holdd
we will show that the z-dependent partition Si, S2 and permutation
Assuming the support of the task distribution contains task families 2₁ as described in the main text, we will proceed in two steps:
1. Given k we will show that if the student achieves zero loss on all tasks in Za, then the partition S1, S2 and permutation a for all x E Zg is identical.
2. Using the property connected support, we will "glue" these partitions and permutations together to show that they are in fact the same for all k.
Showing the first step requires Pa to verify number of properties. First of all, clearly need that P, is such that the probability of choosing a value from Zn is non-zero for all k. However, additionally need some properties how the probability mass distributed on each Zz. The key idea of our proof is to show that generalization within Zμ actually occurs as soon as we have zero loss on an appropriate basis of Zg. For this to work an obvious necessary condition is that (strict) linear subspace of span( Zx) should not contain probability mass larger than zero. Note that this implies that the probability that any d, i.i.d. samples of Zz. where dj dini(span the dimension of Z1. will not form a basis of the space spanned by 7 is zero. In fact, we need slight generalization of this idea, as we also need that certain other non i.i.d. sampling of da points have zero probability. For this to hold require that the induced probability distribution obtained by projecting the points in Z on the unit sphere admits a density function.
denotes
Formally, thus want the training task distribution P: have the property that can extract family Z from the support Z Ps, with an associated family of binary vectors {b...... bx} and spans span( (Za)
(o @ by RRM J, such that the following three conditions satisfied:
P PANP, = 0
For any hi P PPPPP_ (z 7x)
Pieza projected on the unit sphere admits probability density function on the unit sphere
30

---

Published as a conference paper at ICLR 2024
where we denote by PaeZ, the probability distribution P. conditioned drawing value from Zy-
The above conditions would hold, for example, for any distributions Ps which is a discrete mixture of distributions PreZ. such that each PzeZ, is uniform distribution defined on the set of convex combinations of the canonical vectors corresponding to the non-zero entries of bg. as we do in our experiments.
We are now ready to state and prove the theorem. The proof for the general class of smooth activation follows the same steps, omit the details.
function used in Appendix follows
Theorem 9. Assume Pa has full support in the input space, M S n, M M, h that E, agO; is full rank, and (as) are all non zonozero. Assume furthermore that for all modules @ no two rows colinear, that
extract family Z& from the support Z of P. with the above property, and that P= has connected and compositional support. Then,
Ez-Ps min (0, (); D2) = = Vzz min L(2, (),
D2)
(61)
Furthermore, for any k, we denote dN. dim(span(Zz)) Then for any set of task latent variables F= Un F& such that Fv 1k are sampled i..dd from PzeZu we have
Vz E F, min L(2, (Ô, D2) =0 = Vz, min L(2, ())
(62)
In the following, we will prove the result in equation 61 The result in equation 62 requires a more involved intermediate result which we provide at the end of the section.
Proof of Theorem first part. In a first step we prove the existence of an invertible matrix F, and for each k. the existence of permutations of [1..h], and partitions S of [1..h] such that Zz E span (Zz).
Vi E St -oo
(63)
Vi St =,0 aox(
We start by arguing that we may assume without loss of generality that all 2 E
Z₁ sat-
isfy some nice properties. First we claim that may assume that all points E Z, sat-
isfy min; L(2, a); D2) = 0. This holds as Re-P, min; T(2, (Ô, à); Dz) = 0 and
our assumption that PZ-P,, [z E 7 > implies that the set of points z of Z such that min; L(2, (0, à); D=) f 0 has zero probability.
Next we argue that we may assume that for all 2 E Zμ we have that W z) contains no two rows that are colinear. To see why this is true recall that we assumed that all modules (()) contain no two rows that are colinear. In other words, all canonical basis vectors corresponding to the non- zero entries of by have this property and the claim thus follows from Lemma 6 together with our assumption that P.ez, projected on the unit sphere admits a probability density function on the unit sphere.
Taken together, this allows us to apply Theorem 4 on all elements of Zg. That is, for all in Z there exists σ2, S2(x) such that we have
Vi E S1(2),
à,,, ar, (1) Oz2())²
(64)
Vi E S2)x,,
0 = -a,(ii) Oz)
(65)
Σ 2,,02 = Σ a, a,
(66)
where here, and henceforth, & argming. Uhe2 Dz).
A priori these a2, S(z) (S, (2()) need not be the same. We prove this next.
31

---

Published as a conference paper at ICLR 2024
We group the set of all points in Zy into sets (((,5) so that the points in U70,)) all obtain (a, S) when applying Theorem 4 Clearly, U(a.s) U(a,5) Zz. Furthermore, because Px when projected onto the unit sphere admits density function, it follows that Zμ contains an infinite number of points of which any subset of dk elements form a basis of span(Zx). Since the set of possible permutations and partitions (a, S) is finite, there has to exist at least one pair (a, S) so that U(a,5) contains basis of Z.) By linearity, this implies in turn that Theorem 4 holds for this (o, for every point in span(Zz)
We can thus fix σg, St S such that for all E span(Zs))
Vi Sf. 0,O,2 (1) (1)
(67)
E Sh à,O,² = -0 (1) (o)())
(68)
La.O,a Σ a,
(69)
remains to show the existence of the desired matrix F. Towards this end, recall that we assumed that E, a, is full rank. As we also assumed that span(U,, Zx) R it follows from equation 69
that Σg a, also has full rank. Thus, letting F EE3 ,, Σ a,Or, where At denotes the pseudoinverse of A, we obtain a matrix that is full rank and satisfies & = Fz for all z in span for all k. Taking everything together, the equality in equation 63 is obtained.
We now show that the St S2, oz are the same for all k. By our assumption on the connectedness of the sets Zk. it suffices to show the following implication: if k # I are such that there exists a in span( Zx) n span Zi) so that W(O, x) has no two rows that are colinear, then 0g dr, and St = St. St = Sz.
To see this, fix some in span Zx) n span(Z)) so that W W(O, x) has no two rows that are colinear. From equation 63 we obtain the equality
Vi E [1..h], (o) 2 (0) (-1) 1108 Oz
(70)
(-11) - (-1) - a (0) o,)(1)
(71)
i.e. we get that for any i, θ (1) θ (z)()) are colinear, which is contraction unless oT (i
Therefore, 0 = ot as desired.
We now show S Sh.- Without loss of generality, assume there exists such that i S and
i & Sy. We fix such an i. Then, using equation 70 as well as (i) = (i),,
(o2 4 og (1) = (1)
(72)
-a, (1) (1)
(73)
and thus ar O0, (1)= 0, which again is contradiction of the non-colinearity of the rows of W(O, x). The claim thus follows.
The upcoming lemma is crucial for demonstrating that students almost surely generalize composi- tionally after achieving zero loss on finite task sample from P.. Loosely speaking this Lemma is a generalization of Lemmal0Jin that it asserts that under the condition that no two elements from (O,); are identical up to scaling, a randomly sampled set of M + task latent vectors has the following property: there exists a rescaling of the points such that the sum of rescaled points is zero (this trivially holds as we have more points than the dimension of the space) and such that by applying arbitrary sign inversions and multiplication by arbitrary elements of (Ot); to each scaled vector, the sum would not be zero unless the matrices and the sign inversion applied to all vectors are the same. In the lemma we choose the M + 1 points from the unit sphere, and one can think of them being chosen i.i.d. uniformly. Within the proof of Theorem [we will then argue that our assumptions on P. allow us to replace this uniform sampling from the sphere by a sampling according to P..
Henceforth, measures on the sphere are spherical measures (resp. product measure of spherical measures for Cartesian products of spheres), and Lebesgue measures in Rd for all d.
Lemma 10. Let a family of Rud matrices (O;) isigh Assume furthermore that there exists least one 20 E Rdd such that two elements of (0;zn); are colinear. Let S be the unit sphere of Rdd
32

---

Published as a conference paper at ICLR 2024
Then, for almost every set of +1 points _....., 20,1 in Sd+ there exists A1 da+1 such that
(zele == (de&ele verifies the following conditions:
Σ 1+11
Ze =0
v(te)e E [1.,6] M(se)e
E 1-1, 111,1
(74)
Σ
=. =0 = V 11, = i1) (8x = s)).
c6(1.d+1]
Proof. We define the matrix = (11, 12.(a4))) where I, for all is the M x M identity matrix, the Deel d+1] Re which concatenates d + vectors of R into
isomorphism & (ze [eee..(_.+1] vector of [d((++1), and the projection to the unit sphere / x T defined everywhere except the origin. Finally, let C ((z))e((L..d+)) FQ & [1..d+ 1] (Ze) linearly dependent), i.e. the set of family of vectors such that not any subset of cardinality d form basis of Rd
The proof will proceed in the following steps:
1. First, we show that for almost any x I) \ ((C)), property 74 holds for 4 (z))
2. Then, we show that there exists surjective function from ga+1 \C'to P(Ker (C),, such that the pre-image of any zero-measure set is also zero-measure, and such that for all (Ze)eE[1.d+1) E 5u++1 ' C. ((ze)e) is transformation which scales each member by a scalar and concatenates them, i.c. (((z)))) E (Veeez))c) (2)))e(I..d+IIJEeR1
We remind the reader that if a surjective function f A B between two metric spaces A, B equipped by metrics Já' verifies that all pre-image of nullsets are nullsets, i.e. AX C B :::44'( (X) = O, then we have that if some property is true for almost every x B, it must be true for f(x) for almost every x E (B) = A. If the above 2 points are true, then by application of this reasoning, we get that for almost any (ze)e 5 C. there exists (Ae)e such that property 74 holds for (ee.e) Since it can be shown that the measure of gu+1 C is 0, we have that the assertion holds almost everywhere on Sd+11 which would conclude the proof.
O (x)=
In the remainder we will prove the two points above.
Let some i (ie)e E [I./t]d+1 1 and s = (se) E (1-1,1]d+) such that there exists a pair (n, m) such that in - im 8n f 8m. We fix such a pair. Let Mi,- = (s1 O ...4 O ). We will now show that I & Ker Mr iee that there exists a point in Ker which is not in Ker Mr Let the family (ze)c = (1 (1e=nve=m( (-1) 20)a. Clearly, Σe 2e 20 - 20 0, thus ((2)))) belongs to the Kernel of I. However, E. S,O 0 = 8, 9
Step 1:
8m P,.,. 20- Because by assumption no two elements (0,z)): are colinear, necessarily E. s,O, Ze
(sn = sm) A (in, im), which contradicts the definition of i, Therefore, necessarily E. Di Ze + 0. ie. (((2))c) & Ker Mi.--
Ker M.,, is linear subspace, and P(Ker I) is unit sphere living in linear subspace. Since Ker I ạ Ker Mr,, then I) of Ker M, and thus P(Ker Z) n Ker Mi, has zero mea- sure in P(Ker Z). Because finite union of nullsets is nullset, it is then clear that P(Ker n U.. Ker Ms,s has zero measure in P(Ker Z), where the union is done over = (ie)a E [1..1] and (s.) (-1,1] such that there exists pair (n, m) such that i, # im or s,, # 8m. On the other hand, it is clear that for any family i, such that all elements of i resp. 8 are the same, we have P(Ker I) C Ker M..-
This now proves that for almost any P(Ker Z), & (x) must satisfy the property 74 see that (CC has zero measure P(Ker I), which finalizes the claim.
We can
Step 2: Let S+ = {(ze)e(()..4+1) E gd+1 3(Ac) : (Ee Aete 0) >e
We will define two surjective functions, E \C + S \C and T: C (CC) such that their respective pre-image of nullsets are nullsets, and each of which multiplies each member of
33

---

Published as a conference paper at ICLR 2024
the input (ee)e by scalar, before eventually concatenating them. If so, their composition vP ToE will inherit these properties, which finalizes the proof.
We first define Y. Notice that on St \ C, the scalar family (Ae)ee(I..dt)) is unique up to scalar multiplication. The family becomes unique we furthermore restrict it to be of norm 1. We can thus define the function Y which maps each (Ze)eE(1..at)) the unique & (((AA=)))e(1..att) such that (Ee Ae 0) A (Ve, >. > 0) A (E. X2 1). Clearly, Y defines bijection between S+ C and P(Ker I) \ ((C),, with the inverse function Y-1, - ( ( with (ze)e = (z), defined on P(Ker Z) \ (C)). Furthermore, Y-1 is differentiable everywhere on P(Ker I) (CC), and thus preserves mullsets.
We now define E. For all points (ze)e in Sd+1 \ C, there exists a family of non-zero scalars (Ae)er unique up to scalar multiplication, such that Σ. AgZe 0. Given an arbitrary such (Ae). for which
A1 > O. we can define the function E from Sd+1 C St C as E (ze)e - (sign( ) )ze)e- E is surjective, as it leaves all elements of S+ \ C C Sd+1 \ C unchanged. Furthermore, the preimage of any set X by is included in the finite union U, Xs, where E {-1,1) and
X, = {(ze)c (soza)e E X). Since X, is trivially a nullset if X is nullset, the preimage of nullsets by E are nullsets.
We can now prove the part of the theorem stating that a final sample from the training task distribu- tion will almost surely allow the student to generalize compositionally.
Proof of Theorem 2 second part. Observe that it suffices to show that for any k. equation 63 holds almost surely when sampling dy + 1 points i.idd from the probability distribution Pse The re- mainder of the proof is then identical to the continuous case.
Fix some set Zk- As in the continuous case we may assume without loss of generality that all points z in 7μ satisfy that no two rows of W z) are colinear and that they achieve zero loss, i.e. ming L(2, (0, (): Dz) =0 0.
Now, consider Fs = (*k.e) disesd +1, family of dę + task vectors sampled i..d. from Z₁ follow- ing Prez, Due to the assumption that P=ezz projected onto the unit sphere admits probability density function, any subset of size de of Fs form a basis of span almost surely, and in partieu- lar span((= ( to each vector individually. we can show that the resulting permutations o and partitions S1, S2 are identical across these vectors, then that would suffice conclude the proof.
span( Zx) with probability one. Furthermore, we can apply Theorem 3
In the previous scenario where the loss was zero at an infinite number of points in Zg. it was pos- sible to extract a basis from span(Zx). with each basis element having a consistent permutation and partition. This strategy, however, is not applicable in the current context. To establish the uni- formity of permutations and partitions for all vectors, that with probability 1, F& has the property that there exists a family of scalar (Ae,e) such that (At,ezk.e
must first refer to Lemma 10 and show
) seedd--1 verifies equation 74 We will prove this next.
Because we have that any vector z in Zx satisfy that no two rows of W(O, 2) are colinear, equiv- alently, two elements of (a,O,z); have two colinear elements. We can thus apply the lemma to (a. which would guarantee that almost everywhere on S where S is the unit sphere of R.M. the property holds. In fact, if we consider the sphere S₁ of the linear subspace span(Z,, Zzx), is possible to show that the property holds everywhere on Sd++1 except for a set of measure zero. Because PLez. projected onto the unit sphere admits probability density function, it is easy to see that when sampling i.i.d. dç + 1 samples from Pzez. and projecting each of them on the unit sphere, the family will be on this set of measure zero with probability zero. Because whether the property holds or not is left unchanged by normalization of each vectors in the family, clearly the property must hold with probability for Fr. This concludes the claim.
We now show that all the permutations and partitions associated with the vectors in F& are the same.
For all x, we fix E arg ming L(2, (Ô, à); D2). As argued above already, it follows from Theo- rem 3 that for any in Fk. there exists permutation and partition of (1..A), a and (8), S2(=)
34

---

Published as a conference paper at ICLR 2024
such that we have
Vi (,),, aoz(1) (1)
(75)
Vi S2(z), 0,0, = -0) (12())
(76)
Σ 0,0, = Σ 0,0,z
(77)
Similarly before, since span(U; FR) span Zx) Rdd, the last equality be used to obtain an invertible matrix F such that for any k, and any in Fk. 2 = Fz, which establishes linear mapping from 2.
In the remainder, we drop the index in for notational simplicity. We fix the family of scalars (Ae) such that equation 74 holds for (NeZe)e.
In particular, we have Ee Aeze = O. By multiplying by F, we have Ee Acãe 0.
For any i, multiplying by â,O,
=
(78)
which in turn implies
Σ (-1)ttess( ag 2 1o) =0
(79)
Since equation 74 holds for (Aeze)e, clearly, vi E [1..k), Ve, f:0 (1) = (())
= S1(z1)
and S2(ze) = S2(2)), as desired.
35

---

Published as a conference paper at ICLR 2024
Discrete
Continuous
32
Out-of-distribution loss
0.002
Module overparamett zeeppaaaattizaiinnn
600
16
0.0015
400g
8
0.001
200μ
4
0.0005
8
16
32
2
16
32
Figure A2: Hidden overparametrization
Hidden overparametrization
loss is sensitive to overparameterization. Numbers denote the factor by which the student dimension is larger than the teacher. Error bars denote the standard error the mean over seeds.
Table A1:
Effect of weight decay on module alignment in the overparameterized regime. Numbers module alignment averaged over 5 seeds.
M 16 M
M 32 M
Weight decay strength
h 16 h
h h= 32 h
h = 16 h
h 32 h
0.001
0.7704
0.7165
0.8369
0.8144
0.0001
0.7699
0,6752
0,8483
0,8043
0.00001
0.7633
0.6727
0.8346
0.8292
0
0.7542
0.7109
0.8421
0,8076
B ADDITIONAL EXPERIMENTS
B. MULTI-TASK MAACHER-STUDENS
Overparameterization hurts compositional generalization. We show in Figure A2 the out-of- disribution (OOD) loss for the experiment presented in Figure We see that the degradation in OOD loss mirrors that of identification.
Weight decay does not improve module alignment of overparameterized models. We repeat the overparameterization experiment with varying weight decay strengths, for the discrete task dis- tribution. We report the results in A We see that regularization using weight decay does not seem to help to improve module alignment for overparameterized models.
B.2 HYPERTEACHER
Sample complexity scales with the number of tasks. We investigated empirically, how many training tasks need to present in order to obtain specified OOD accuracy. Figure shows this scaling behavior.
Sensitivity to finite data samples for task inference. In contrast to the infinite data regime con- rely gradient-base meta-learning from finite sam-
sidered in the theory, in the hyperteacher ples. We investigated the dependence of our results on the number of train shots in the support and query set for each task in Figure A4 finding that redueing the number from the N 256 consider throughout our experiments leads to a decrease in performance.
Sensitivity to fraction of held-out module combinations during training. In our main results hold-out 25% of possible module combinations during training to evaluate OOD accuracy and quantify compositional generalization. In Figure A5 we vary this fraction for M E [4, 8, 16) and K revealing that for larger numbers of modules, OOD accuracy stays increasingly robust across
36

---

Published as a conference paper at ICLR 2024
Tasks until 90% OOD accuracy
2M
1-to-K hot
1.5M
0.5M
5
10
15
Number of factors
Figure A3: Number of tasks encountered during training before achieving 90% OOD accuracy across modules M and maximum number of modules per task K in the hyperteacher.
Linear hypernetork
Nonlinear hypernetwork
teacher modules
teacher modules
1-
OOD accuracy
0.8
1-to-2 hot
1-
OOD accuracy
140-4 hot
0.6
262
0g
25g
4622
6g
Shots train
Shots train
Figure A4: Influence the number of train shots on OOD accuracy in hyperteacher for M {4.8)
and K E 4).
37

---

Published as a conference paper at ICLR 2024
Nonlinear hyperneteerk
hypernetwork
Learned initializzaiion
0OD Train
tack factors
8 tack factors
16 task factors
0.8
Accuracy
0.4
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
Fraction OOD tasks
Fraction of OOD tasks
Fraction of OOD tasks
Sensitivity OOD accuracy to the fraction held-out module combinations during training for the OOD set in the hyperteacher.
Figure A5:
Linear hypernetwork Nonlinear hypernetwork
Linear decoder modules (R9)
modules
2 4
2
4
hot
1-to-K hot
1-to-K hot
Figure A6: Linear decodability for varying number of modules E {4, 6, 8) and maximum num- ber of modules per combination K {1,2, 4, 8} in the hyperteacher.
larger fractions of tasks held-out during training. Note, that at the same time expected training accuracy remains comparably stable.
Linear decodability for varying number of modules and combination size. To complement Figure 3E in the main text, we show the corresponding linear decodability across modules M e (4, 6, 8) and maximum number modules per combination K E (1,2, 4, } in Figure A6
Detailed metries. For completeness we provide detailed training metrics for all methods across range of settings for M E (4, 8) and K E (1,2,4, 8) in Tables A2
A3
A4
A5
A6
A7
A8
and
present training curves in Figure
B.3 COMPOSITIONAL PREFERENCES
Overparameterization in the compositional preference environment. Motivated by the sensi- tivity to overparameterization in the multi-task teacher-student setting, we investigated the effect of varying the hidden dimension and number of modules in the hypernetwork models on the OOD loss
Table A2: Hyperteacher comparison of accuracies across models for M=4 K=1.
ANIL
MAMI.
Linear hypernetwork
Nonlinear hypernetwork
Train accuracy
79.68 0.25
85.14 0.57
80.98 + 8.92
97.31 + 0.76
Test accuracy
79.72 + 0.12
85.13 0.54
80.89 9.1
97.15 = 0.93
OOD accuracy
46.52 0.8
20.53 = 0.72
20.42 + 1.64
20.93 = 2.22
OOD accuracy (K=1)
46.48 0.83
20.58 0.74
20.59 1.52
20.9 + 2.19
0OD accuracy (K=2)
50.47 + 0.26
26.71 0.26
27.8 2.39
26.78 + 3.33
0OD accuracy (K=4)
48.18 = 1.2
22.74 = 0.59
24.57 2.57
23.04 + 2.86
38

---

Published as a conference paper at ICLR 2024
ANIL
Linear hypernetwerk
MAML
Nonlinear hypernewoork
teacher modules
teacher modules
teacher modules
OOD accuracy
0.8
0.6
1
OOD accuracy
0.8
1400-2 hot
0.4
0.2
OOD accuracy
1:to-4
0.6
0.4
0.2
1-to-K hot
1-to-K hot
1-to-K hot
Figure A7: OOD accuracy the course of training for various models, number of modules M (4, 6,8) 8) and maximum number modules per combination E (1,2,4).
Table A3: Hyperteacher comparison of accuracies across models for M=4 K=2.
ANIL
MAML
Linear hypernetwork
Nonlinear hypernetwork
Train accuracy
66.45 + 0.23
68.36 1.3
88.32 + 4,83
98.05 = 0.09
Test accuracy
66.17 + 0.06
68.11 = 1.02
88.29 + 5.09
98.0 + 0.05
OOD accuracy
60.78 0.24
55.41 = 0.8
76.35 + 10.11
90.9 0.26
OOD accuracy (K=1)
59.12 + 0.98
53.33 = 1.17
71.32 + 9.97
89.24 = 0.88
0OD accuracy (K=2)
61.38 0.21
55.98 0.53
77.41 = 10.18
90.79 = 0.58
OOD accuracy (K=4)
60.61 + 1.46
50.33 = 1.36
71.96 + 14.18
68.31 = 6.31
Table A4: Hyperteacher comparison of accuracies across models for M=4 K=4.
ANIL
MAML
Linear hypernetwork
Nonlinear hypernetwork
Train accuracy
65.63 + 0.32
68.91 0.28
92.22 + 4.68
97.43 = 0.36
Test accuracy
65.58 + 0.25
68.72 = 0.3
92.29 + 4.66
97.49 + 0.27
OOD accuracy
63.14 + 0.7
63.29 0.6
88.61 + 6.15
94.9 0.23
OOD accuracy (K=1)
59.77 + 1.23
56.29 1.54
82.55 + 8.17
91.79 = 0.52
OOD accuracy (K=2)
62.82 0.9
63.16 0.83
88.67 6.5
95.13 = 0.06
OOD accuracy (K=4)
65.55 + 3.67
66.69 = 3.38
89.0 + 7.52
95.63 = 1.46
Table A5: Hyperteacher comparison of accuracies across models for M=8 K=1.
ANIL
MAML
Linear hypernetwork
Nonlinear hypernetwork
Train accuracy
66.22 + 0.64
65.41
1.16
88.13 + 2.69
95.85 = 0.12
Test accuracy
66.01 + 0.49
65.16 = 1.35
87.91 +: 2.87
95.72 + 0.08
OOD accuracy
55.8 0.26
47.42
2.16
26.67
+ 3.12
26.91 + 0.4
OOD accuracy (K=1)
56.06 0.24
47.68
2.25
26.73 3.12
26.9 + 0.49
0OD accuracy (K=2)
56.9 0.21
49.21
2.35
30.6 3.38
30.45 = 0.52
0OD accuracy (K=4)
55.77 + 0.28
47.63
= 2.14
26.51 + 3.28
26.5 0.75
00D accuracy (K=8)
54.07 0.31
45.77 2.24
21.85 + 2.93
22.81 + 1.3
39

---

Published as a conference paper at ICLR 2024
Table A6: Hyperteacher comparison of accuracies across models for M=8 K-2.
ANII
MAMI.
Linear hypernetwork
Nonlinear hypernetwork
Train accuracy
60.0 + 0.49
61.33 = 0.68
89.9 + 2.77
96.91 = 0.25
Test accuracy
60.07 0.31
61.33 0.45
90.06 + 2.56
96.95 = 0.22
OOD accuracy
59.76 0.27
87.45 + 2.55
96.05 = 0.34
OOD accuracy (K=1)
59.53 0.56
60.76 0.71
86.73 + 2.56
96.01 = 0.53
OOD accuracy (K=2)
39.76
60.48 0.31
87.49 + 2.35
95.9 + 0.39
0OD accuracy (K=4)
0.28
60.01 0.31
82.56 + 1.65
77.79 = 1.32
0OD accuracy (K=8)
0.26
57.99 0.32
60.6 + 6.61
33.08 = 1.94
Table A7: Hyperteacher comparison of accuracies across models for M=8 K=4.
ANIL
MAML
Linear hypernetwork
Nonlinear hypernetwork
Train accuracy
60.9 + 0.44
63.02 0.62
93.64 + 0.18
95.78 = 0.61
Test accuracy
60.64 0.3
62.93 0.37
93.57 + 0.25
95.71 = 0.55
OOD accuracy
60.63 0.21
62.78 0.24
93.5
95.68 = 0.55
OOD accuracy (K=1)
58.77 0.31
58.81 0.38
0.31
92.6 1.47
OOD accuracy (K=2)
59.83
61.39 = 0.45
92.3
95.25 = 0,7
00D) accuracy (K=4)
61.2 0.19
64.0 0.2
93.92 + 0.17
95.88 + 0.51
00I) accuracy (K=8)
62.32 0.26
62.85 0.23
88.18 + 1.03
84.43 + 4.49
Table A8: Hyperteacher comparison of accuracies across models for M=8 K=8.
ANIL
MAML
Linear hypernetwork
Nonlinear hypernetwork
Train accuracy
61.41 + 0.37
63.87 0.38
91.54 0.89
94,83 0.1
Test accuracy
61.41 0.28
64.05 = 0.19
91.56 = 0.71
94.71 = 0.24
OOD accuracy
61.53 0.26
64.11 0.24
91.65 + 0.72
94.75 + 0.24
OOD accuracy (K=1)
57.71 0.1
57.75
0.22
83.5 0.9
89.26 + 0.95
0OD accuracy (K=2)
58.99 + 0.42
59.77
0.56
88.32 1.06
92.97 + 0.23
0OD accuracy (K=4)
61.37 + 0.31
64.12 0.2
91.75 + 0.69
94.82 + 0.15
0OD accuracy (K=8)
63.63 1.52
1.86 66.73
92.04 + 0.54
95.49 + 0.36
40

---

Published as a conference paper at ICLR 2024
Nanlinear hypernework
Linear hypernetwork
128
0.1
Dut-of-distributton loss
Number of moduless
32
16
0.01
128
256
512
128
256
512
Hidden dimension
Hidden dimension
Figure A8: Effect of overparamettrrizaioon on the OOD loss in the compositional preference envi- ronment.
Nonlinear hypernetwork
Linear hypernework
ANIL
MAML
Fraction OOD policy optimal
Fraction OOD OOD policy optimal
0.6
0.6
0.4
0.4
0.2-
0.2
0-
50k
100k
150k
200k
Non Nan.compositional
Compasitional
Training steps
Figure A9: Complementary figure to Figure 4: and reporting additional metric that measures the fraction times the learned policy exactly follows the optimal policy along the whole path and therefore obtains the only environment reward by reaching the goal object. Error bars denote the standard error the mean over seeds.
achieved in the compositional preference environment. Figure A8 shows that within the range of values were able accommodate given the maximum GPU memory available us, overparam- eterization seems to have no negative influence on the OOD loss. This in line with our empirical observations for the continuous multi-task teacher-student setting, preferences for each task continuous combinations of the selected modules.
B.4 COMPOSITIONAL GOALS
Modular architectures outperform monolithic architectures by leveraging compositional structure. Figure 48 and 4F report the out-of-distribution accuracy of the learned policy with respect to the ground-truth optimal policy. In Figure A9 complement it with another metric that measures the fraction of times the learned policy exactly follows the optimal policy along the whole path and therefore obtains the only environment reward reaching the goal object. It is possible that the learned policies only deviate from the optimal trajectory in a way that would also yield the reward which would not count towards this fraction.
Overparameterization the compositional goal environment. Similarly, we investigated the effect of varying the hidden dimension and number of modules in the hypernetwork models on the OOD accuracy achieved in the compositional goals environment. Different from the compositional preference environment, teacher-student student setting OOD performance was sensitive to overparameterization, Figure Alloshows relatively stable OOD accuracy across varying hidden and module dimension.
have discrete set of goals in this case. While in the discrete multi-tate
41

---

Published as a conference paper at ICLR 2024
Linear hypernetwork
Nonlinear hypernetwork
Out-of-distributiin accuracy
32
Number
0.6-
0.4
128
255
128
256
Hidden dimension
Hidden dimension
Figure A10: Effect overparameterizatioo on the OOD accuracy in the compositional goals envi- ronment.
C EXPERIMENTAL DETAILS
C.I MULTI-TASK TEACHER-STURE STUDENT SETUP
C.I.1 DATA GENERATION
We generate data by first initializing the teacher parameters once. Then, for each task, we sample a task latent variable z which induces a noiseless mapping from inputs æ to targets y following equation 3 modulo rescaling factor by the square root of the input dimension
y = avp(
Vn W
(80)
For all experiments in the multi-task teacher student, we fix the activation function v to the ReLU nonlinearity, the task latent variable dimension to M 6, input dimension to 16, hidden dimension of the teacher to h 16, and output dimension to = 4. The teacher parameters θ and are generated by sampling the entries i.i.d. from the centered truncated normal distribution, with standard deviation resp. VM and J. We define the distribution over inputs to be the uniform distribution with mean and standard deviation on RR. Finally, we specify the distribution over task latent variable z. We distinguish 2 settings:
Continuous task distribution. Here, we consider tasks where modules are sparsely, linearly com- bined. A task distribution is specified by set of masks, that are binary vectors of R M Given mask, we sample a task z as follows. We first sample an M-dimensional random variable following the exponential distribution. Then, zero-out entries corresponding to the mask being O. We then normalize the vector such that the sum equals 1. This procedure simulates the uniform sampling from the simplex spanning the directions in which the mask is non-zero. Finally, add the mask to the vector and rescale the outcome by 0.5. This ensures two tasks generated by distinct masks do not have intersecting support (but intersecting span). See Algorithm for the pseudocode.
Algorithm Algorithm to sample the tast latent variable from given a mask.
Require: mask m of size M return task latent variable z sample M-dimensional vector from the exponential distribution f- C m
x +- 1+|
+ 0.5 (z m)
return z
The task distribution is then generated as follows: first, a mask is sampled randomly and uniformly from the pre-specified set. Then, the vector is sampled following the above procedure.
42

---

Published as a conference paper at ICLR 2024
Table A9: Training task distribution used for our experiments. Each column defines a distribution. For discrete settings, the set of vectors specify those in the training set. We omit the normalizing factor for simplicity. For continuous settings, the vectors specify the set of masks used in the training set to generate the distribution.
Discrete
Continuous
Comnected
Disconnected
Connected
Disconnected
(1,0,0.0.0,0)
(1,1,0,0,0,0) (1,0,0,0,0,0)
(1,1,1,0,0,0)
(1,0,0,0,0.0) (1,1,0,0,0,0)
(1.1,1.0,0,0)
(0,1,1,0,0,0) (0.100,0,0.0) (0,1,0,0,0.0)
(0.1,0,0.0.0) (0.1,1,0,0.0) (0,0,1.1,1.0)
(0,0,0,1.1.I)
(0.0,1,1.0.0) (0.0.1.0.0.0) (0,0,1,0.0.0)
(1,0,0,0.L.))
(1.0.1,0,0.0) (0.0.1,0.0.0) (1,1,0,0,000)
(0,0.0.1,0.0) (0,0,0.1,0.0))
(0.0,0,1,1,0)
(0,0,0,1,1.0)
(0,0,0,1.0.0) (0,1,1,0,0.0)
(0.0,0.0.1.0) (0,0,0.0.1.0)
(0,0,0.0.1.1)
(0,0,0,0,1,1)
(0.0.0.0.1.0) (1,0.1,0,0,0)
(0,0,0,0,0.1) (0,0,0.0,0.I)
(1,0,0,0.0.1)
(0,0,0,1,0.1)
(0,0,0,0.0.1) (0,0,0,1,1.0)
(1.1,0,0.0.0) (1.1.0.0.0.0)
(0.0,0.0.1,I)
(0.1,1,0.0,0) (0.1,1,0,0.0)
(0.0,0.1.0.I)
(0,0.1.1,0.0) (1,0,1,0,0.0)
(0,0.0.1.110)) (0,0,0.1.1.0))
(0.0,0.0.1.1)) (0.0.0.0,1.1)
(1,0.0.0.0.1) (0.0.0.1.0.1)
Discrete task distribution. Here, we focus on task latent variables that are simple normalized many-hot vectors, For example, for M = 4, we would consider vectors such as (0.0, D, O, 1), v2 J. (0, I, 0, 1) or 1 (0, 1). By only considering such variables, we wish to model tasks where 73 modules are combined sparsely without varying the magnitudes with which modules enter the com- position. There are finite number of such vectors, and thus the task distribution is a uniform mixture of diracs. We define the training task distribution by manually specifying which of these vectors are in the distribution.
C.1.2 STUDENT MODEL
We use the same parameterization and parameter initialization scheme for the student as for the teacher, but we vary the hidden layer width and inferred task latent variable dimension M.
C.1.3 TRAINING AND EVALUATION
Algorithm 2 Bilevel training procedure
Require:
outer- outer-batch size Bouters inner batch-size Bunner number of outer steps Nower, number of inner steps Nimerr outer learning rate Nouters inner learning rate Minner
for Nauter iterations do
Sample Bowter task latent (z+)
AO 0
For k do
① x 60
for Nuner iterations do
Sample Binner data ((xa, yi from Druppon(22)
@ + Φ - V/inner L , Φ, ((ri, y:));)
end for
Sample B Binner data ((I,, y:)) from Davery AO + -0 + VoL Φ, ((z4, ()))))
end for
A Notter A0
end for
Since there can be an infinite amount of distinct tasks for continuous task distributions, we adopt training procedure that allows for training in a multi-task setting with infinite tasks. For this purpose, we take the standard algorithm to optimize the cross-validation bilevel problem, detailed in Algorithm 2 and adapt it to the infinite-data multi-task case, using the fast parameter the
43

---

does not involve any second-order gradients and can be com- Published as a conference paper at ICLR 2024
inferred task latent variable & and the meta parameters 0 the remainder of the student parameters, θ, á.
Concretely, to stay close to the theory, we use the whole task distribution for both support and query set, in particular i.e. Deupoot Diverry (z In this case, equilibrium (over the distribution), we know that the total derivative puted efficiently with partial derivatives. Therefore we ensure the inner loop converges by allowing for a large number of inner steps Nimer- and apply the stop-grad operation on & before computing the gradient w.r.t 0.
While typically the fast parameter initialization Φo can also be learned, here for simplicity we ran- domly sample the initial value from a normal distribution and keep it fixed throughout training.
In order to measure out-of-distribution performance at the end of training, we wish to sample new task latent vectors z involving unseen combinations of modules, and evaluate the ability of the student to fit the corresponding data. We follow the same inner loop procedure as during training, and report the average outer loss over the task distribution, as well as the module alignment metric described in Section 3.3
C.1.4 EXPERIMENTS
We conduct 2 experiments, corresponding to Figure and C.
In the first experiment (Figure 2B), we wish to investigate the effect of the connected support on module identification. We manually select task distributions that have, or do not have the connected support property, for both continuous and discrete task distributions. See Table A9 for the chosen
tasks.
In the second experiment (Figure 2C), investigate the effect of overparameterization identifia- bility when the task distribution has connected and compositional support. For both continuous and discrete distributions, we the task corresponding the first column of the appropriate section in AA
The training loss all experiments reached a value lower than 10-?
C.1.5 HYPERPARAMETER:
For all experiments in the multi-task teacher student, we set Bower 64, Binner = 256, Nouter 60000, Nimner 300. We optimize the inner loop using the Adam optimizer with Vinner 0.003, and the outer loop using the AdamW optimizer with various weight decay strengths and an initial learning rate Nowser 0.001 annealed using cosine annealing down to 10 at the end of training. We observed that this set of hyperparameters gave enough time for the student to practically reach ( loss, which is the setting studied in the theory.
C.2 HYPERTFACHER
C.2.1 DATA GENERATION
We generate data by first initializing the teacher parameters once. The teacher network is linear hypernetwork for which the parameterization and initialization scheme are described Section [ For all our experiments, unless specified otherwise, we fix the input dimension to = 16. hidden dimension of the teacher to 32, and output dimension to o We define the distribution over inputs æ to be the multivariate uniform distribution over the hypercube 1]".
Targets are generated by passing the input æ through the generated network. To make sure the distribution of the outputs is not too dissimilar across tasks, normalize each output neuron to have unit variance and zero mean by numerically estimating the first and second moment of the distribution over random inputs. The tasks sampled following the same procedure as in the discrete task distribution description in Section C.I.I That is, a finite set of binary vectors of size M fully specifies the task distribution.
44

---

Published as a conference paper at ICLR 2024
hidden layer
2 hidden layers
32
Number of teacher modules
Number of tasks
10k
100
1-to-K hot
1-to-K hot
Figure All: Number of possible tasks in the hyperteacher as function of the number of teacher modules M, the number of sparse combinations K and number of hidden layers L given by
Σ K M L-1 growing exponentially in M.
C.2.2 TRAINING AND EVALUATION
We train the models following the bilevel cross-validation objective introduced in Section Specif- ically, we follow the algorithm outlined in The loss function for both inner and outer loop the KL divergence between the log-softmax of the target and the log-softmax of the prediction of the network.
For each task, we sample ii.d 256 input-output pairs for each of the support and query set. The evaluation follows the same procedure and inner optimization, except for the fact that it involves tasks that combination of modules that have not been seen during training. We report the outer KL divergence averaged over Bower tasks, as well as the accuracy, which measures the average over tasks of fraction of input for which the index of the largest value of the prediction matches that of the largest value of the target.
C.2.3 EXPERIMENTS
Here we provide additional details for the experiments presented in Figure
Compositional vs non-compesitional support. We perform this experiment for different values of K. Given K, we compute the set of binary masks with up to non-O entries, and randomly select three-quarters of them and use the obtained list for generating training tasks. We make sure that the obtained list has compositional support, that is, all modules appear at least once. The remaining quarter is used as the OOD task for evaluation. For conpooopoiiii support, all binary masks with up to non-O entries, but split it into sets such that for training, only masks used, while for the evaluation, masks which always have the
again enumerate
where the last module is inactive last module activated, are used to generate the tasks.
Connected vs disconnected support. For these experiments, we consider M = 8.
For connected task support, we used the task distribution defined by the set of masks {e: (e:isisM {e: + Cs+1 mod A hgsga U {e, + 6:+2 mod M hisism where (c:)svgM represent the canonical basis of R For disconnected task support, we used the task distribution defined by the set of masks (ei U (er + e))..jj(....M/
(e+ + J Ml.ijj These tasks were chosen such that the number of tasks for the connected and disconnected experiment are comparable. For the OOD task distribution used for evaluation, we used all binary masks with K that have been neither used by the connected nor disconnected task distribution.
Other experiments. For the other experiments, generate tasks following the same split for the compositional experiment. For the overparameterizztion experiment, we further set the teacher hidden layer width to h 16.
45

---

Published as a conference paper at ICLR 2024
Table A10: Hyperparameter grid-search space for the hyperteacher. Best found parameters are
marked in bold.
Hyperparameter
Nonlinear hypernet.
Linear hypernet.
ANIL
MAML
1r.outer 1r ir.inner
{0.001, 0.003)
{0.001. 0.003)
(0.001, 0.003)
(0.001, 0.003)
{0.1. 0.3}
{0.1. 0.3)
{0.01. 0.033
0.001, 0.03}
grad clip
(1.2) (0.01, 0.0001)
[.2]
(1,2)
[1,2]
we weight_decaa
(0.01, 0.00013
(0.01, 0.0001)
{0.01, 0.0001
C.2.4 HYPERPARAMETEAS
For all hyperteacher experiments, we set Bouter = 128, Nouter = 200000 and use full-batch inner optimization. For all models, we optimize both the inner and outer loop with the AdamW optimizer. Unless specified otherwise, we set the following architectural hyperparameters, ensuring the total number parameters across models comparable.
Linear Hypernetwork: base network hidden layer L = 3, hidden units à 128. Number of modules M = 4M
Nonlinear Hypernetwork: base network hidden layer 1. 3, hidden units = 128. Num- ber of modules M
4M
MAML: network hidden layer I, = 3, hidden units h = 368
ANIL: network hidden layer = 3, hidden units h = 512
We fix the inner loop steps Nineer 10 for all models, except for ANIL where Nimer 100. For all models, we tune the hyperparameters independently for the inner learning rate, the outer learning rate, outer weight decay well as the outer gradient clipping on the range specified in Table(A1)) The hyperparametens are tuned to optimize for the compositional generalization loss on the compositional support experiment detailed above. The same set of hyperparameters were then used for all other experiments.
C.3 COMPOSITIONAL PREFERENCES
C.3.1 DATA GENERATION
Environment. We consider grid world of size 5
in Figure 4A. There are actions:
move up, right, down, left, and terminate. Black cells represent walls. The agent deterministically moves from one cell another following the direction if the grid is wall(or border)-free, and oth- erwise remains in place. The episode terminates when the terminate action is taken. Finally, there are 4 objects placed on wall-free cells of the grid. Each object has exactly one of the colors. Each object remains static throughout the episode, but disappears when the agent steps on it, potentially producing a reward in task-specific manner, which will explained below. The state space then the grid, which shows the positions of the walls, the objects, their respective feature, and the position of the agent, as in Figure 4A.
Task. A task is defined by task latent variable z, of dimension M 8. Such variable defines the reward function of the environment in the following way: first, a preference vector is computed by linearly combining set of M pre-computed preference template vectors of dimension 8, using the task latent variable as the coefficient. The resulting preference vector specifies mapping from the 8 colors to a preference value. During the task, when the agent steps on object, it gets a reward equal to the preference value of that color.
For each task, several instances of the environment are then created. While the wall configuration remains fixed, for every new instance of the environment, we sample the position of all objects uniformly from all wall-free cells (making sure there is no overlap), as well as the initial position compute the optimal action-value function with discount factor
of the agent. For each instance, 0.9 and time horizon 8. We then let an agent behave greedily following the obtained action value function (forcing the agent take the terminate action when positively rewarding objects are left
46

---

Published as a conference paper at ICLR 2024
in the grid), and build the task dataset ((z) ((z2, 9s)), by collecting the observations along the trajectory, and the corresponding action value vector y, over all instances. Finally, tasks are sampled following the same procedure as in the continuous task distribution description in Section C.I.1 A finite set of binary masks, each of size M, thus fully specifies the task distribution.
C.3.2 TRAINING AND EVALUATION.
We train the models following the bilevel cross-validation objective introduced in Section Specifi- cally, we follow the algorithm 2 The loss function for both inner and outer loop the Mean Squared Error (MSE) loss between the action values and the prediction of the network.
For each task, we instantiate 32 different instances, 16 of which being used for the support set and the remaining 16 for the query set, and create the dataset following the procedure outlined in the previous section. The evaluation follows the same procedure and inner optimization, except for the fact that it involves tasks that are combinations of modules that have not been during training. We report the outer MSE loss averaged over Bouter tasks.
C.3.3 EXPERIMENTS
We detail the experiments corresponding to Figure 4 and C.
Connected vs disconnected support. For connected task support, we used the task distribution defined by the set of mask {e) U (e, + e++1 mod MIsIgM fe, + 64+2 mod M IS where (e:)ssigM represent the canonical basis of RRM For disconnected task support, we used the task distribution defined by the set of mask {eihisiM U {ei ej),jeen....//)..d. U {es + ey eyh,, h. Mjj. These tasks were chosen such that the number of tasks for the connected and disconnected experiment are comparable. For the OOD task distribution used for evaluation, we used all binary masks with K = 2 that have been neither used by the connected nor disconnected task distribution.
Compositional vs non-compositional support. For compositional support, we compute the set of binary masks with up to 3 non-O entries, and randomly select three-quarter of them and use the obtained list for generating training tasks. We make sure that the obtained list contain a composi- tional support. The remaining quarter is used as the OOD task for evaluation. For non-compositional support, we enumerate all binary masks with up to 3 non-O entries, and split it into sets following the value of the last entry of the mask. For training, we use the set of mask where the last module is inactive, while for the evaluation, we use the set of mask which always has the last module acti- vated, to generate the tasks. Again, these settings were chosen such that the number of tasks for the compositional and non-compositional experiment are comparable.
C.3.4 INTERPARAMETERR
For all experiments, we set Bouter 128, full-batch inner optimization, Nouter 100000. For all models, we optimize the inner loop with Adam optimizer with the default weight decay of 0.00001, and the outer loop with AdamW optimizer with tuned weight decay. Unless specified oth- erwise, we set the following architectural hyperparameters to ensure the total number of parameters across models comparable.
Linear Hypernetwork: base network hidden layer = 3, hidden units h = 64. Number of modules M = 32
Nonlinear Hypernetwork: base network hidden layer L = 2, hidden units = 64. Number of modules M = 32
MAML: network hidden layer L 3, hidden units h 368
ANIL: network hidden layer L = 4, hidden units h = 512
We fix the inner loop steps Nuner to 10 for all models, except for ANIL where Nimer = 100. For all models, tune the hyperparameters independently for the inner learning rate, the outer learning
47

---

Published as a conference paper at ICLR 2024
Table All: Hyperparameter grid-search space for the compositional preference environment. Best found parameters are marked in bold.
Hyperparameter Nonlinear hypernet.
Linear hypernet.
ANIL
MAML
1r.outer
(0.0003. 0.001}
(0.0003, 0.001]
(0.0003, 0.001}
{0.0003, 0.001] 0.01,, 0.03) (1.2)
1r innin
0.31.
00.1 0.3)
0..11. 0.03}
grad clp
(1.2)
[1,2]
(1,2)
rate well the outer gradient clipping from the range specified in table A11 The hyperparam- eters are tuned to optimize for the compositional generalization loss on the compositional support experiment detailed above. The same set of hyperparameters were then used for all other experi- ments.
C.4 COMPOSITIONAL GOALS
C.4.1 DATA GENERATION
For the compositional goals environment we consider mazes of size 11 11 with the same movement dynamics as described above the compositional preference environment. In addition moving down, left and right, the agent can choose one of the 'object interaction" actions. Each environ- ment has randomly placed objects, one of which the target object. The agent is placed randomly and has to reach the target object followed by performing the target interaction to obtain a reward. Each task defined by a goal vector that specifies which of the 5 maze configurations used, which of the 5 object types is assigned as the target object, which of the 2 possible target interactions is the correct one and finally in which the 4 quadrants the maze the target object will be located (c.f. Figure P.)
Goals are compositional in the sense that any of these factors can be arbitrarily combined leaving a total of 5-5- 2 200 possible goals. Given goal, first sample agent and all objects, compute the optimal behavior policy and use it to draw sample demonstration of the optimal policy for the support set. We then sample random configuration of the objects and the agent given the same goal for the query set.
random position of the
In the setting of compositional support, randomly hold-out 25% of the possible goals ensuring that every goal factor appears at least in one of the goals of the remaining 75% of goals used for training.
To produce training distribution with non-compositional support, we consistently hold out one of the goal quadrants from the of goals and use all goals that contain the held-out goal quadrant for the OOD set.
C.4.2 TRAINING AND EVALUATION
We train the models following the bilevel cross-validation objective introduced in Section [ Specif- ically, we follow the algorithm outlined in Algorithm? 2 The loss function is the cross-entropy loss between the optimal action and the prediction the network.
For each task, we sample a goal, instantiate two environments consistent with the goal and draw samples containing the optimal trajectory from the agent location to the goal for both. We use the demonstration of one instantiation of the environment for the support set and the second one for the query set.
For all experiments, use full batch optimization in the inner loop, and fix the outer batch size to 128, i.e. Bouter = 128.
The evaluation follows the same procedure the inner optimization, except for the fact that it involves tasks based on goals that have not been seen during training. We report the outer loss averaged over 1024 tasks.
48

---

Published as a conference paper at ICLR 2024
Table A12: Hyperparameter grid-search space for the compositional goal environment. Best found parameters are marked in bold.
Hyperparameter Nonlinear hypernet.
Linear hypernet.
ANIL
MAML
1r.outer
(0.001, 0.003}
{0.001. 0.003)
0.0001, 0.003)
(0.001, 0.003)
ir innin
10.1..33
0..1.033
0.01, 0.03}
0.011 0.03)
grad.cli
[1.2)
[1,2]
(1,2)
(1,2]
C.4.3 HYPERPARAMETERS
For all experiments, we set Bouee 128, Bimer 256, Nouter 200000. We optimize both the inner and outer loop with the AdamW optimizer. For the outer loop we use a weight decay of 0.001 for the inner loop we use weight decay of 0.0001. Unless specified otherwise, we set the following architectural hyperparameters, ensuring the total number of parameters across models is approximately comparable.
Linear Hypernetwork: base network hidden layer = 2, hidden units h = 32, Number of modules
Nonlinear Hypernetwork: base network hidden layer L. = 2, hidden units = 32. Number of modules M
8
MAML: network hidden layer L hidden units h 384.
ANIL: network hidden layer L 2, hidden units h = 512.
We fix the inner loop steps Numer to 10 for all models, except for ANIL where Ninerr = 100. For all models, we tune the hyperparameters independently for the inner learning rate, the outer learning rate as well as the outer gradient clipping from the range specified in Table A12 The hyperparam- eters are tuned to optimize for the compositional generalization loss on the compositional support experiment detailed above, for all models. The same set of hyperparameters were then used for all other experiments.
D META-LEARNING MODELS
In this section provide details of the models train in meta-learning experiments. We the same architectures across experiments and settings.
We keep the notation of 0 and ø to respectively denote meta- pett-parrmeters and fast hast-parameters, used in Section 2 and Algorithm
D. MAML
D.1.1 PARAMETERIZATTION
We consider a standard ReLU MLP with L hidden layers of width h, with bias. During meta learning, all parameters of the network are adapted and thus part of the fast parameter é. The only task-shared parameter is thus the value at which we initialize the fast parameters at the beginning of each task, i.e. 0 do-
D.1.2 INITIALIZATION
We initialize d such that parameters corresponding to weights are initialized following the truncated normal distribution with mean 0 and standard deviation VIT where H is the size of the input to the layer. Parameters corresponding to biases initialized
49

---

Published as a conference paper at ICLR 2024
D.2 ANIL
D.2.1 PARAMETERIZATION
We consider a standard ReLU MLP with L hidden layers of width h, with bias.
Here, during meta learning, only the readout layer of the network is part of the fast parameters & The parameters for the rest of the network, well the value do which initialize the fast parameters at the beginning of each task are part of the task-shared parameters 0.
D.2.2 INITIALIZATION
We initialize 0 and such that parameters corresponding weights are initialized following the truncated normal distribution with mean 0 and standard deviation сн input to the layer. Parameters corresponding to biases are initialized at O.
- where H is the size of the
D.3 LINEAR HYPERNETWORK
D.3.1 PARAMETERIZATION
We consider a base network consisting in a ReLU MLP, with hidden layers of width Crucially, the NTK parameterization of the MIP, where at each layer, the pre-activation units rescaled by /H where H is the dimension of the input to the layer. As such, at initialization the hypernetwork outputs weights that are centered and unit variance, the forward pass would be approximately variance preserving.
The weights are generated by the linear hypernetwork in the following way:
First, given an embedding the linear hypernetwork normalizes the embedding to the unit norm.
Then, the normalized embedding is linearly multiplied by the parameter of the hypernet- work θ to generate the flattened weights
Finally, the flattened weight are reshaped and plugged into the base network
The biases not generated by the hypernetwork.
Crucially, we use 3 hypernetworks to generate respectively the first layer weight, all hidden layer weights, and the last layer weight. There therefore 3 embeddings in total, one for each hyperntt- work.
When training the model in the meta-leaming setting, we as fast parameters the of embed- ding vectors as well as the biases in the base network. The meta parameter 0 are the hypernetwork parameters, well as the initialization value for the fast parameters.
D.3.2 INTIIALIZATION
The hypernetwork weights are initialized following the truncated normal distribution with mean and standard deviation +
VM where M is the embedding vector dimension.
The meta- parameters corresponding to the initialization values of the embeddings are initialized following uniform and unit variance distribution, and those corresponding to biases are initialized to O.
D.4
NONLINEAR HYPERNETWORK
D.4.1 PARAMETERIZATTIO
The parameterization of the nonlinear hypernetwork closely follows that of the linear counterpart (see Section D.3.I Instead of using a linear layer to produce the weights the base network given the fast parameters ø, the fast parameters are first nonlinearly transformed in ML.P.
50

---

Published as a conference paper at ICLR 2024
Specifically, the hypernetwork is 4-layer MLP with the ELU activation function applied every hidden layer. Layer normalization is applied before each activation function, as well as to the input.
When training the model in the meta- meta-learnng setting, we as fast parameters ø the set of embed- ding vectors as well as the biases in the base network. The meta-parameters are the hypernetwork parameters, well the initialization values for the fast parameters.
D.4.2 INITIALIZATION
The hypernetwork parameters are initialized such that all biases are set to and MLP layer weights are initialized following the truncated normal distribution with mean O and standard deviation VH where H is the size of the input to the layer.
The meta parammere corresponding to the initialization value the embeddings are initialized fol- lowing a uniform and unit variance distribution, and those corresponding biases initialized to 0.
E EXTENDED RELATED WORK
We complement the related work discussed in the main text with a more comprehensive review in the following.
The debate on the aptitude of neural networks for compositionality dates back to the first wave of connectionist models (Rumelhart & McClelland 1986; Fodor & Pylyshyn 1988 Smolensky 1991 Hadley 1994, Phillips) 1995) fueled by the proposition that combinatorial structures are a key property of human cognitive abilities. The widespread success of deep learning has lead many efforts to evaluate the capacity of deep networks for compositional generalization.
A number of benchmarks have been designed specifically investigate the ability for compositional generalization primarily in the context of language (Bastings et al. 2018
1018:Lak Lake & Baroni 22018] 8 Kim
& Linzen 2020 Ruis et al. [2020] Hupkes et al. (0220) Keysers
2020 Dankers al. 2022
sual question answering (Johnson et 2011 BBBdanau et al. (020). visual scenes (Xu et al. 2022)
and reinforcement learning (Barreto et al.) 2020] IZhao et al.] 2022). Commonly, these benchmarks
demonstrate shortcomings of current deep networks at achieving compositional generalization from
recurrent neural networks (Lake Baroni 2018 Loula et al. 2018), convolutional networks
(Dessi & Baroni 2019) and transformers [Keysers et al. 2020).
large number of architectures have been designed with compositionality in mind. A common theme among them is to attempt to decompose skills reusable atoms of knowledge, across tasks in order to recombine them in novel ways for new tasks (Alet et al. 2018 8 Ponti et al. 2022] Kingetsu et al. 2021). Compositional plan vectors (Devin al.) 2020 learn composable embeddings to encode different tasks and condition a policy. Similarly, Frady et al. (2023) use composable vec- tor symbols to represent compositional visual scenes given latent supervision. To impart strong bias for compositionality. some works aim to directly learn symbolic expressions (Liu et al.) 2020 Vankov Bowers 2020). Similar our setup, meta seq2seq (Lake) (019)) conditions a sequence model on demonstrations to facilitate compositional generalization. Building modular systems that exhibit compositional generalization is also main goal of graph neural networks (Battaglia al. 2018). In this context compositional structure is typically built into the system while here we attempt discover this structure from unstructured observations. Complementary this, Ghazi et al. (2019) have proposed framework to make the function implemented by individual modules interpretable.
With the advent of pretrained large language models, compositional abilities have seemingly come into reach (Zhou et al., 2023 Orhan 2022 Furrer et al.) 2021] Csordás et al.) 2021). Many spe- cialized architectures that show compositional behavior on one the aforementioned benchmarks (Russin et al. 2019; et al. 2219 Gordon al.) 2020 Andreas 2020 Liu et al., 2020 Lake 2019)
Nye et al., 2020 Kaiser Sutskever] 2016) are outperformed by pretrained large language models (Furrer et al., 2021). However, it remains open question whether the ability for compositional
generalization extends beyond the training distribution (Srivastava et al. 2023 Press et al. 2023
Dziri et al. 2023).
51

---

Published as a conference paper at ICLR 2024
F ADDITIONAL DETAILS
F.1 COMPUTE RESOURCES
We used Linux workstations with Nvidia RTX 2080 and Nvidia RTX 3090 GPUs for development and conducted hyperparameter searches and experiments using 5 TPUv2-8, TPUv3-8 and 1 Linux server with Nvidia RTX 3090 GPUs over the course of months. In total, we spent an estimated amount of 6 GPU months.
F.2 SOFTWARE AND LIBRARIES
For the results produced in this paper we relied on free and open-source software. We implemented experiments in Python using JAX (Bradbury et al. 2018 Apache License 2.0) and the Deepmind Jax Ecosystem (Babuschkin et al. 2020 Apache License 2.0). For experiment tracking we used wandb (Biewald 2020 MIT license) and for the generation of plots we used plotly (Inc 2013 MIT license).
52

---

Published as a conference paper at ICLR 2024
53


Reviews:
[{"review_id": "8hqkGa1Q2A", "reviewer": "Reviewer_4knX", "summary": "This paper studies the problem of compositional generalization in modular architectures. The authors show that in the teacher-student setting, it is possible to identify the underlying modules up to linear transformations purely from demonstrations. They further show that meta-learning from finite data can discover modular solutions that generalize compositionally in modular but not monolithic architectures. The authors also demonstrate how modularity implemented by hypernetworks allows discovering compositional behavior policies and action-value functions.", "questions": "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?\n2. Could you explain why the network modules have to be hypernetworks in your setting?\n\n\n-----------\nupdate: raised to 6.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "1. The paper theoretically shows that students can learn the underlying modules from the teachers under certain conditions.\n2. The results are supported by empirical experiments.", "weaknesses": "1. 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.\n2. There is no section that explicitly discusses related work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WV2t2Q0rMB", "reviewer": "Reviewer_bgW6", "summary": "This work explores the ability of hyper-networks to detect the ground-truth task generating functions and have modules (rows of the linear hyper-network weights) specialise to these underlying factors of variation. Moreover these modules can then be composed in previously unused combinations to generalise to new tasks. Experiments demonstrating the benefit of modular meta-learning algorithms over monolithic meta-learners are also shown.", "questions": "I have raised a number of questions in my discussion above. I think the only outstanding question or point of clarification I have at this point is the following:\nCould 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.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "## Originality\nThis paper uses an established technique, the teacher-student setup, to understand the necessary conditions for a dataset to promote systematic generalisation. Moreover, a connection between the teach-student setup and meta-learners is drawn, which I have not seen been made frequently before. Also considering ANIL, MAML and hyper-networks as the class of monolithic and modular meta-learners is interesting a new. Thus, the combination of concepts, theoretical technique and models considered in this work is  new.\n\n## Quality\nThe theoretical setup of this work seems appropriate for addressing the main concerns of this work - under what dataset conditions do meta-learners naturally generalise compositionally. The overall structure and logical setup of the work is good and sections naturally lead from one to the next. Section 3.3 directly tests the theory which is presented and in the case of Discrete task distributions the experiments support the theoretical findings. As far as I can tell the assumptions of the setup are clearly stated and the setup is clearly defined.\n\n## Clarity\nFigures are visually clear and well designed.\n\n## Significance\nI reiterate here that I think this work uses a very interesting combination of previous ideas. As a result the reported findings are interesting and I could certainly see the results leading to future work and guiding research on meta learning. Thus, I think this work does have the potential for high significance. I would say this is contingent on some of the weaknesses discussed below being mitigated.", "weaknesses": "I have two broad concerns for this work. One is on clarity and the other is on quality. I will begin with clarity as I think this may be a factor in the concerns on quality.\n\n## Clarity\nOverall I found this work to be relatively unclear. Notation is used but not properly introduced, for example $(U_k)$ where it is not mentioned what the double subscript $k$ refers to, or why the second subscript is necessary. Another example is where it is said $\\forall k,l$ where it is necessary to infer from context what $k,l$ is referring to. Also where it says $U_{k_i}$ what does the $i$ or $k$ refer to and why is this necessary. Similarly in Theorem 3.1 where it says \"The if $\\mathcal{L}(W_2,\\Theta)$...\", what is this new $\\mathcal{L}$ referring to? Is it the loss and this is just a mistake in the font? If so, why does this loss function not accept the same  parameters as previous loss functions. Finally, the notation introduced in the paragraph beginning \"We can now present our main result\" is particularly confusing because the superscript $(i)$ is overloaded three times, once referring to a row, the other referring to a column and the third for a full matrix sliced out from a tensor.\n\nSecondly, Definitions 3.1,3.2 and 3.3 are not clear and no intuition or interpretation is provided. For example, where defining irreducibility, the fact that all rows of $W_1$ 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 $W_2$ to be $0$ is also not mentioned. Why would a full column of $W_2$ be $0$? Similarly, \"Compositional Support\" seems to just be saying that $U$ is a basis of $\\mathbf{R}^M$ 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? Definition 3.3 is just extremely difficult to parse in general and so is Theorem 3.1. This is due to a lot of terminology being mentioned without definition and needing to be inferred based on context. For example, what is $\\hat{M}$ and why is Theorem 3.1 making a distinction between even and odd values of $n$? The difficulty in following these definitions alone makes the rest of the work difficult to follow.\n\nMy final point on clarity is that the figures, while visually neat and well done, are vague and their captions unhelpful. This is particularly bad when the figures are relied on heavily to explain concepts. 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. Essentially, every figure caption should be elaborated on and potentially more information be placed in the figures to depict what is actually going on. 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)$.\n\n## Quality\nMy concerns on quality are likely due to misunderstanding from the above. I would like to reiterate that I do believe this work has potential significance. I am, however, struggling to connect this work to the general literature on compositional generalisation. For example, assuming that $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? 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.\n\nThe 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)?\n\nI would be open to increasing my score quite substantially if it is shown that I have indeed missed something crucial. Alternatively if the clarity issues are addressed and the connection to prior work made more explicit I would also increase my score. Likely to a 7. I would certainly prioritise improving the clarity, with the figures being particularly low hanging fruit which would make quite a big difference if improved upon.\n\n## Minor Points\n1. \"with each row representing one parameter module\" - I believe I understood what you were saying, but this was not an easily understandable sentence.\n\n[1] Hadley, Robert F. \"Systematicity in connectionist language learning.\" Mind & Language 9.3 (1994): 247-272.\n[2] Ruis, Laura, et al. \"A benchmark for systematic generalization in grounded language understanding.\" Advances in neural information processing systems 33 (2020): 19861-19872.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7zKiTW0Px3", "reviewer": "Reviewer_9UrZ", "summary": "This paper proposes a multi-task teacher-student approach with modular architecture for compositional generalization.\nIt uses hypernetworks to convert the generalization to the identification of modules and theoretically show the result up to linear transformation.\nExperiments also support the ability.", "questions": "Please respond to the weakness section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- It focuses on an important question of how to learn modular structure for compositional generalization.\n- The hypernetwork and modularity approach cast the generalization problem into an identification\nproblem.\n- Both theory and experiments support the result.\n- It proposes connected support to address permutation invariance.", "weaknesses": "My largest concern is that the framework is very constrained, and some constraints, e.g., linearity, may be essential to deriving the results.\nIt indicates that there may be difficulty when generalizing the result to more complex situations.\n\n(1) Linear assumption in hypernetwork.\n\n(2) It uses two two-layer neural networks.\n\n(3) It assumes knowing the correct (teacher) architecture.\n\n(4) The theory assumes knowing the number of modules (M) and hidden units (h).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ERGOKlFual", "reviewer": "Reviewer_KGfv", "summary": "The authors provide a theoretical analysis of the compositional generalization capabilities of linear hypernetworks under a teacher-student setup where the teacher network is the ground-truth data-generating model, and is itself a linear hypernetwork. The key result is that under 3 assumptions (compositional support, connected support, and no over-parameterization), the student can provably _identify_ the modules of the teacher. This identification is further demonstrated to be a necessary and sufficient condition for compositional generalization. The authors provide a series of experiments to validate their theoretical findings and test how sensitive the empirical results are to various violations of the assumptions.\n\n######## Post-rebuttal update ########\n\nGiven the authors' response and updated manuscript, I am updating my score to an 8 (this 8 should be considered closer to a 7, but 7 is not an option).", "questions": "######## Additional feedback ########\n\nThe following points are provided as feedback to hopefully help better shape the submitted manuscript, but did not impact my recommendation in a major way.\n\nAbstract\n- 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.\n\nSec 3\n- Up to this point, it was still unclear to me what the authors meant by identification of modules\n    - Is it that the student can determine which modules to use, or that it can find the right set of modules?\n    - Is this in a setting where the student knows the latent codes?\n    - These are later clarified, but it might be worth doing so earlier on\n\nSec 5\n- It's odd that the authors combined a related work section with the discussion, but I think it works okay. \n\nTypos/style/grammar\n- Intro, page 1 first paragraph: series of tasks --> set of tasks [a series is sequential]", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "######## Strengths ########\n- The problem of compositional generalization is of great interest to the ML/AI communities, and there are very few theoretical treatments of the problem\n- The combination of theoretical results and empirical evidence that tests beyond the bounds of the theory is well balanced", "weaknesses": "######## Weaknesses ########\n- The experimental setting is unclear, especially the sequential decision-making portion, which makes it hard to assess its impact\n- It is unclear whether the three assumptions (compositional support, connected support, and no over-parameterization) are necessary conditions or just sufficient\n- The choice of linear hypernetworks as the base modular model somewhat limits the intepretability of the compositionality\n\n######## Recommendation ########\n\nI perceive this work as borderline acceptable. The balance of strengths and weaknesses, I believe, works out in favor of the strengths slightly. That being said, I encourage the authors to address my comments below to improve their manuscript.\n\n######## Arguments ########\n\nThere 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). The two works study very distinct settings, and so I do believe that this submission proposes a novel piece of work. Moreover the authors execute a fairly comprehensive experimental evaluation of their ideas, which is rare for theoretical works. To me, this merits acceptance of the work.\n\nHowever, there are a few shortcomings that the authors should address.\n- Clarity of the experimental setting\n    - Sec 4.1\n        - How is compositional generalization measured? Does the agent train the task-specific coefficients on the new task without training the basis vectors? This isn't really described anywhere.\n        - \"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. \n    - Sec 4.2\n        - Am I reading right that the agents are trained on _tens of millions_ of tasks? Though it seems that convergence occurs much earlier on. Could the authors zoom in to the left portion of 4E?\n        - The training setting for 4.2 is very unclear. Is there a teacher? If not, what is the training process? Why is \"accuracy\" used as the metric in a sequential decision making problem? How many total tasks are there (training on 10M out of 1B maybe isn't that bad, but 10M out of 30M isn't super encouraging)? Does accuracy correlate with actual performance on the task (i.e. does the agent actually compositionally generalize to new RL tasks)? How are we measuring \"OOD generalization\" in this setting?\n- Necessary or sufficient conditions \n    - The authors' aim, as stated in the introduction, is to understand \"what properties of the data generating process allow to discover modular solutions enabling compositional generalization\"\n    - However, the phrasing of the first lines of page 4 suggests that assumpions i-iii are _sufficient_ conditions, but not _necessary_ conditions for identification\n    - This means that what we really get is an understandong of _some_ of the properties that allow to discover modular solutions that compositionally generalize\n    - The authors do take steps to understand empirically what the \"cost\" is of violating those assumptions, and the assumptions themselves are fairly intuitive, so that's a positive\n- Interpretability of linear hypernetworks\n    - 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_\n    - Linear parameter combinations within a hypernet don't seem to have such a straightforward intuition\n    - 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.\n\n\n[1] Ghazi et al., \"Recursive Sketches for Modular Deep Learning.\" ICML, 2019.\n\n[2] Andreas et al., \"Neural Module Networks.\" CVPR, 2016", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)