Attention for Compositional Modularity
Abstract
Modularity and compositionality are promising inductive biases for addressing longstanding problems in machine learning such as better systematic generalization, as well as better transfer and lower forgetting in the context of continual learning.
Here we study how attention-based module selection can help achieve compositonal modularity – i.e. decomposition of tasks into meaningful sub-tasks which are tackled by independent architectural entities that we call modules.
These sub-tasks must be reusable and the system should be able to learn them without additional supervision.
We design a simple experimental setup in which the model is trained to solve mathematical equations with multiple math operations applied sequentially.
We study different attention-based module selection strategies, inspired by the principles introduced in the recent literature.
We evaluate the method’s ability to learn modules that can recover the underling sub-tasks (operation) used for data generation, as well as the ability to generalize compositionally.
We find that meaningful module selection (i.e. routing) is the key to compositional generalization.
Further, without access to the privileged information about which part of the input should be used for module selection, the routing component performs poorly for samples that are compositionally out of training distribution.
We find that the the main reason for this lies in the routing component, since many of the tested methods perform well OOD if we report the performance of the best performing path at test time.
Additionally, we study the role of the number of primitives, the number of training points and bottlenecks for modular specialization.
Introduction
Modularity and compositionality are appealing inductive biases for addressing several long-standing problems of machine learning such as generalization under distribution shift, a.k.a out of distribution (OOD) generalization [19, 7] and continual learning (CL) [6, 20, 17].
Given a set of modules the goal is to find a decomposition of knowledge into these modules s.t.
a new task can be solved efficiently through recombination of modules (compositionality), addition of new, or/and fast adaptation of existing modules.
Achieving compositionality presents a number of challenges. These include (a) decomposing tasks into reusable modules.
While related to sub-task discovery [14, 23], this has a promise of improving sample complexity through positive transfer and improved OOD generalization [19].
(b) Routing samples through a set of modules.
The latter is especially challenging when the routing procedure can be subject to forgetting as is the case continual learning.
Additional challenges not addressed in this work include e.g. pruning and addition of new modules to ensure enough capacity when new tasks must be learned.
Different forms of attention have been proven useful for achieving meaningful modularity in the mixture-of-experts (MoE) setup [21, 7, 10, 9, 22].
Here we are interested in compositional modularity — the training tasks should be decomposed into meaningful and reusable sub-tasks that can be executed in parallel or sequentially. In our instantiation each such subtask is assigned to an architectural entity
M 1 M 2 M…
X
X X
X
Functional Output …
Query Key …
OOD score 0.
+ -0. sin sin(0.12+(0.48))
Input e.g. signature (a) Cross Attention
M 1 M 2 M…
X
X X
X
Functional Output …
Query Key …
OOD score
Layer 1 Layer(b)
Structural Oracle
X 1 X
X 3 X…
Query
Key E n c ode r / C on c a t ena t i on
M 1 M 2 M… … … (c) Attention with encoder
Layer signatures
X 1 X
X 3 X…
E n c ode r / C on c a t ena t i on
M 1 M 2 M… … … …
L=1L=2 Query
Key (d) Attn.
+ encoder + Sbtl.
— a module, and the network is dynamically constructed for each input sample through a sequence of module selection & application steps.
In contrast to the MoE, where the supervision is provided after each module’s application, in compositional modularity the supervision is provided after the final selection step — i.e. after the desired depth of the modular network is achieved.
While MoE is only limited to parallel sub-task execution, the main advantage of compositional modularity is that it can decompose a task into sub-tasks, where the sub-tasks are executed both sequentially or in parallel.
For example, given a task of solving additions and subtractions (e.g. 4 + 5 − 3 = 6 ), a MoE would learn a single module (i.e. expert) to execute the whole expression, while a compositional system should be able to learn a module specialized on "+" and a module specialized on "-" executing both sub-tasks sequentially.
This can allow such system to achieve better compositional OOD [4, 22] generalization and efficient adaptation to new tasks in CL [25].
The main question we study here is how compositional modularity can be achieved.
To this end we first design an experimental setup based on math equations.
This naturally compositional domain provides the possibility to generate simple compositional problems for fast experimentation.
We study a wide range of key-value attention [3] based module selection mechanisms in the compositional modularity setting.
We find that only the structural oracle baseline — a baseline which is given the privileged information about which part of the input should be used for module selection, can reliably achieve modular specialization and reduce the compositional OOD loss.
Further, we find that the main reason for the poor OOD performance lies in the pure quality of the routing component, which we demonstrate through the excellent OOD performance of modular baselines in the test-time oracle setting — a setting where we report the OOD loss of the best module selection path after checking all possible paths.
Finally, we discuss the role of task diversity, number of samples per primitive operation, as well as the role of bottlenecks for preventing modular collapse.
Models
Here we describe the modular systems used in this work.
Similarly to neural production systems (NPS) [9], we aim to learn a set of modules, where each module learns a simple function that can be applied to the input variables.
In contrast to NPS and similarly to [22], we want the modules to form a chain (modules are applied in superposition to each other), where the supervision is only applied at the end of the chain.
Input data.
Similarly to [9], we assume that the input variables represent some meaningful entities of the real world.
In general terms, these entities can correspond to, e.g., objects in vision, words in NLP, or even to other high level variables.
The task of extracting such variables from high-dimensional input data is considered external to the system that we study here and can be accomplished by any method built for extracting object-centric representations [8, 15, 16].
We assume that input data is given as a set of N variables { X 1 , X 2 , . . .
X N } , hence a sample x = { x 1 , x 2 , . . .
, x N } consists of values of each of the N underlying variables, s.t.
each x n ∈ R
d is a d -dimensional vector.
Module selection.
We test several approaches for module selection.
These approaches share the following module selection principle, that differ in detail.
We consider a set of M modules, each module is equipped with a signature vector s m ∈ R r .
Additionally, we consider a modular system with L layers (we use l to index these layers).
At each layer the system must select a module to apply to the input of the current layer l conditioned on this input. The pool of possible modules is shared across layers.
Module selection is performed in the following 2 steps using key-value attention mechanism:
[Equation]
where W ( k ) l ∈ R p × r , W ( q ) l ∈ R p × d ′ are learnable parameters and p l ( m |· ) is the probability of the m th module at l th layer.
E is an encoder that is optionally applied to the input x .
As detailed next, encoder can be a learnable function such as neural net, in which case it is parameterized with η l .
Encoder can also represent a simple concatenation operation.
The encoder projects x into d ′ dimensional space.
We use θ ( st ) m :=
( W ( k ) , W ( q ) , η m ) to refer to structural parameters.
Next, we describe four variations of the above mechanism used in this paper. (1) Cross-attention: in this case, the projector E is an identity function, i.e x = E ( x ) .
Additionally, we encode each x n separately with a shared set of weights W ( q ) l , which gives a separate query q n per input variable.
As visualized in Fig.
1a)
, we use dot-product to calculate the attention score matrix R ∈ R N × M , s.t. R n,m = k
Tm q n .
(2) Concatenation: the projector E performs a concatenation operation producing a single vector x ′ =
E ( x ) = concat ( x ) ∈ R d ′ ( d ′ =
dn here).
Queries are then produced using Eq. 1.
The score matrix R ∈ R 1 × M , s.t. R 1 ,m = k
Tm q (Fig. 1c).
(3) Transf.
Encoder: in this case E is an encoder neural net parameterized with η l , whose goal is to summarize information present in all variables in x into a single vector x ′ =
E ( x ; η m ) ∈
R 1 xd ′ , where d ′ < dn (Fig. 1c).
In our experiments, E is instantiated with a transformer encoder network [24].
The score matrix R ∈ R 1 × M , s.t. R 1 ,m = k
Tm q .
(4) Structural Oracle: the model has privileged information about which of the N input variables is relevant for module selection at each layer.
E.g. for the math expression (0 . 4 + 0 . 01) ∗ 0 .
4 , instead of having to learn which input variable to use for module selection at each layer (one of the real numbers or one of the two operations), this baseline will use operation “ + ” for selecting module at layer l = 0 and “ ∗ ” at layer l = 1 (Fig. 1b).
We use the straight-through Gumbel softmax trick [13] to perform differential hard-module selection over the score matrix R .
In our experiments we always perform hard selection at test time.
At train time, we additionaly consider a baseline (“ +all path ”)
that, instead of a hard selection, checks every possible path in the modular system and weights the training loss of each path with it’s probability.
A path probability can be calculated as an exponential of the sum of the log-probabilities of each module in a given path, where module’s activation probability is calculated as in Eq.
Module application and functional bottleneck.
The selected module m ∗ receives all N input variables.
Similarly to NPS [9] we instantiate another bottleneck for the number of variables that can be processed by the functional component of each module – the functional bottleneck .
More specifically, we only allow two variables to be processed by the module’s functional parameters.
This design decision is motivated by the math equations domain as detailed in Sec. 3.
This bottleneck can be learned using another attention mechanism (query coming from each variable and key coming from the selected module’s signature).
As we demonstrate in Sec.
4.2, the functional bottleneck is essential to prevent module collapse.
Compositional modular systems vs. MoE.
In our experiments we consider compositional modular systems.
Its main difference to MoE is that the number of module applications L is greater than 1, whereas in MoE L = 1 .
Here, we refer to a single application of modules as a modular layers, where at each layer a module is selected from the same pool of modules.
Vanilla MoE system has the disadvantage of not being able to compositionally generalize to unseen combinations of attributes.
Parameter sharing.
Unless stated otherwise, modules at each layer share structural parameters W ( q ) l as well as η l (if applicable).
Motivated by the CL use case, where less parameter sharing can reduce forgetting, we ablate the role of sharing of W ( k ) l in Sec. 4.
Note, not sharing W ( k ) l between modules is essentially equivalent to instantiating s m ∈ R p and learning it directly.
Similarly to NPS, the inputs are pre-processed by a simple MLP encoder (before application of each of the modular layers), which is also shared across all modules and layers.
Training.
We derive our training objective by treating the optimal module selection path as a hidden variable denoted as z .
We consider a single path through the modular system as a sequence of L consecutive module selection and application steps.
The number of possible paths is M L .
We first derive a lower bound for the log-likelihood as (cf.
App.
A for details):
[Equation]
where H q denotes the entropy of the distribution over paths q ( z ) .
Our training objective is to find parameters θ ∗ =
{ θ ( st ) , θ ( f ) } ∗ that maximize the right-hand side in the Eq.
4, that is:
[Equation]
) where D tr is the training dataset.
By default, in our experiments we ignore the entropy term and approximate the expected log-likelihood with a single sample of z obtained by sampling modules at each layer from the Gumbel-softmax distribution parameterized with score matrix R .
We also consider baselines where loss is computed exactly by checking every possible path for a given pool of modules.
This is not a scalable solution in a general case, as the number of paths grows as O ( M L ) .
Data generation
Math equation tasks.
Here we consider the input data to be comprised of a set of N attributes X 1 ...N x paired with labels Y . Both are sampled from a joint distribution p
( Y, X 1 ...N ) =
p ( Y | X 1 ...N x ) p t ( X 1 ...N x ) .
We instantiate this setting in the math equations domain similar to [9, 19] with N = 5 .
We select this domain due to it’s simplicity, which enables fast experimentation, as well as the compositional nature of this domain.
Specifically, X 1 , X 2 ∈ R [ − 1 , 1] follow a uniform distribution U ([ − 1 , 1]) over the domain of real numbers between − 1 and 1 .
Variables X 3 , 4 , represent the primitive math operations to be performed on the first two variables, X 5 is a constant.
Hence, p X 3 , 4 =
U ( S ) , where S is a set of all possible operators (e.g."+", "-", "sin" etc., the list of operations is in App. B).
The dependant variable y is generated by applying sampled operations x 3 and x 4 to the inputs x 1 and x 2 sequentially, i.e. the result of one operation becomes the input to the next: in the reverse polish notation “ x 1 x 2 x 3 x 1 x 4 ” (e.g. ".1 .04 + .1 *" evaluates to . 014 ).
In this setup, given L = 2 , at each step the modular system has to learn to select modules given all the samples of all five input variables X 1 ...N .
The information about which module to select is encoded in X 3 for the first layer and in X 4 for the second. The privileged information about which variable to use for selection [-(i.e. l = 
 →-]( X 3 or X 4 ) is only provided to the structural oracle baseline.
The goal is to minimize the MSE loss by learning the mechanism p ( Y | X 1 ...N x ) .
Data generation.
The training data D tr is generated by considering all possible combinations of length L = 2 of | S | math operations, i.e. | S | L combinations. In the follwoing, we refer to one such combination as a task, and an underlying math operations as primitive, sub-task or mechanisms.
For each sample we first randomly select 1 task out of | S | L tasks that can be composed with L =operations encoded in input variables X 3 and X 4 .
To generate y , we first sample x 1 , x 2 ∼ U ([ − 1 , 1]) and x 3 , x 4 ∼ U ( S ) . We then apply the operations x 3 and x 4 as described above.
We sample the total of ( | S | L ∗ C ) /
( | S | ∗ L )) points, where C defines the expected number of training points to be seen by the algorithm per operation 1 .
Out of | S | L tasks, we select | S | tasks for the compositional OOD test set (i.e. one task per primitive operation), which are excluded from the training data.
In our experiments we use C ∈ [150 , 500 , 1000 , 5000 , 10000 , 20000] for the training dataset.
We use C = 200 for the in-distribution (ID) test set and a larger C = 2000 for the OOD test set.
Experiments
In this section we investigate how different module selection approaches can achieve compositionality.
We first present the baselines and metrics we use.
Baselines.
We consider the following baselines.
Cross-Attention (X-Attn.)
baseline separately encodes each input variable and performs cross-attention operation to obtain an N × M matrix of scores, it selects the module with the highest score (Fig. 1a).
Structural Oracle (Str.
Oracle): similar to the X-Attn.
baselines, but it is given the information about which variable to use for module selection at each step (i.e. x 3 at layer 1 and x 4 at layer 2).
Transformer Encoder (Trans.): processes the input variables with a transformer encoder, producing one encoding vector that is then matched with the modules’ signatures using key-value attention.
This encoding incorporates information about all input variables but can become selective throughout the learning process.
2 Concatenation + Attention (C-Attn) baseline simply concatenates all the input variables and performs key-value matching with the resulted concatenated vector.
Functional Oracle : only learns the structural parameters of the modules and uses a set of modules with a hard-coded ability to perform perfectly one of the mechanisms (math operations) used for data generation.
For all modular architectures we use modules with functional component being an MLP of width 300, depth 2, and a ReLU [2] activation.
Unless stated otherwise, we fix the number of modules to 30 in all cases independently from the number of primitives | S | .
We use two modular layers, i.e. L = 2 , which is equal to the number of primitive operations per sample in our experiments.
Monolithic baseline represents an MLP with 2 layers and the width of 9000 (30 times the width of each module in modular case).
This baselines receives as input a concatenation of all input variables (i.e. no functional bottleneck).
We ablate the importance of sharing W ( k ) L between modules at a layer in Figs.
8a and 8d and come to the conclusion that sharing these weights does not result in performance improvement, and in many cases it actually results in performance decrease.
For this reason, in our experiments modular baselines do not share W ( k ) l between modules in a layer unless stated otherwise.
By default, modular baselines described above only apply the most likely module at each layer.
Hence, at test time the most likely module corresponds to the module with the highest score in R ; at training time, the module index is samples from the Gumbel-softmax distribution parameterized with the score matrix R .
Additionally, we consider a version of the modular baselines that optimizes the objective in Eq. 4 exactly. Here, the final loss (MSE) is calculated by weighting the losses obtained with each possible module combination (i.e. path) with the probability of this particular path.
The probability of each path is calculated as an exponential of the sum of the log probabilities of each module selected in a given path.
These baselines are marked with +all.pth. .
We optionally add another hard attention layer to select which input variable to use as input to the module selection mechanism (as illustrated in Fig. 1d).
This additional selection bottleneck is signaled with +Sbtl. .
Here, instead of passing all variables to the selection mechanism at each layer, we only pass a single variable selected through matching it with layer signatures using cross-attention.
We include this baseline as it results in slight performance improvements in some cases.
The intuition is that in this way the model can easier learn the structure of the underlying math expressions – i.e. x 3 always encodes operations applied first and x 4 always encodes the operations applied thereafter.
Metrics.
We use the following metrics in this section: (1) ID loss – in-distribution test loss; (2)
Compositional OOD loss ( OOD loss ) — loss on the compositional OOD test set containing novel combinations of operations not seen during training; (3) Specialization score – consider an activation matrix A R × M , s.t.
A i,j is proportional to the probability that module j activates if the input contains operation i .
We define specialization score as inverse average normalized entropy of each the rows of
A , i.e. 1 − 1 R (cid:80) Ri = 1 log( M ) H ( A i, : ) .
A specialization score of 1 signifies that the module selection mechanism could assign each primitive operation in the data generative process to a single module.
This module, however, can be potentially the same module for all primitives.
To control for this, we consider the collapse score.
(4) Collapse score : consider ⃗a =
(cid:80) Ri =
A i, : — a vector representing unnormalized module activation distribution.
We define collapse as the normalized entropy of ⃗a , i.e.log( M ) H ( ⃗a ) .
Lower collapse score means that many operations activate the same modules.
An ideal
OO D l o ss
Str.
Oracle
C-Atn.
Transf.
+ Sbtl.
X-Atn.
+ Sbtl.
I D
l o ss 0.50 0.
S
pe c
i a li z a t i on (a) Functional Oracle 0.0 0.
OO D l o ss
Str.
Oracle + all.pth.
C-Atn.
+ all.pth.
Transf.
+ all.pth.
+ Sbtl.
X-Atn.
Monolithic 0.0 0.
I D
l o ss 0.4 0.
S
pe c
i a li z a t i on (b) End-2-end training system should have specialization and collapse scores equal to note that as shown in Fig. 8(c) none of baseline severely suffers from module collapse ( discussion in sec
4.1 Results
Functional Oracle setting.
We first assess different module selection strategies in a functional oracle setup.
We present consolidated view of selection strategies in Fig.
2a.
We observe that even in the functional oracle setting only the structural oracle baseline can achieve perfect specialization, which is also reflected in low comp.
OOD loss.
Second best is the “ Trans. + Sbtl. ”
baseline which can achieve reasonably low OOD loss on average as well as good specialization.
This result demonstrates that even if the knowledge about the sub-tasks is already perfectly encoded in the modules, learning how to combine these sub-task modules to solve the global task is a difficult problem on its own.
[Figure]
Only the structural oracle minimizes compositional OOD loss.
Results are shown for the end-2-end learning setting in Fig. 2b.
This is inline with [18], who use one-hot task descriptors identifying which module to use for each task (our oracle only get information about which variable to use for module selection, the latter is still learned, and not told directly which module to use).
The good performance of Str.
Oracle can be attributed ton the perfect specialization as shown in Fig.
2b (right).
The second best selection strategy is X-Atn. , which also achieves second best OOD loss.
All modular methods achieved OOD loss lower than that of the Monolithic baseline, which proofs the utility of modular solutions.
We consider the loss ( ↓ ) of the monolithic baselines as an upper bound for good performance.
Poor OOD performance in end-2-end setting is due to sub-optimal module selection.
To find why modular methods fail to generalize compositionally OOD in the end-2-end regime we measure the compositional OOD loss in the so-called test-time oracle regime.
In this regime we check every possible path in the modular network during the test time and report the minimal loss.
This testing strategy assumes the ground truth knowledge for the test set.
In Fig.
3 we see that for all modular solutions we can achieve very small loss in this regime.
This suggests that functional components in these models have learned meaningful information and that the reason for bad performance in the non-oracle settings lies in the structural components responsible for routing.
Given that good OOD loss can be achieved purely through recombination of existing modules, this approach can be useful for meta-learning and continual meta-learning settings like OSAKA [5].
Better performance of +all.pth .
As expected, we observe in Fig. 8(a)
that performing training in the + all.pth.
regime, where each possible path is checked during each training step, results in better ID for all methods (with exception of the Str.
Oracle).
Bad specialization of Transf.
baselines.
We observe in Fig. 8(c), that good OOD performance generally goes along with high specialization score, c.f.
X-Atn. and Str.
Oracle baselines.
Nevertheless, there is a clear exception: Transf.
+ all.pth.
baseline reaches OOD loss comparable to that of the
X-Atn.
+ all.pth.
baseline, while the later one reaches specialization score double as high.
In fact, we observe that both Transf.
+ all.pth.
and Transf.
baselines tend to have very bad specialization while
20000 C 0.
OO D l o ss (a)
Str.
Oracle
X-Atn.
Transf.
+ all.pth.
+ Sbtl.
C-Atn.
+ all.pth.
Transf.
+ all.pth.
Monolithic    20000 C 0.
S
pe c
i a li z a t i on (b)     14 | | 0.
S
pe c
i a li z a t i on (c)     14 | |
OO D l o ss (d)   15 | | 0.
OO D l o ss (e)    20000 C
OO D l o ss (f) | |  maintaining a high collapse score. After analysing the module activation patterns (c.f. examples Fig. 9) of these methods, we observe that these two methods activate modules much less sparsely than other methods making decision about which module to activate not only based on operations (vars x 3 and x 4 ) but also based on the digits ( x 1 and x 2 ).
This can be attributed to the fact that module selection is strongly over-parameterized for these baselines and misses any bottleneck.
Hence, simply adding a selection bottleneck ( +Sbtl. ) improves specialization considerably.
The role of the number of primitives | S | and training points per operation C .
We now turn our attention to the role of the number of primitives | S | and the number of training points per primitive C .
We start with two hypotheses: (1) Increasing the number of primitives results in better compositional generalization ability of modular methods due to the increase in the number of tasks (i.e. number of combinations in which each primitive is seen); (2) Increasing number of training points per task can have a harmful effect on the compositional generalization ability.
Intuitively, for large number of samples per operation, the bias towards solutions with specialized and reusable modules is weak.
This is because enough training samples are given for achieving low training loss even when a separate module per primitive per task (as opposed to only per primitive) is learned. On the other hand, when the number of samples per primitive is limited, the number of training points may be not enough for achieving low training loss with solution without sharing, and hence, sharing modules across tasks can result in lower training loss and more bias towards learning reusable and sharing sub-tasks.
Regarding our first hypothesis, indeed in Fig. 4(d) we observe a reduction in the OOD loss as we increase the number of primitives | S | from 3 to 5 consistently for all methods emphasizing the importance of diversity of tasks for specialization and compositional generalization.
Interestingly, for | S | > 5 the effect of | S | is insignificant, and as we show in Fig.
4(e), OOD loss of some methods actually rises as more primitives are introduced, with only Str.
Oracle reliably performing well.
This can be attributed to the fact that increasing | S | also leads to more OOD tasks (unique combinations of primitves) in the OOD test set, which makes the OOD problem harder.
Furthermore, in Fig. 4(c)
we observe that increasing | S | can result in decrease in specialization score, a trend especially pronounced for | S | between 10 and 15 (except Str.
Oracle ).
Hence, disentangling sub-tasks into specialized modules becomes harder when more sub-tasks are present in the underlying data.
We now turn our attention to the second hypothesis regarding the role of the number of training samples per operation C .
We make the following observations in Fig.
(a) First, even for the monolithic architecture, that has no inductive bias favouring specialization and OOD generalization, increasing C results in decreased OOD loss.
This is inline with observations made in the context of large scale training [12, 1], where good amount of generalization and transfer is achieved through simply increasing the amount of training data 3 .
(b) For Str.
Oracle and X-Atn.
we observe in Fig. 4(b) that modular specialization improves as we increase C (this is against our hypothesis), which is reflected in the tendency of decreasing OOD loss for larger C (Fig. 4(a)).
For the C-Atn. , Transf.
+ all.pth.
and Transf.
+ all.path.
+Sbtl.
baselines, that generally do not result in good specialization and OOD loss, we observe an opposite trend: i.e. larger C leads to worse specialization, and only for C-Atn.
we observe the tendency to larger OOD loss with increased C (Fig. 4(a)).
(c) In Fig. 4(f) only for small | S | we observe the overall trend of growing IID loss with increased C (in line with out hypothesis).
This again emphasizes the point that if the data is not diverse enough (is the case for | S | = 3 ), increasing C worsens the performance.
Overall, from observation (b) and (c) we can conclude that keeping C low does not result into enough bias towards specialized solution, hence the main source of such bias must originate from the architectural/algorithm choice.
On the other hand, if the algorithm under consideration does not provide enough inductive biases (e.g. in form of bottlenecks) towards specialized solution, increasing C indeed worsens the performance.
4.2 Ablation: the role of functional bottleneck
[Figure]
Central to our modular methods to work is the functional module bottleneck: each module can only take a subset of variables as input which prevents collapse as we show next.
In the experiments above, for the sake of simplicity we manually selected only x 1 and x 2 (i.e. the input digits) as inputs to the functional parts of the modules.
Here we show that, similarly to NPS [9], this bottleneck can also be learned.
In Fig.
5 we plot the collapse score ( ↑ ) for three application options (Str.
Oracle, Transf., X-Attn).
We test three types of functional bottlenecks: the “oracle” knowledge that only x 1 and x 2 should be selected as input to the functional component, a learned attention bottleneck that selects which two variable to take as input, and without the functional bottleneck at all.
We find that across all three application modes, both oracle and attention bottlenecks perform similarly in terms of collapse.
This shows that the functional bottleneck can be instantiated with a learnable key-value attention.
Removing this bottleneck results in increasing collapse in all application modes.
Conclusion
In this work we studied different attention-based module selection approaches for compositional modularity.
We have shown that information bottlenecks play an important role for module specialization (functional bottleneck) and module selection (selection bottleneck).
The main take-away is that in our experiments only the structural oracle baseline could reliably achieve good specialization, resulting in good compositional OOD performance.
We also showed that poor performance of non-oracle solutions can be attributed to attention-based routing failing to correctly route samples at test time.
These findings motivate application of these methods in settings where a small number of samples are available at test time to adapt the routing mechanism, or using self-supervised routing mechanisms that would facilitate test-time adaptation.
Another important research direction is identifying and implementing information bottlenecks necessary for specialization in modular networks. A [Figure] [Equation]
[Figure] [Equation] [Figure] [Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure] [Equation] [Figure] B Basic operations
We sue the following list of basic expressions for training and testing data generation.
For some input x and y: x+y , x-y , x*y , x*x , y*y , (x+y)*x , min(x,y) , max(x,y) , (x+y)-x , (x+(2*y)) , sin(x) , sin(y) , cos(x) , cos(y) , cos(x+y) . [Table] [Figure] [Table]
[Figure]
[Table]
X
[Table]
t n
[Table]
A
[Table]
[Table]
l .
X
[Table]
+
[Table]
h
[Figure]
[Figure]
r .
O r a c l e
S t r .
O r a c l e + a ll .
p t h .
Application option 0.0 0.
Str.
Oracle
C - A t n . C - A t n .
+ a ll .
p t h .
Application option 0.0 0.
Concat.+Atn.
T r an s f. + S b t l .
T r an s f.
T r an s f. + a ll .
p t h .
+ S b t l .
T r an s f. + a ll .
p t h .
Application option 0.0 0.
Trnsf.
Enc.
+ Atn.
M ono li t h i c
Application option 0.0 0.
Monolithic
X - A t n .
+ S b t l .
X - A t n . X - A t n .
+ a ll .
p t h .
+ S b t l .
X
[Table]
t n
[Table]
+ a ll .
p t h .
Application
[Table]
[Table] [Table]
S
pe c
i a li z a t i on
Cross-Atn.
S t r .
O r a c l e
S t
[Table]
O
[Table]
a c l e + a
[Table]
p
[Table]
C -
[Table] [Table]
h
[Table]
[Table] [Table]
[Figure]
C ∈
[Figure]
] , | S | ∈ [5 , 7 , 10 , 15] , 5 seeds.
[Figure]