Scalable Modular Network: A Framework for Adaptive Learning via Agreement Routing
Authors: Minyang Hu
Hong Chang
Bingpeng Ma
Shiguang Shan
Xilin CHEN
Venue: ICLR 2024 poster (2024)
Keywords: modular neural network; module selection process; adaptive learning
Primary area: representation learning for computer vision, audio, language, and other modalities
TL;DR: In this paper, we propose a novel modular network framework, called Scalable Modular Network (SMN), which enables adaptive learning capability and supports integration of new modules after pre-training for better adaptation.

Abstract: In this paper, we propose a novel modular network framework, called Scalable Modular Network (SMN), which enables adaptive learning capability and supports integration of new modules after pre-training for better adaptation.
This adaptive capability comes from a novel design of router within SMN, named agreement router, which selects and composes different specialist modules through an iterative message passing process.
The agreement router iteratively computes the agreements among a set of input and outputs of all modules to allocate inputs to specific module.
During the iterative routing, messages of modules are passed to each other, which improves the module selection process with consideration of both local interactions (between a single module and input) and global interactions involving multiple other modules.
To validate our contributions, we conduct experiments on two problems: a toy min-max game and few-shot image classification task. 
Our experimental results demonstrate that SMN can generalize to new distributions and exhibit sample-efficient adaptation to new tasks. 
Furthermore, SMN can achieve a better adaptation capability when new modules are introduced after pre-training. 
Our code is available at https://github.com/hu-my/ScalableModularNetwork.

SCALABLE MODULAR NETWORK: A FRAMEWORK FOR ADAPTIVE LEARNING VIA AGREEMENT ROUTING
Minyang Hu?? Hong Chang Bingpeng Ma², Shiguang Shan Xilin Chen
1 Institute of Computing Technology, Chinese Academy of Sciences University of Chinese Academy of Sciences
minyang. hu@vipl. ict.. ac. en, bpma@ucas ac. [changhong, sgshan, alchen]@icte ac. en
ABSTRACT
In this paper, we propose a novel modular network framework, called Scal- able Modular Network (SMN), which enables adaptive learning capability and supports integration of new modules after pre-training for better adaptation. This adaptive capability comes from novel design of router within SMN, named agreement router, which selects and composes different specialist mod- ules through an iterative message passing process. The agreement router itera- tively computes the agreements among set of input and outputs of all modules to allocate inputs to specifie module. During the iterative routing, messages of modules are passed to each other, which improves the module selection process with consideration of both local interactions (between a single module and in- put) and global interactions involving multiple other modules. To validate contributions, we conduct experiments on two problems: a toy min-max game and few-shot image classification task. Our experimental results demonstrate that SMN can generalize to new distributions and exhibit sample-efficient adapta- tion new tasks. Furthermore, SMN can achieve better adaptation capability when new modules are introduced after pre-training. Our code is available at hww.://owwoooooayyooooooooaoaeee
INTRODUCTION
Inferring new knowledge from existing information is hallmark of human intelligence. For in- stance, an individual who understands two concepts of "wooden table" and "iron chopsticks" can concept of "wooden chopsticks", even without having seen any amples of (Atzmon et al. 2020). This learning efficiency comes from the cognitive ability in human brain that decomposes knowledge into abstract sub-concepts and then recomposes these quired sub-concepts to understand a new concept (Sternberg) 2011 Meunier et al.) 2009). Through this modular cognition process, humans adapt well to the change of environment and excel at dealing with new tasks (Sporns& & Betzel 2016).
immediately understand
To realize such learning ability in machines, modular neural networks are introduced as a promising solution. A modular neural network involves several specialist modules, which serve as computa- tion units process specific sub-concepts. A natural question arises: Нow compose the specialist modules for different inputs? Prior works utilize metadata to control the information flow to mod- ules and make discrete routing decision before training. For example, some works (Andreas al., 2016 Hu et al.) 201 7a by Buch et al. (021) route different specialist modules for visual question answer with linguistic syntax tree. Besides, some vision- language methods (Radford et al. 2021 Tsimpoukelli et al. 2021 Alayrac et al. 2022; Najdenkoska et al.) 2023) design deterministic and fixed routing based on modality information, thus the visual and linguistic modules are responsible for different modality streams. Although these works decompose knowledge into modular parts and then recompose, they require extra metadata for routing, which does not always exist.
Learning the routing function from training data with some inductive biases is more promising for modular neural networks. One popular routing strategy is routing by input, where router selects
different modules to activate conditioned the input data (Shazeer al. 2017 [i al. 2022
Ma

---

et al. 2018
Rosenbaum et al.) 2018 Kirsch al.) 2018: Chang et al. 2019). To avoid trivial
solution where a router always selects all modules or few modules repeatedly for all inputs, these
methods often integrate reinforcement learning algorithm (Rosenbaum al.] 2018: Kirsch et al. 2018 Chang et al. 2019) or a noisy top-k selection mechanism (Shazeer et al. 2017] Li et al. 2022 Ma al., 2018). Despite the promising results shown by the matiing-by-inutt approach, these works ignore the role of module states in routing process and cannot effectively handle situations where new modules need to be inserted. To overcome this limitation, recent works (Goyal et al. [02
Rahaman et al. 2021 2022) propose routing based on the pairwise interaction between each module state and input. NI (Rahaman et al.) 2021) and NAC (Rahaman al.) (022) divide module param- eters into two parts: a signature vector and a computation part, where the former decides routing through a pairwise attention mechanism. RIMs (Goyal et al., 202Tb) extend the top-k mechanism based on all pairwise interactions between module states and model inputs. However, these studies ignore the global interactions during module selection process, which plays an important role for selecting appropriate modules to differentiate similar samples.
In this work, we introduce novel modular network framework, called Scalable Modular Network (SMN), to remedy the above limitations with adaptive learning capability. To this end, we design agreement router for dynamic module selection with consideration of both local interaction (be- tween input and single module) and global interactions (among all modules). More specifically, the agreement router selects and composes different specialist modules through an iterative message passing process, during which top-down feedback is iteratively obtained thus both local and global interactions are counted. In this way, the agreement router enables adaptive learning capability, i.c. adapting specific input and refining the module composition by iterative explorations.
Moreover, SMN supports integration of new modules after model training for better adaptation. When new modules are introduced, the agreement router dynamically reassigns coefficients of all modules based on the inputs and the functionality of existing modules (i.e. their corresponding outputs), enabling scalable routing process to enhance the adaptation ability of the modular net- work. Overall, leveraging the adaptive learning capability, SMN can effectively select and compose distinct modules to represent new samples within the representation space, thereby facilitating the out-of iitto---disributionn generalization and task adaptation.
Our primary contributions are as follows. (a) We introduce Scalable Modular Network, a novel modular framework that enables adaptive learning capability and supports integration of new mod- ules for better adaptation. (b) We propose agreement router for module composition that considers both local and global module interactions with iterative message passing. (c) We evaluate the pro- posed SMN architecture and agreement router in two learning problems that requires generalization and fast adaptation capabilities. We show that SMN is capable of sample-efficien adaptation and scalable with new modules. In addition, SMN can effectively distinguish similar sub-concepts by utilizing global information.
RELATED WORKS
Modular Neural Networks. A modular neural network comprises multiple specialist modules that exhibit properties of reusability and combinability. These properties enable modular architec- tures to possess good compositional generalization ability. To incorporate networks with modular inductive bias, previous research efforts (Andreas al. 2016; Hu et al.) 2017a| Rahaman et al.) 2021 2022; Goyal et al., 2021b) have explored various approaches. A direct solution is to employ the metadata (Andreas al. 2016 (Hu et al. 7a Buch et al. 202T), such natural language or modality information, to decide the routing before learning modules separately. Alternatively, other approach (Goyal et al. 202 aalaab b Rahaman al. 2021 2022) explores neuro-symbolic architectures and utilize distinct strategies to simultaneously learn modules and routing functions without rely- ing on extra metadata. In our work. the proposed SMN aligns with the latter approach, but with enhanced adaptation and scalable ability.
Module Composition. An important problem in modular network is how to select and compose specialist modules for different input samples, which is related to the compositional generalization of modular neural networks (Bahdanau et al.] 2019 Rosenbaum et al. 2019). Various routers have
been designed for this purpose. For instance, MoF. (Shazeer et al. 2017 et al. 2022 Ma et al.
2

---

Module Processor
Medsle
Module
Medale1
Tokenizer
0000
Clasalter
Medsle
Module
Module
logits
0Oo0!
Medule
Module
Medste
Image
t=
t=1
t=T
Figure I: A sketch the proposed SMN. The SMN contains three components: (a) a tokenizer that transforms input image into local features; (b) module processor containing specialist modules and an agreement router; (c) classifier that aggregates module outputs for final classification.
2018) employs a gating network select specialist modules based module input. Top-k selection mechanism is incorporated to restrict the number of modules to be activated and updated for each sample. RIMs (Goyal al. 2021b) extend the top-k mechanism to select RNN-based modules according to the pairwise interactions between each module state and input. Besides, NI (Rahaman al. 2021 divides the module parameters into a signature vector and computation part. The signature vector specifies the most relevant input vectors for the corresponding module by a pairwise kernel function. NAC (Rahaman al. 2022) inherits the same router design but imposes graph structural prior on the learning process. In contrast to RIMs and NAC, SMN selects modules based on agreement between inputs and outputs of all modules. This routing strategy incorporates global interactions during module selection process, and is compatible with arbitrary module designs.
Fast Adaptation on New Tasks. How learn new tasks with very few labeled samples by lever- aging experience from related training tasks is central question in few-shot learning (FSL. area (Hu al. 2023). Previous FSL approaches based on meta-learning framework can be categorized into two categories. The metric-based approach (Vinyals al., 2016 Snell et al., 2017 [Sngg al.) 2018 Hou et al. 2019 Allen et al.) 2019) aims to learn a cross-task embedding function and pre- dict the query labels based on the learned distances. The optimization based approach focuses
learning some optimization states, like model initialization (Finn et al. 2017 [unn et al.) 2019) or step sizes (Antomiou et al. 2019 Li et al., 2017). rapidly update models with very few labeled samples. Recent FSL works (Dong et al. 2022 Hiller et al. 2022 Lin et al., 2023) have shown that standard transfer learning procedure, which involves initial pre-training with self-supervised loss and subsequent fine-tuning, achieves state-of-the-art performance on various benchmarks. In this work, demonstrate that the transier-learnigg based approach can be combined with SMN framework and results in a further performance improvement due to the modular inductive bias.
SCALABLE MODULAR NETWORK
In this work, present a novel modular framework, Scalable Modular Network (SMN), that enables adaptive learning capability and supports integration new modules after pre-training. In Section 3.1, will introduce three components that constitute the SMN framework. Next, in Sections 3.2 and 3.3, we will describe how to route different specialist modules based on input- output agreements and explain how the agreement routing algorithm handles scenarios where new modules are added to the pre-trained network. Section 3.4 will focus on the training process for both modules and agreement router within the SMN framework. Finally, in Section 3.5, will discuss the relationship between agreement routing and other works.
3.1 OVERVIEW OF SMN
The SMN consists of three components: (a) a tokenizer that transforms an input image to set of local features, (b) module processor which contains multiple modules and an agreement router that composes a subset of specialist modules for specific input, (c) a classifier that aggregates the outputs from the module processor make final predictions. A sketch of SMN is presented in Fig.
3

---

Tokenizer. The tokenizer is any component that converts the input image to set of representa- tion vectors. It can be the feature extractor of a ConvNet without global pooling, or a transformer encoder. We try both two choices in the experiments.
Module Processor. The module processor is composed of an agreement router and multiple spe- cialist modules. The agreement router dynamically selects a subset of specialist modules to be activated through an iterative process. After that, each module computes using its own parameters with the acquired input vectors and then outputs single vector. These output vectors of modules will be forwarded to the classifier for further processing.
The classifier is responsible for receiving the output vectors from the module processor and subsequently mapping them to the label space. feasible to aggregate these output vectors from the module processor through average or concatenation operation to acquire a single vector for final classification.
Classifier
3.2 ITERATIVE ROUTING BY AGREEMENT
The kernel part SMN is the agreement router, which selects different specialist modules through iterative agreement routing algorithm. Assume that the tokenizer outputs N vectors (5.....,5N] and the module processor contains M modules [J.....JM.. The agreement routing algorithm maps from a set of N input vectors to the out- put vectors {U1, ..., (UI,..,UMM of M modules, vy (( The routing process iterates with two steps: computing input-output agreements and up- dating output vectors. First, the agreement router transforms each input vector through a learnable matrix and computes the agreement between trans- formed vector and each module output. Then, the router selects different input vectors for M mod- ules and refines their outputs.
Molie
t=
t=1
Figure 2: Agreement routing with two steps.Left: selecting imputs for distinct mod- ules and refining their outputs. Right: comput- ing the agreements between transformed in- put vectors (dark squares) and module outputs (three colored circles).
An illustration of agreement routing is shown in Fig. Each module output can be regarded as a cluster and the input vectors can be viewed as samples to be grouped. More formally, for a set of input vectors (s...... S. we first apply all available modules to each input vector s; get the processed vectors {ui = (s:)
We then calculate the average of the corresponding processed
vectors o) (0) the agreement score of between each input vector and module output oft-1) by cosine similarity:
N Σ 1 Uiy as the initial output of each module fy. At the t-th iteration, we calculate
0, t=
Was:
+ a
1111)
+ O
10
sali
where W. a learnable matrix that maps 8; into the module output space, and * denotes the dot- product operation. We normalize the agreement score aij with softmax function to get the agree- ment coefficient by which is then used to update the output vector of the j-th module:
(t)
exp(o (0)
N
(e)
Uiy-
M E 5
Through the softmax function, each module competes with others in every iteration to decide to what extent the input vector can be used to acquire its output vector:Thus, the agreement coefficients determine how the modules are activated, conditioned on the input-output agreements. The agree- ment routing process will alternate between Eq. I and Eq. 2 until the number of iteration reaches a threshold or the agreement coefficients no longer change. The process of agreement routing is summarized in Algorithm
4

---

Algorithm 1 Agreement Routing.
Input: {uij...., UNj (
AM f=11 (,,,.., SN]
Parameter: Wa
I: Initialize v/(0) N/2 Ei= Wij E RD,, aij (0)
0.
2: while / - + T do
of .W.*.
3: for all input vector 8, and output vector oj: agj t- aij
(1)
(f-1)
+
110 W
(()
exp (a
4: for all input vector 8, and output vector vy: Cjj
Σ, exp (a))
5: for all output vector uy: o3
Es off Uij
6: end while
7:
return fu] K
Sparse Module Selection. As the agreement routing algorithm does not explicitly impose con- straints the number of activated modules, natural question arises that how does it enable sparse module selection so as to avoid trivial solutions? We explore this question theoretically and discover that, as the iterative computation proceeds, the agreement routing naturally tends to produce a sparse module selection.
Theorem 1. Let {s1, SN) be N imput vectors and {ug = {( M modules [f...... fur) As for the agreement routing alternating between Eq. Mand Eq. if any two inpul vectors are not identical, the coefficient cij + 0 or 1 as the number of ierattion + oc.
output vectors of
Proof. Please refer to Appendix A.1.
Theorem shows that as the iteration number ~> oo, each input vector will ultimately be assigned to single specialist module based on the sign of agreements. This implies that there will only exist N connections between input vectors and specialist modules, avoiding the risk of trivial solutions.
3.3 ROUTING WITH NEW MODULES
After training, the agreement router within SMN should dynamically select and compose distinct specialist modules for different input samples. However, is possible that the downstream tasks introduce unseen visual primitives have new feature distributions. Under such circumstances, the trained modules may be inadequate in adapting to the environment changes. Therefore, becomes necessary to incorporate also adaptable to the environment changes.
modules that are not only complementary with existing modules but
The agreement routing algorithm can naturally handle the situation where modules are added into a trained model. Assume that a SMN with M modules [Ji...,,M and an agreement router has been trained, and new modules {fi...., fD} are being added into this model. We utilize the trained agreement router to route the existing and new introduced modules simultancously. For each module J/, we calculate the agreement score aiy and utilize softmax function across all
(t) M + D modules determine the agreement coefficient of and refine its output vector uz) as:
exp(a Σー= exp + EN exp(a
N
Σ (()).
(3)
Through the softmax function, the D new modules introduced into the trained model, where they compete with the existing M modules. This competitive process serves to allocate the new modules' attention distinct inputs, achieving complementary effects with the existing modules.
3.4 LEARNING MODULES AND AGREEMENT ROUTER
We jointly train specialist modules and an agreement router using a combination of two losses: clas- sification loss L.s and importance loss Limp- The classification loss is computed the classifier
5

---

component with aggregated outputs from the modules {0) (T) T
f=11- This loss guides the agreement
router to select appropriate modules to classify the input sample:
Ceta Σ vE(
(4)
J.
where o is output logits from the classifier for input sample x, and CE denotes the cross entropy loss. However, to select the same set of few modules for different samples, as illustrated in (Shazeer et al. 2017). To address this issue, an importance loss introduced to balance the utilization of specialist modules based on their agreement coefficients Cij. Following (Shazeer et al. 2017), we define the importance loss on a set of importance values {imp(f) )= as the square of the coefficient of variation (CV):
have observed that only using the classification loss consistently leads the router
Limp
(5)
where imp(fi) Σ Ggg with ojj being the agreement coefficient for input s; and module fz iteration T. The final objective L combines both two losses:
== Lets + Ahimp-
(6)
Here, À is
that balances the contributions of the two losses.
Fine-tuning with New Modules. When a pre-trained SMN needs to adapt to new task by adding new modules, fix all existing pre- pre-ttrined parameters and only fine-tune the newly added modules. This partial fine-tuning strategy reduces the required number of samples and avoids catastrophic forgetting enable the boarder generalization. Furthermore, most of the parameters of the new modules can be shared with pre-trained modules, and only a small set of parameters have to be tuned for their unique computations. Consequently, the total number of parameters does not noticeably increase with the number of new modules, resulting in a sample-efficient adaptation process.
3.5 DISCUSSION
Agreement Routing as Parse Tree Constructing. Modular neural network aims to factorize knowledge into multiple modular sub-concepts, which can be reused and recomposed to solve out- of-distribution tasks. This idea of knowledge factorization is similar with a cognitive theory, known as "recognize by components" (Biederman 1987), which suggests to construct a visual parse tree to represent an image. One challenge is how to dynamically assign the low-level nodes to high-level nodes in the visual parse tree. To address this issue, clustering like process can be applied (Sabour et al. 2017). which regards each high-level node as a cluster and low-level nodes as samples to be grouped. By iteratively computing the similarity between samples and clusters, each low-level node will be dynamically assigned to appropriate high-level nodes. From this point of view, the proposed agreement routing within SMN is equivalent to constructing a visual parse tree to represent an image. Each node in the tree corresponds to a specialist module responsible for specific sub-concept.
Agreement Routing and k-mmaans Clustering. Agreement routing similar with k-means al- gorithm, classical unsupervised learning approach which divides the training set into k distinct clusters examples that are close to each other. In contrast to k-means, the agreement routing algo- rithm is driven by both supervised classification loss and importance loss. This routing approach aims to adaptively select and compose different specialist modules for classification. Furthermore, in agreement routing process, the computed agreements previous step are accumulated for the next step, which encourages a sparse module selection to mitigate the risk of trivial solutions.
4 EXPERIMENTS
In this section, we empirically evaluate Scalable Modular Network in various problem settings, aim- ing to address the following key questions: (1) Can the agreement router improve module selection process by incorporating global interactions? (2) Can SMN exhibit better generalization and adap- tation capabilities? (3) Can SMN support integration of new modules after pre veettainnng for better adaptation? We first address these questions through a toy example and subsequently validate the applicability of SMN in more realistic scenario.

---

4.1 ToY PROBLEM: MIN-MAX GAME
We first consider simple Min-Max Digit Game, where model is required to find the minimum maximum digit according to the summation all digits in the image. Our hypothesis is if a modular model can incorporate global interactions into the module selection process, it can differentiate distinct samples even if some parts of them are very similar.
Method
Params
Test Ace
CNN
128.9K
42.34 (2.7)
Transformer
125.7K
45.76 (1.8)
Top-K Truncated
124.6K
42.94 (2.1)
131.9K
43.77 (3.4)
Figure 3: Visualization of synthetic images. Left and Middle: two synthetic images from training digital compositions (1,2) and (6.2). Right: a synthetic image from test digital composition (6, 1).
T
44.55 (2.2)
SMN
T =
131.9K
47.60 (4.1)
T = 4
54.46 (5.0)
Table I: Results on out-of-distribution samples.
Task Construction. Consider a set of N digital images, denoted as {(rrg d)) where each a, represents image sample and d, [0.9] signifies its corresponding digit label. We create syn- thetic image along with a corresponding label using any two elements with the following operation:
( dz); d, + d max dy), d, + dj < 10
{
min(d;
concatenate(1i, 2z);
In the above operation, we combine two digital images and assign the resulting synthetic image la- bel equal to the minimum or maximum value of the two digits according to their summation. We se- lect 55 different compositions of digits to construct a Min-Max Digit task, and split 40 compositions for training and another 15 compositions for test, respectively. To evaluate the out-of-distribttion generalization ability, constrain the synthetic image for test only include new compositions of two digits, but each individual digit must have been encountered during the training phase. See some synthetic examples in Fig,
Comparative Methods. We compare our SMN with two non-modular network architectures: CNN and Transformers. Furthermore, we replace the agreement router in SMN with two module selection methods (Shazeer et al. 2017 Rahaman et al. 202I): noisy top-k selection mechanism (Top-K) and truncated kernel (Truncated), ). We set the number of modules M to for SMN, with the expectation that these two specialist modules will learn find the minimum maximum digit in digital composition. For a fair comparison, we set the same number of modules for the two modular network variants. See more experimental details in the Appendix C
Out-of-distributien Generalization. A modular model should excel at composing specialist modules for out-of-distribution generalization. Thus, the agreement router improves module se- lection process, SMN is expected to exhibit better generalization performance compositions. Tab. shows the results of different models. First of all, when the the iteration num- ber T' is 0, the performance of SMN is similar with non-modular models, CNN and Transformers. However, SMN can significantly outperform them with a large T. This phenomenon shows the im- portance of iterative computation in agreement routing. Secondly, comparison with two modular variants (Top-K K and Truncated), our SMN generalizes better on new compositions than them. We attribute this to the global interactions incorporated in agreement router, which is absent for the two methods. Futhermore, we visualize the activation map of two modules within SMN. In Fig. 4 we observe that the module fi focuses on digits and 4 in the first example while becomes non- activated in the second example, even though the digit 6 appears in both examples. This verifies our hypothesis that incorporating global information helps the specialist modules to separate unseen similar samples, thus improving the generalization capability of the model.
the new digital
7

---

f1 f2
11
(a)
(b)
Figure 4: Activation maps of two modules within the SMN (T 4). (a) example from composition
(6,4) and label y (6, 4). (b) example from composition (6,0) and label y/ max(6, 0).
Test Accuracy
Method
h: = 15
k 30
k = 45
k = 60
k 75
k: 90
CNN
82.9 (2.4)
84.8 (2.0)
85.9 (2.6)
87.8 (1.6)
88.6 (1.3)
89.5 (1.3) 88.6 (0.9)
Transformer
82.2 (2.4)
85.0 (3.1)
86.4 (1.6)
87.4 (1.6)
88.8 (1.2)
Top-K Truncated
83.2 (3.9)
85.3 (2.3)
86.5 (1.6)
87.4 (1.1)
88. (1.0)
88.2 (0.6)
82.7 (3.3)
85.9 (1.9)
86.9 (1.3)
87.7 (0.9)
87.8 (1.2)
88.0 (0.7)
T = 0
84.4 (2.7)
86.9 (2.2)
88.2 (1.7)
89.7 (1.0)
89.6
90.3 (1.3)
SMN
T 2
86.4 (3.1)
89.0 (2.2)
90.4 (1.5)
91.7 (0.8)
91.8 (0.9)
91.8 (1.0)
T 87.8 (1.6)
87.3 (1.6)
89.0 (1.6)
90.6 (1.7)
91.3 (1.0)
92.3 (1.0)
Table 2: Results of various models for few-shot task adaptation. We run each experiment 10 times with different random seeds and report the average results with standard deviation.
Few-shot Adaptation. We further explore the task adaptation
capability SMN on a transfer learning scenario. We design a new task based on parity code of digital compositions. In this task, the label of a digital composition depends on the count of value in the binary encoding of minimum maximum digit.
Text acuuaay
Specifically, if the count is odd, the corresponding label is 1, otherwise it equals Max Digit task with sufficient training samples, and then fine- tune them on the Parity Code task with ie labeled samples. Tab, 2 shows that SMN is better than all other models with different values of And, this performance improvement becomes more significant with the increase of iteration number T. This compar- ison shows the importance of agreement routing on task adapta- tion capability for SMN. Besides, verify that SMN supports integration of new modules after pre-training for better adaptation in Fig. 5) We observe that the performance can be further improved by incorporating a small number of modules into pre- trained SMN. However, as we increase the number of added modules, the improvement diminishes. We attribute this the overfitting, as the incorporation of new modules introduces more parameters.
O. We train various models on the Min-
Number of new added modules
Figure 5: Test accuracy for SMN (T' 4) in terms of the number of new added modules.
4.2 FEW-SHOT IMAGE CLASSIFICATION TASK
We have shown that SMN exhibits better generalization and adaptation capabilities on new compo- sitions. While these improvement have been striking in synthetic dataset, raises question: Can SMN still perform well on more realistic scenario? To answer the question, we consider the Few- Shot Image Classification (FSIC) task, which aims to learn novel categories with very few labeled samples by exploiting experience from related training categories. Our intuition lies in the attribute compositionality hypothesis: each category can be represented as composition of attributes, which reusable in huge assortment of meaningful compositions (Hu et al. 2023). We conjecture that when applied to FSIC task, SMN should have a better task adaptation capability as its modular inductive bias leads to improved generalization previous unseen compositions.
8

---

Method
minilmageNet
CUB
RelationNet
49.31 +: 0.85
62.34 +: 0.94
MAML
46.47 0.82
54.73 = 0,97
ProtoNet
44.42 0.84
50.46 0.88
Basclinc++
48.24 0.75
60.53 0.83
SMN
54.35 0.77
67.33 0.90
ofnen adsed modsls
Table 3: Results on 5-way 1-shot classifica- tion with Conv4-64 backbone.
Figure 6: 1-shot improvement with different
number of added modules on minilmageNet.
minilmageNet
CIFAR-FS
Method
Backbone
5-way 1-shot
5-way 5-shot
5-way 1-shot
5-way 5-shot
SUN
VIT
67.80 + 0.45
83.25 0.30
78.37 + 0.46
88.84 + 0.32
FewTURE
Swin-Tiny
72.40 # 0.78
86.38 # 0.49
77.76 + 0.81
88.90 + 0.59
SMKD
ViT-S
74.28 0.18
88.82 0.09
80.08 0.18
90.63 0.13
SMN
VIT-S
75.76 0.84
89.28 0.47
82.87 0.79
91.35 0.56
Table 4: Results on 5-way 1-shot and 5-shot classification on miniImageNet and CUB.
Datasets and Setup. Following (Chen et al. 2019; Lin et al. 2023 we consider three popular benchmark in FSIC task: minilmageNett (Vinyals al., 2016), CIFAR-FS (Bertinetto et al., 2018) and CUB-200-2011 (CUB) (Wah et al. 2011). As previous works, we use the N-way K-shot setting construct training and test episodes from dataset, where labeled samples are provided for each of N classes. We report the average accuracy 600 test episodes with 95% confidence intervals. See more details of datasets and experimental setup in the Appendix C.5
Comparative Methods. We compare our model with four classical few-shot methods: (1) MAMI.
(Finn et al. 2017). (2) ProtoNet (Snell et al. 2017). (3) RelationNet (Sung et al.) 2018) and (4)
Baselinc++ (Cen et al. 2019). These methods use the Conv4-64 network backbones. Further more, we include three SOTA few-shot methods: (1) SUN (Dong et al. 2022). (2) FewTURE (Hiller
al. 2022 (3) SMKD (inn al. 2023), which use the vision transformer backbones.
Results. Tab. and Tab, show the results of various models on three datasets. Compared to other FSL method, SMN achieves the best performance all datasets, which shows modular archi- tecture design do improve the fast adaptation ability. With the same transformer backbone, our SMN can achieve new state-of-the-art performance, demonstrating the generality of our modular frame- work. Furthermore, in Fig. 6 we observe that the performance of SMN can be further improved by adding new modules. When we use the outputs from module processor for classification, integrating a small number of new modules improve the adaptation capability. We argue this is because new modules focus on distinct parts of images, which are complementary with existing modules.
CONCLUSION
We propose novel Scalable Modular Network (SMN) framework. which enables adaptive learning capability and supports integration of new modules after pre-training for better adaptation. These capabilities come from the novel design of agreement router in SMN, which improves the module selection with consideration of both local and global interactions. Experiments demonstrate that SMN exhibits a better generalization capability and supports integration of new modules after pre- training. For future research, it is promising to train large-scale modular network which supports the integration of new modules for better adaptation. This property enables to adapt general knowl- edge specific downstream tasks without forgetting the learned knowledge from pre-trained data.
ACKNOWLEDGMENTSS
This work is partially supported by National Key R&D Program of China no. 2021ZD01 1901, and National Natural Science Foundation of China (NSFC): 62376259 and 62276246.

---

REFERENCES
Man-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, lain Barr, Yana Hasson, Karel Lenc, Arthur Mensch, Katherine Millican, Malcolm Reynolds, et al. Flamingo: a visual language model for few-shot learning. 2022.
Kelsey Allen, Evan Shelhamer, Hanul Shin, and Joshua Tenenbaum. Infinite mixture prototypes for few-shot learning. In ICML, 2019.
Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Dan Klein. Neural module networks, In CVPR, 2016.
Antreas Antoniou, Harri Edwards, and Amos Storkey. How to train your MAMI.. In ICLR, 2019.
Yuval Atzmon, Felix Kreuk, Uri Shalit, and Gal Chechik. A causal view of compositional zero-shot recognition. In NeurIPS, 2020.
Dzmitry Bahdanau, Shikhar Murty, Michael Noukhovitch, Thien Huu Nguyen, Harm de Vries, and Aaron Courville. Systematic generalization: what is required and can be learned? In ICLR, 2019.
Luca Bertinetto, Joao F Henriques, Philip HS Torr, and Andrea Vedaldi. Meta-leaming with differ- entiable closed-form solvers. arXiv, 2018.
Irving Biederman. Recogrition-bb-compooenss a theory of human image understanding. Psycho- logical review, 94:115, 1987.
Shyamal Buch, Li Fei-Fei, and Noah D Goodman. Neural event semantics for grounded language understanding. Transactions of the Association for Computational Linguistics, 9:875-890, 2021.
Michael B Chang, Abhishek Gupta, Sergey Levine, and Thomas L. Griffiths. Automatically com- posing representation transformations means for generalization. In ICLR, 2019.
Weiyu Chen, Yencheng Liu, 7solt Kira, Yuchiang Frank Wang, and Jiabin Huang. A closer look few-shot classification. In ICLR, 2019.
Maxime Chevalier- Boissvert, Lucas Willems, and Suman Pal. Minimalistic gridworld environment for openai gym, 2018 https::/githii www. mox wwwpa//wwmmm///mmmnmmma.aaaaaa/a
Bowen Dong, Pan Zhou, Shuicheng Yan, and Wangmeng Zuo. Self-promoted supervision for few- shot transformer. In ECCV, 2022.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In ICMI., 2017.
Anirudh Goyal, Aniket Didolkar, Nan Rosemary Ke, Charles Blundell, Philippe Beaudoin, Nicolas Heess, Michael C Mozer, and Yoshua Bengio. Neural production systems. In NeurIPS, 2021a.
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. In ICLR, 2021b.
Markus Hiller, Rongkai Ma, Mehrtash Harandi, and Tom Drummond. Rethinking generalization in few-shot classification. In NeurlPS, 2022.
Ruibing Hou, Hong Chang, Bingpeng Ma, Shiguang Shan, and Xilin Chen. Cross attention network for few-shot classification. In NeurIPS, 2019.
Minyang Hu, Hong Chang, Zong Guo, Bingpeng Ma, Shiguang Shan, and Xilin Chen. Under- standing few-shot learning: Measuring task relatedness and adaptation difficulty via attributes. In NewrIPS, 2023.
Ronghang Hu, Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Kate Saenko. Learning to reason: End-to-end module networks for visual question answering. ICCV, 201
Ronghang Hu, Marcus Rohrbach, Jacob Andreas, Trevor Darrell, and Kate Saenko. Modeling rela- tionships in referential expressions with compositional modular networks. In CVPR, 2017b.
10

---

Louis Kirsch, Julius Kunze, and David Barber. Modular networks: Learning to decompose neural computation. In NewrIPS, 2018.
Alex Krizhevsky et al. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
Bo Li, Yifei Shen, Jingkang Yang, Yezhen Wang, Jiawei Ren, Tong Che, Jun Zhang, and Ziwei Liu. Sparse mixture-of-experts are domain generalizable learners. In ICLR, 2022.
Zhenguo Li, Fengwei Zhou, Fei Chen, and Hang Li. Meta-sgd: Learning to learn quickly for few- shot learning. arXiv, 2017.
Han Lin, Guangxing Han, Jiawei Ma, Shiyuan Huang, Xudong Lin, and Shih-Fu Chang. Supervised masked knowledge distillation for few-shot transformers. In CVPR, 2023.
Jiaqi Ma, Zhe Zhao, Xinyang Yi, Jilin Chen, Lichan Hong, and Ed H Chi. Modeling task relation- ships in multi-task learning with multi-gate mixture-of-experts In SIGKDD, 2018.
David Meunier, Renaud Lambiotte, Alex Fomito, Karen Ersche, and Edward Bullmore. Hier- archical modularity in human brain functional networks. Frontiers in neuroinformattic, 3:571,
2009.
Ivona Najdenkoska, Xiantong Zhen, and Marcel Worring, Meta learning to bridge vision and lan- guage models for multimodal few-shot learning. In ICLR, 2023.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In ICML, 2021.
Nasim Rahaman, Muhammad Waleed Gondal, Shruti Joshi, Peter Gehler, Yoshua Bengio, Francesco Locatello, and Bernhard Schölkopf. Dynamic inference with neural interpreters. In NeurIPS, 2021.
Nasim Rahaman, Martin Weiss, Francesco Locatello, Chris Pal, Yoshua Bengio, Bernhard Schölkopf, Erran Li, and Nicolas Ballas. Neural attentive circuits. In NeurIPS, 2022.
Clemens Rosenbaum, Tim Klinger, and Matthew Riemer. Routing networks: Adaptive selection of non-linear functions for multi-task learning. In ICLR, 2018.
Clemens Rosenbaum, Ignacio Cases, Matthew Riemer, and Tim Klinger. Routing networks and the challenges of modular and compositional computation. arXiv, 2019.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma. Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. International journal of computer vision, 11::211-252 2015.
Andrei Rusu, Neil C Rabinowitz, Guillaume Desjardins, Hubert Soyer, James Kirkpatrick, Koray Kavukcuoglu, Razvan Pascanu, and Raia Hadsell. Progressive neural networks. arXiv, 2016.
Sara Sabour, Nicholas Frosst, and Geoffrey E Hinton. Dynamic routing between capsules. In NeurlPS, 2017.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv, 2017.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In ICLR, 2017.
Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning. In
NeurlPS, 2017.
Olaf Sporns and Richard Betzel. Modular brain networks. Annual review of psychology, 67: 613-640, 2016.
11

---

Saul Sterbberg. Modular processes in mind and brain. Cognitive neuropsychology, 28:156-208, 2011.
Qianru Sun, Yaoyao Liu, Tat-Seng Chua, and Bernt Schiele. Meta-transfer learning for few-shot learning. In CVPR, 2019.
Flood Sung, Yongxin Yang, Li Zhang, Tao Xiang. Philip HS Torr, and Timothy M Hospedales. Learning to compare: Relation network for few-shot learning. In CVPR, 2018.
Maria Tsimpoukelli, Jacob L Menick, Serkan Cabi, SM Eslami, Oriol Vinyals, and Felix Hill. Mul- timodal few-shot learning with frozen language models. In NeurIPS, 2021.
Oriol Vinyals, Charles Blundell, Timothy Lillicrap, Daan Wierstra, et al. Matching networks for one shot learning. In NeurIPS, 2016.
Catherine Wah, Steve Branson, Peter Welinder, Pietro Perona, and Serge Belongie. The caltech-uesd birds-200-2011 dataset. 2011.
Haohan Wang, songwei Ge, Zachary Lipton, and Eric P Xing. Learning robust global representa- tions by penalizing local predictive power. NeurIPS, 2019.
12

---

CONTENTS
Introduction
2
Related Works
3
Scalable Modular Network
3.1
Overview of SMN
3.2
Iterative Routing by Agreement
3.3
Routing with New Modules
3.4
Learning modules and agreement router
Discussion
4
Experiments
4.1
Toy Problem: Min-Max Game
4.2
Few-Shot Image Classification Task
5
Conclusion
Theoretical Proofs
13
Modulated Modules
14
Experimental Details
15
Implementation Details for Min-Max Game
15
More results for Min-Max Game
15
Effect of Fine-tuning Agreement Router on Min-Max Game
15
Computational Cost Analysis
16
Details for Few-Shot Image Classification
16
Visualization Iterative Dynamics
17
Evaluation on Diverse Tasks and Datasets
17
Reinforcement Learning.
17
E.2
Continual Learning
18
Cross-Domain Few-Shot Learning
19
A THEORETICAL PROOFS
Theorem 1. Let {81,..., Sx } be N input vectors and {uy = f a set of output vectors of M modules {f...., fMr ). As for the agreement routing alternating between Eq. 7 and Eq. p if any two input vectors identical, the coefficient eı 0 or 1 as the number of iterations T
13

---

(T) Proof. With T iterations, the output vector 2 of module fy can be expanded as
exp (ak (2))
(T+1) Uij exp ) Wij-
Divide both the numerator and denominator by the term exp (a ( we have
(T+1)
1 Enrts exp (a
(T)
Uij-
(9)
According to Eq. have
(T-1)
exp exp(a.1) (a (2)) = exp
W
8i
Ojj (1-11)
|00
(T'
Wa
W
T
(1)
W.
= exp
(10)
1=
When the iteration number T - oc, the value of oj thus the above equation can rewritten as
(T)
will converge fixed vector, denoted as vj.
exp (a ( exp (T
W
(11)
and this result becomes O or oo, if the sign of by Was; is negative or positive, Next we prove the
theorem by considering two cases: (1) di (T) is the largest value among the set {aly 11'=1" Assume
(T) that any two elements in the set (,1, 8N) are not identical, then o aiy - is negative number
(T)
for any j' # j. Thus, the term Σj' #) exp (a ai
(2)) in Eq. tends to 0 and Gt
(2) oij
(T)
not the largest value among the set {ap =1-=1" In this case, at least one element in set {a fj=1
(T) 1.M
(7)
larger than Oijz Thus, the term Σn exp (o ) + oo and Cij - 0.
B MODULATED MODULES
SMN does not specify the architecture of each module, as the agreement router utilize the agreement between input vectors and the output of each module. Thus, SMN can be implemented with arbitrary module designs, such as MLP, CNN, RNN or self-attention blocks. Inspired by previous works, we utilize modulated module design in our SMN, which is complementary with existing module architecture and useful to adapt to new tasks with few labeled samples. The modulated modules allows module condition its computation a small set of parameters. As the basis of other modulate module, we define Modulated Fully-Connected (ModFC) layer as
ModFC(a; e) = W( O LayerNorm (oyeeroom((Wec) + b,
(12)
where @ denotes element product, W Ridaut x Rd is weight matrix, is a bias vector, E Rdoods the conditioning code vector, and We E Rd x Rdoode is the conditioning weight matrix that maps code vector into input space. It replaces the traditional fully-connected layer and enables the input vector x to be multiplicatively conditioned by a learnable vector c. As FC layers be stacked form MIP, multiple ModFC layers and activation functions also be stacked to form Modulated MLP, which can be defined as:
y ModMLP(2;c)) c)
(M;dFC Activation o ModFC; (x)
(13)
In each modulated module, distinct computations can be executed conditioned on unique learnable vectors Consequently, when new modules introduced, only the new introduced code vectors need to be fine-tuned to adapt to new tasks. This property allows for sample-efficient fine-tuning on newly introduced modules.
14

---

Train from scratch
SMN w/o AR
SMN
1000 Conteiee mae aisenoaenn 1oM
training samples
6 XX Number training samples
Lm Number training samples
Figure 7: The number of training samples acquired for models to achieve test accuracy of 0.90 on Parity Code task. Left: we train a SMN without agreement router (SMN w/o AR) only on the training data of Parity Code task, which means knowledge
be pre-acquired. Middle: We pre-train SMN w/o AR on the Min-Max Digit task with 60,000 training data, and then fine-tune on Parity Code task. Right: We pre-train a SMN and then fine-tune it on the Parity Code task. We run 600 times for each setting with different random seeds and then report the frequency that achieve 0.90 test accuracy on Parity Code task.
C EXPERIMENTAL DETAILS
C.I IMPLEMENTATION DETAILS FOR MIN-MAX GAME
Generalization. We two convolution layers with ReLU activation fune- tion as the tokenizer for SMN. Besides, an average aggregation operation is used in the classifier component aggregate output vectors from modules and acquire single vector for classifica- tion. For fair comparison, we use the same convolutional network as backbone to extract features for other models, except the Transformers. We SGD optimizer optimize model parameters with learning rate 0.01. In all experiments, we train different models with enough epochs from (10, 15, 20} to achieve the training accuracy of 0.99. We set 0.001 in both the Min-Max Digit task and Parity Code task, and run all experiments 10 times with different random seeds then report the average results with standard deviation.
Few-shot Adaptation. We use the same model architecture as previous experiments. We train these models with 60,000 training images Min-Max Digit task and fine-tune them with & labeled samples on the Parity Code task. For fine-tuning, we fix the parameters feature extractor and only fine-tune the classifier (or classification head). For our SMN, when new modules added, fix the parameters of tokenizer and existing learned modules, and fine-tune the new added modules with classifier for adaptation. We use Adam optimizer with learning rate 0.001, and fine-tune each model with 100 epochs.
C.2 MORE RESULTS FOR MIN-MAX GAME
Sample Complexity. We also explore the task adaptation capability of SMN from the aspect of sample complexity. Fig. 7 shows that for SMN, learning from Min-Max Digit task helps reduce the required number of labeled samples in the Parity Code task (30-45 labeled samples in the right subfigure), greater than the number of samples for the same model but without agreement router (75- 105 labeled samples in the middle subfigure), and both outperforms than de novo learning of Parity Code task from scratch (22000-24000 labeled samples in the left subfigure). This comparison shows that the agreement router helps reduce the sample complexity in the transfer learning scenario.
C.3 EFFECT OF FINE-TUNING AGREEMENT ROUTER ON MIN-MAX GAME
In the min-max game, when a pre-trained SMN needs to adapt to the parity code task by adding new modules, fix all existing pre- pre-trained parameters, including the weight W. in the agreement router, and only fine-tune the newly add modules. We freeze W.. when new modules are introduced, with the consideration of scalability. Because fine-tuning W, could lead to additional computation load
15

---

and loss of learned routing information on old tasks. However, when we only focus on the current task, fine-tuning W, a good choice enhance the performance. Here we give more experimental Test Accuracy
Variants
add=0
add=2
add=4
Freeze Wa
87.9 (4.1)
88.7 (2.0)
88.7 (1.7)
Fine-tune W
91.5 (3.8)
92.6 (2.2)
89.0 (3.7)
Table 5: Test accuracy of SMN variants on parity code task with different number of added modules.
Method
FLOPs
Training Time
Inference Time
Accuracy
Top-K
7.6M
140s
2.3s
42.94 (2.1)
Truncated
10.8M
120s
2.3s
43.77 (3.4)
T 0
10.8M
80s
2.3s
44.55 (2.2)
SMN
T = 2
12.2M
130s
2.68
47.60 (4.1)
T 4
13.6M
160s
2.7s
54.46 (5.0)
Table 6: Computational cost analysis of SMN and other comparison methods in the min-max game experiment.
results to show the impact of fine-tuning Wa. In the Tab. [ we adapt the pre-trained SMN to the parity code task by finetuning both Wa and newly added modules. We observe that fine-tuning Wa leads to further performance improvement for this task. We argue this is because fine-tuned agreement router adapts to the parity code task, facilitating the selection of specialist modules.
C.4 COMPUTATIONAL COST ANALYSIS
For analyzing computational cost, report the FLOPs, training and inference time (scconds), and accuracy for our SMN and other modular variants in the min-max game experiment. The training time is caleulated 60,000 images with 10 epochs, and the inference time is calculated on 20,000 test images. We acquire these statistical information based on a single 3090 GPU.
Form Tab. 6 We can find that when the iteration number T 2, SMN exhibits comparable training and inference times to Top-K and Truncated methods, while outperforms them in term of accuracy. Notably, increasing iteration number T can improve the performance of SMN further, while bringing more computational burden.
C.5 DETAILS FOR FEW- SHOT IMAGE CLASSIFICATION
We consider three popular benchmark in FSIC task: minilmageNet (Vinyals et al. 2016), CIFAR-FS (Bertinetto et al., 2018) and 2UB-200-201 (CUB) (Wah et al.) 2011] minilmageNet is subset of ImageNet (Russakovsky et al., 2015) consisting of 60,000 images uni- splits consist of 64/16/20 object classes, respectively. CIFAR-FS is derived from CIFAR 100 (Krizhevsky et al., 2009), and it con- tains 100 classes with class split as 64, 16, and 20. CUB is fine-grained bird classification dataset, consisting of 100/50/50 bird classes for train/ httn//alldaiiottesss splits.
Dataset Details
formly distributed over 100 object classes. The train/validation/ee
Experiment Setup. As most previous works (Finn et al. 2017 Snell et al. 2017 Chen et al. 2019) we use the N-way K-shot setting to construct training and test episodes from dataset. means that we first randomly choose N classes from dataset, then choose K support samples and Q query samples for each class to construct few-shot episode. In each episode, few-shot model designed to classify these N X Q query images into N classes based on N x K support samples. Note that the classes of training episodes Ctrasn will not overlap with the classes of test episodes Clest, which means Ctrain n Ctest 0. More specifically, for the 5-way 1-shot (and 5-shot) evaluation, we test 15 query samples for each class in an episode and report the average accuracy with 95% confidence intervals of randomly sampled 600 test episodes.
16

---

f1
f2
f1
f2
f1
fz
(a) t = 0
(b) == = 1
(c) t
20
(d) 3
(e) t 4
Figure 8: Activation map of two modules within the SMN (T 4). We visualize the activation map at each iteration an example from composition (6,0).
Implementation Details. Note that our Scalable Modular Network do not specify the choice of tokenizer, any set-valued network that output set of vectors is feasible. We adopt the following two popular backbones as our tokenizer: (a) Conv4-64 (Vinyals et al. 2016). convolutional neural network which consists of 4 convolution layers with 64/64/64/64 filters for total 1.3. param- eters; (b) VIT-S (Lin et al.) 2023), a small size vision transformer which consists of 12 transformer encoders for a total of 21M parameters. For the former choice, we train the conv4-64 with module processor and classifier from scratch. For the latter, we first pre-train the vision transformer and projection head with a self-supervised loss following (Lin et al., (023). After that, the pre-trained model is used as our tokenizer, and we train it with other parts by the combined loss in Eq.
D VISUALIZATION OF ITERATIVE DYNAMICS
To better understand what occurs in each iteration for the agreement router, We visualize the acti- vation map of two modules fi and fo per iteration from trained SMN model. As shown in Fig & both modules fi and f2 are activated in response to the input image at the initial stage (t 0). As the iteration proceeds, module f1 gradually deactivates, while module f2 remains unchange. In the iterative dynamics, the input-output agreement
be treated as if it (t-1) as expressed Line of Algorithm added to the previous logit a,
liglliill and
!
Guided by the ac-
cumulated loglikelihood {a f the router models a conditional distribution to classify which specailist modules can acquire the input vector (as expressed in Line of Algorithm IT). although in fact, there does not exist ground-truth labels for the classification task.
E EVALUATION ON DIVERSE TASKS AND DATASETS
We conduct additional experiments to demonstrate the generality and effectiveness our SMN in diverse tasks and datasets, including reinforcement learning, continual learning and cross-domain few-shot learning tasks.
E.1 REINFORCEMENT LEARNING.
We consider a reinforcement learning task path finding from openai gym (Chevalier-Boites Boisvert et al. 2018) for evaluating the out-of-distribution generalization ability of models. More specifically, in the path finding task, an agent must open door using a key and then get to the goal with only partial observation of the environment. See example of the path finding task within 8x8 room at Fig 9 We train an RL agent using Proximal Policy Optimization (PPO) (Schulman et al. 2017) building upon the representations from modular network. The training environment is within
17

---

6
Figure 9: An example of the path finding task within a 8x8 room. The red triangle represents an agent, and the green square is the goal. Yellow square locked door, which should be opened with key.
5x5 room, while evaluations are conducted in different environments with 6x6 and 8x8 rooms. From Tab. 7 we can find that SMN performs better than other methods, and the agreement router plays an important role in the generalization ability of SMN.
Method
5x5 Room
6x6 Room
8x8 Room
Top-K
0.88 (0.05)
0.88 (0.05)
0.80 (0.16)
Truncated
0.88 (0.05)
0.88 (0.05)
0.74 (0.18)
T = 0
0,81 (0.09)
0.81 (0.09)
0.71 (0.22)
SMN
2
0.88 (0.05)
0.88 (0.05)
0.82 (0.20)
T = 4
0.91 (0.03)
0.91 (0.03)
0.85 (0.18)
Table Results SMN and other methods in the path finding task. We scored different models by computing the average total reward of the last 100 episodes.
min-max
parity
min-max
parity
add=0
96.12 (0.13)
87.48 (3.98)
add=0
78.17 (5.27)
91.5 (3.8)
add=1
95.42 (0.35)
88.61 (3.31)
add=1
75.19 (6.47)
92.2 (3.3)
add=2
94.50 (1.73)
89.32 (2.34)
add=2
73.57 (5.51)
92.6 (2.2)
Table 8: Performance of SMN on min-max digi- tal task (min-man and parity code task (parity) with different number of added modules.
Table 9: Performance of SMN on two tasks. We fine-tune the pre-trained Wa and newly added modules together.
E.2 CONTINUAL LEARNING.
We also construct a continual learning scenario to verify the adaptation ability of SMN without for- getting previous knowledge. In this scenario, min-max digital task and parity code task comes sequentially. The first task is comprised with sufficient training data, while the second task only contains limited training data. Solving this problem with SMN straightforward by applying dif- ferent specialist modules based on the task ID, parameter-isolatio methods in task incremental learning (Rusu et al. 2016). However, we explore more difficult setting to demonstrate the scal- abity of agreement router, where the task ID is unavailable in the module selection process. More specifically, We train SMN with modules on the first task, and then adapt to the second task by adding 1-2 new modules. During test, the agreement router camnot acquire the task information to select specialist modules. Tab. & shows that introducing new modules for the parity code can im- prove the performance, meanwhile only causing little decrease in the min-max task. We attribute
18

---

this to the scalablity of agreement router, observe that agreement router adaptively deactivate the newly added modules for the old task. Moreover, we explore the impact of fine-tuning Wa on performance. Tab. show that fine-tuning Wa the parity code task will hurt the performance in the min-max digital task, as the learned routing information on old task loses.
minilmageNet--CUU
minilmaggeNe--Skete
Method
Backhone
5-way 1-shot
5-way 5-shot
5-way 1-shot
5-way 5-shot
ProtoNet
33,91 (0.67)
53.74 (0.72)
41.05 (0.76)
63.26 (0.72)
RelationNet
38.19 (0.69)
52.57 (0.66)
46.33 (0.81)
59,46 (0.74)
MAML
Conv4
36.97 (0.69)
51.60 (0.70)
42.95 (0.83)
58.51 (0.78)
Baseline++
37.11 (0.66)
52.42 (0.67)
36.46 (0.75)
51.79 (0.77)
SMN
42.14 (0.72)
61.20 (0.73)
48.33 (0.86)
69.48 (0.75)
Table 10: 5-way 1-shot and 5-shot performance of different FSL methods on cross-dataset scenario (minlmageNet-cCUB and mindmmaeeNet-Skeee report the average accuracy on 600 novel tasks with 95% confidence interval.
Conv4-64 is used as the backbone model. We
E.3 CROSS-DOMAIN FEW-SHIO LEARNING
To demonstrate the effectiveness of our SMN more challenging tasks, we consider the cross- domain few-shot learning problem, where the source and target domains are dissimilar.
Datasets and Setup. Based on the setup Section 4.2, we consider two cross-domain few- shot learning scenario: (1) from minilmagenet CUB (mintlmagenet>>CUBB) minilmageNet general object recognition dataset and CUB fine-grained classification dataset; (2) from minilmagenet ImageNet-Sketch Uminllmagentt-SSetchh, ImageNet-Sketch (Wang et 2019) contains 50,000 sketch images belonging to the same 1,000 ImageNet classes, but exists large do- main gap with ImageNet images. For the former scenario, we use 100 classes of minilmageNer as training classes, and the 50 validation and 50 test classes from CUB. For the latter, we also 100 classes of minilmageNet training and novel classes non-overlapped.
training classes, and 900 classes from ImageNet-Sketer keep the
Results. Tab. shoos that our SMN owns a better adaptation ability than other FSL methods in both two domainn scenarios, which demonstrates the adaptation and oft-of- distribbutionn gener- alization ability of SMN.
19


Reviews:
[{"review_id": "zAHVbRmfa8", "reviewer": "Reviewer_nZgT", "summary": "The paper presents a novel approach for scalable modular framework. They propose a unique method for learning routing among a set of modules by averaging outputs across input tokens, computing agreement with each of the input tokens, refining the outputs based on this agreement, and finally aggregating these refined outputs for final  classification. The method showcases adaptability to a new task (parity code task) after training on a different task (min-max detection). However, the scalability and adaptation mechanisms, especially when new modules are introduced or when there are multiple layers of modules, are not clearly explained.", "questions": "1. How is \\(W_a\\) adapted when new modules are introduced, and how does this adaptation impact the performance ?\n2. Can you explain the relationship between \\(y\\) from Equation 4 and the outputs from the modules?\n3. How does the method work when there are multiple layers of modules?\n4. In Equation 12, is W_c frozen or trainable when new modules are added?\n5. With only two modules, what is the value of K in the Top-K method in Table 1, and how does this compare to an Ensemble approach where all modules are activated?\n6. Could you provide more insight on what occurs in every iteration of the proposed method, and elaborate on the notion of agreement in this context?\n7. Could you add more on the trade-off between computational cost and accuracy of the proposed method vs prior methods?\n8. Is there a plan to conduct experiments showcasing the scalability of the method, particularly when adapting to a new task by adding new modules without forgetting old tasks? You can choose the same setup of learning min-max detection and parity code task one after the other by adding new modules.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The method displays remarkable results on a toy dataset and adapts well to new tasks (parity code task)\n2. Better performance on the ConvNet backbone for few-shot classification task on real-world datasets as compared to previous methods.", "weaknesses": "1. The scalability of the proposed method is not clearly demonstrated, especially regarding how weights W_a in agreement router and W_c  in module parameters are adapted with the introduction of new modules.\n2. The paper proposes a method that works on a single layer, which is quite restrictive. But doesn't dive deeper into how the method translates when using multiple layers of modules.  \n3. A detailed computational cost analysis is missing, which is crucial for understanding the trade-offs, especially since the proposed method seemingly demands more compute by activating all modules repeatedly.\n4. The performance on the ViT backbone does not mirror the improvements seen on the ConvNet backbone, raising questions on the consistency of the method's performance across different backbones.\n5. The method needs additional losses terms to prevent degenarcy similar to Top-K approaches.", "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": "14VDAE5HBP", "reviewer": "Reviewer_9n4d", "summary": "This paper presents an innovative modular network framework, the Scalable Modular Network (SMN), which represents a significant advancement in the realm of adaptive learning. SMN not only enhances the capacity for adaptive learning but also offers a seamless mechanism for integrating new modules post pre-training, thus markedly improving overall adaptability. This adaptive process is underpinned by the ingenious \"agreement router\" integrated within SMN, which streamlines the module selection process by carefully weighing both local and global interactions. Throughout the experimental validation on two distinct problems, a toy min-max game, and a few-shot image classification task, SMN consistently showcases remarkable generalization abilities to new data distributions and demonstrates sample-efficient adaptation to novel tasks. Notably, SMN's adaptation capabilities are further enhanced when new modules are introduced after pre-training. These results underscore the potential of the SMN to continually evolve and improve its adaptive capabilities.", "questions": "Please kindly find the comments in the weakness section.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper effectively presents the motivation behind the study and articulates the issues addressed by the proposed method in a well-structured manner.\n2. The authors provide a comprehensive review of existing methods related to modular neural networks. They also meticulously analyze the drawbacks and issues that need resolution in these existing approaches.\n3. The Scalable Modular Network (SMN) is presented as a concise and efficient method with remarkable clarity in both its rationale and technical details. The paper offers ample explanations, discussions, and evidence to underpin the theoretical foundation of SMN. \n4. The inclusion of implementation details and experimental settings further substantiates the fairness of comparisons across two distinct problems. \n5. Moreover, the study hints at the significant potential value in large-scale modular networks capable of seamlessly integrating additional modules.", "weaknesses": "1. The experiments conducted on a toy min-max game and few-shot image classification may not be sufficient to definitively establish the overall superiority of SMN, considering that SMN is positioned as a general framework. Additional experiments on a wider range of datasets and tasks would better illustrate the effectiveness of the method\n2. The paper falls short in providing adequate explanations of limitations and training costs.", "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": "d0KqGaUG78", "reviewer": "Reviewer_y9sx", "summary": "The paper introduces the Scalable Modular Network (SMN), a modular framework with adaptive learning capabilities that can incorporate new modules after initial training for better adaptation.\n\n**Key Features:** \n- **Agreement Router:** A unique component in SMN that iteratively selects and assembles specialist modules based on both local and global interactions.\n  \n- **Dynamic Module Selection:** Allows SMN to adjust module combinations adaptively based on input data.\n\n- **Scalability:** Enables the addition of new modules post-training, enhancing adaptability.\n\n**Benefits:** \nSMN efficiently selects modules for new samples, generalizes for out-of-distribution data, and can differentiate between similar sub-concepts using global information.\n\n**Experiments:** \nTests on a toy min-max game and few-shot image classification demonstrated SMN's adaptive capabilities, especially when adding new modules post-training.\n\n**Significance:** \nSMN offers a solution to the challenge of composing specialist modules in neural networks, moving closer to achieving human-like learning efficiency in machines.", "questions": "See the weakness.\nAdditional Question:\n\n1. **Pre-training Paradigm:** The abstract mentions the applicability of SMN in pre-training paradigms, but there seems to be limited experimental evidence or discussion on this. How does the SMN truly perform in a pre-training context?\n\n2. **Comparison with Current Techniques:** Modern pre-training often employs methods like contrastive learning or MLM (Masked Language Model) loss. Is there a significant gap between SMN and these prevalent techniques? How does SMN align or differentiate from these established methods?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. **Agreement Router:** Introduces a dynamic module selection mechanism, mirroring human cognitive abilities for efficient and adaptive learning.\n\n2. **Scalability:** Allows for the integration of new modules post-training, ensuring the network's adaptability and evolution.", "weaknesses": "1. **Limited Experimentation:** The experiments are overly simplistic, conducted only on small datasets.\n \n2. **Lack of Broad Testing:** Results focus solely on classification tasks, with no results provided for large language models like LMM.\n\n3. **Additional Fine-tuning:** The addition of new modules still requires extra fine-tuning, implying integration isn't as seamless as desired.\n\n4. **No Comparison with MOE:** The paper doesn't offer a comparison with established methods like MOE, limiting the understanding of its relative performance.\n\n5. **Efficiency Overlooked:** There's no comparison or discussion regarding computational efficiency or time delays, making it hard to assess the practicality of deployment.", "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": "7VMOshLruG", "reviewer": "Reviewer_Zqvo", "summary": "This paper presents Scalable Modular Network for better adaptive learning by incorporating new modules after pre-training. An agreement router is proposed to select specialist modules using an iterative message passing process. The approach is evaluated with a min-max game task and few-shot image classification task.", "questions": "1. Fix typos, e.g, \"impose constrains on the number of activated...\"", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Modular networks have certain advantages in some machine learning settings such as meta-learning and continual learning.\n2. The proposed agreement router is novel and effective.", "weaknesses": "1. Some evaluation under continual learning setting may be desirable. Evaluation in the paper is not sufficiently strong with one toy task and one few shot learning setting.", "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)