Scalable Modular Network: A Framework for Adaptive Learning via Agreement Routing

Thank you for taking the time to review our work. We are glad that you deem the agreement router to be novel and effective, and the modular networks advantageous. In the following, we address your questions and concerns.

W1: 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.

Thank you for this suggestion. We try to verify the generality of SMN in the following three diverse tasks:

• Continual Learning. We consider a continual learning scenario to verify the adaptation ability of SMN without forgetting previous knowledge. In this scenario, the min-max digital task and parity code task come sequentially. The first task comprises with sufficient training data, while the second task only contains limited training data. Detailed experimental setting and results can be found in the response to Question 8 for the Reviewer nZgT. The experimental results demonstrate that SMN can adapt to new task better with new modules incorporated, without forgetting the old tasks.
• Cross-Domain Few-Shot Learning. To demonstrate the adaptation ability of our SMN in more chanllenging tasks, we delve into the cross-domain few-shot learning (CD-FSL) scenario. In this context, the source and target domains are dissimilar, thereby intensifying the difficulty of adapting to novel tasks. We conduct experiments on two setting: miniImageNet --> CUB and miniImageNet --> ImageNet-Sketch. See the experimental results in our response to Weakness 1 for Reviewer y9sx. We find that our SMN owns a better adaptation ability than other FSL methods, even in the more chanllenging scenario.
• Reinforcement Learning. We consider a reinforcement learning task of path finding to evaluate the out-of-distribution generalization ability of SMN. More specifically, in the path finding task, an agent must open a door by a key and then get to the goal with only partial observation of the environment. We train our SMN and other compared methods in the 5x5 room but evaluate them on the 6x6 and 8x8 rooms. The detailed experimental setting and results can be founded in the response to Weakness 1 of Reviewer 9n4d. We find that SMN performs better than other methods, and the agreement router plays an important role in the generalization ability of SMN.

Q1: Fix typos, e.g, "impose constrains on the number of activated..."

Thanks for pointing out this. It should be "impose constraints...". We have fixed it in the updated version.

Q1: There seems to be limited experimental evidence or discussion on pre-training paradigms. How does the SMN truly perform in a pre-training context?

We have discussed the pre-training paradigms when evaluating the fast adaptation ability of SMN: initially pre-training a model on some base tasks with sufficient data and subsequently fine-tuning it on novel tasks with limited data. More specifically, in the toy min-max game, SMN performs pre-training on the min-max digital task, followed by fine-tuning in the parity code task. Similarly, in the few-shot image classification task, the process begins with pre-training SMN on base training tasks, then fine-tuning on novel tasks. Experimental results demonstrate the adaptation ability of SMN in the pre-training paradigm.

Q2: Is there a significant gap between SMN and contrastive learning/masked language model loss? How does SMN align or differentiate from these established methods?

The SMN framework is complementary with existing self-supervised techniques, such as mask image modeling or contrastive learning. Given that the specialist modules and agreement router build upon the tokenizer, self-supervised techniques can be employed to initialize the tokenizer effectively, thereby enhancing module specialization. In few-shot image classification experiments, we adopt a two-stage pre-training process following [ref]. Initially, we pre-train a ViT-S backbone with a mask image modeling loss as our tokenizer. Subsequently, we pre-train the entire SMN model on training tasks using both supervised loss and an importance loss. We empirically find that using the self-supervised technique can significantly improves the performance of SMN model.

[ref] Supervised masked knowledge distillation for few-shot transformers. CVPR 2023.

Thank you for taking the time to review our work. We are glad that you were interested in the scalability of our method, and we appreciate your suggestions to enhance the impact of our paper. In the following, we address your questions and concerns in the order of being proposed.

W1: The experiments are overly simplistic, conducted only on small datasets.

To demonstrate the effectiveness of our SMN in more chanllenging tasks, we consider the cross-domain few-shot learning problem, where the source and target domains are dissimilar. Following [ref1], we consider a cross-domain scenario from miniImageNet to CUB, as miniImageNet is a general object recognition dataset and CUB is a fine-grained classification dataset. Besides, we consider another cross-domain scenario from miniImageNet to ImageNet-Sketch [ref2], where we select a subset of classes from ImageNet-Sketch to keep the training and novel classes non-overlapped. The results are as follow:

• miniImageNet -> CUB

  Method	5-way 1-shot	5-way 5-shot
  MatchingNet	37.48 (0.68)	49.98 (0.66)
  ProtoNet	33.91 (0.67)	53.74 (0.72)
  RelationNet	38.19 (0.69)	52.57 (0.66)
  MAML	36.97 (0.69)	51.60 (0.70)
  Baseline++	37.11 (0.66)	52.42 (0.67)
  SMN	42.14 (0.72)	61.20 (0.73)

• miniImageNet -> ImageNet-Sketch

  Method	5-way 1-shot	5-way 5-shot
  MatchingNet	44.38 (0.78)	60.76 (0.74)
  ProtoNet	41.05 (0.76)	63.26 (0.72)
  RelationNet	46.33 (0.81)	59.46 (0.74)
  MAML	42.95 (0.83)	58.51 (0.78)
  Baseline++	36.46 (0.75)	51.79 (0.77)
  SMN	48.33 (0.86)	69.48 (0.75)

We find that our SMN owns a better adaptation ability than other FSL methods, even in the chanllenging cross-domain scenario.

[ref1] A Closer Look at Few-shot Classification. ICLR 2019.

[ref2] Learning robust global representations by penalizing local predictive power. NeurIPS 2019.

W2: Results focus solely on classification tasks, with no results provided for large language models like LMM.

We apply SMN in diverse applications, including reinforencement learning and continual learning tasks. The experimental details and results for these two tasks are provided in the response to Weakness 1 of Reviewer 9n4d. The results across diverse tasks demonstrate the generality of our SMN and the effectiveness of the agreement router. As for building a large-scale modular network, such as those seen in large language models, is outside the scope of this work, we leave it into future exploration.

W3: The addition of new modules still requires extra fine-tuning, implying integration isn't as seamless as desired.

Although incorporating new modules requires extra fine-tuning to adapt to new tasks, a limited number of labled samples proves sufficient for this fine-tuning process. In the toy min-max game, we fine-tune the newly added modules using only 1-6 labeled samples per composition. Similarly, in few-shot image classification, we fine-tune the newly incorporated modules with just 1 or 5 labeled samples per class. This sample-efficient adaptation ability stems from the design of SMN framework, where the most parameters of new modules (e.g., $W$ and $W_c$ in Equation 12) is shared with pre-trained modules, as introduced in the Section 3.5 and Appendix B.

W4: The paper doesn't offer a comparison with established methods like MOE, limiting the understanding of its relative performance.

We have provided a fair comparison between our SMN and MoE in both Table 1 and Table 2 of main paper. The compared Top-K method follows the same training strategy as MoE, which contains a noisy top-$k$ gating mechanism and an importance loss during training. For a fair comparison, we adopt the same module structure as employed in our SMN for the Top-K method, and maintain the same number of modules. The results presented in Table 1 and Table 2 demonstrate that our SMN can outperforms Top-K method with the same number of specialist modules.

W5: There's no comparison or discussion regarding computational efficiency or time delays, making it hard to assess the practicality of deployment.

For computational cost, we report the FLOPs, training and inference time (seconds), and accuracy for our SMN and other modular variants in the min-max game experiment. The details and observations can be founded in our response to Weakness 3 & Question 7 for Reviewer nZgT. 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.

Thank you for taking the time to review our work. We are glad that you recognized the potential value of our work in large-scale modular network, and deemed our modular framework to be innovative and ingenious. In the following, we address your questions and concerns in the order of being proposed.

W1: Additional experiments on a wider range of datasets and tasks would better illustrate the effectiveness of the method.

Thanks for this valuable suggestion. We try to verify the generality of SMN in the following diverse tasks:

• Reinforcement Learning. We consider a reinforcement learning task of path finding from openai gym [ref1] for evaluating the out-of-distribution generalization ability of models. More specifically, in the path finding task, an agent must open a door using a key and then get to the goal with only partial observation of the environment. See an example of the path finding task in Figure 9 at Appendix E.1. We train an RL agent using Proximal Policy Optimization (PPO) [ref2] building upon the representations from a modular network. The training environment is set within a 5x5 room, while evaluations are conducted in different environments with 6x6 and 8x8 rooms. The results are as follows:

  	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)
  SMN ($T=0$)	0.81 (0.09)	0.81 (0.09)	0.71 (0.22)
  SMN ($T=2$)	0.88 (0.05)	0.88 (0.05)	0.82 (0.20)
  SMN ($T=4$)	0.91 (0.03)	0.91 (0.03)	0.85 (0.18)

  In the table above, we scored different models by computing the average total reward of the last 100 episodes, as previous RL works [ref2]. We can find that SMN performs better than other methods, and the agreement router plays an important role in the generalization ability of SMN.

  [ref1] Minimalistic gridworld environment for openai gym. 2018

  [ref2] Proximal policy optimization algorithms. Arixv 2017.

• Continual Learning. We also construct a continual learning scenario to verify the adaptation ability of SMN without forgetting previous knowledge. In this scenario, a 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. Detailed experimental setting and results can be found in the response to Question 8 for the Reviewer nZgT. From the experimental results, we demonstrate that SMN can adapt to new task better with incorporated with new modules, without forgetting the old tasks.

• Cross-Domain Few-Shot Learning. To demonstrate the adaptation ability of our SMN in more chanllenging tasks, we delve into the cross-domain few-shot learning (CD-FSL) scenario. In this context, the source and target domains are dissimilar, thereby intensifying the difficulty of adapting to novel tasks. We conduct experiments on two setting: miniImageNet --> CUB and miniImageNet --> ImageNet-Sketch. See the experimental results in our response to Weakness 1 for Reviewer y9sx. We find that our SMN owns a better adaptation ability than other FSL methods, even in the more chanllenging scenario.

W2: The paper falls short in providing adequate explanations of limitations and training costs.

Thank you for pointing out this. Here we will explain the training costs and limitations of our work:

• Train costs: The training time of SMN in the toy min-max game is 160 seconds on a single 3090 GPU, and in the few-shot image classification task is 3 hours on four 3090 GPUs. More details about the computational cost analysis can be found in our response to Weakness 3 & Question 7 for Reviewer nZgT. Notably, increasing iteration number $T$ can improve the performance of SMN further, while bringing more computational burden.
• Limitations: Agreement router needs to activate all modules and then iteratively calculates the input-output agreement, which increases the computational burden. Moreover, in this work, we demonstrate the effectiveness of our SMN on small size of datasets. However, we believe our work lays a foundation for adaptive learning, offering a start point for further research to train a large-scale modular network.

W4: Why the performance on ViT backbone does not mirror the improvements seen on ConvNet backbone?

In Table 1 and Table 2 of the main paper, the performance improvement gap between ViT and ConvNet backbone can be attributed primarily to two factors:

• Stronger backbone of ViT-based methods. ViT is obviously stronger than Conv. The commonly employed Conv4-64 backbone comprises 113K parameters, while ViT-S has 21M parameters. The performance of ViT-based methods surpasses that of the compared ConvNet-based methods by a significant margin. Further performance improvement with ViT backbone is more challenging compared to the ConvNet backbone.
• Relatively small-scale module processor. The commonly employed Conv4-64 backbone comprises 113K parameters, while Compared to the ViT-S backbone, the parameters of module processor are relatively small (444.7K for module processor, 21M for ViT-S), which may result in smaller impact on performance.

W5: The method needs additional loss term to prevent degenarcy similar to Top-K approaches.

Yes, we employ the importance loss to prevent degenarcy. The importance loss serves to balance the utilization of specialist modules, mitigating the collapse problem that the router selects a few modules repeatedly. However, the importance loss defined in Equation 5 is different from the loss in Top-K approach, such as MoE [ref]. In our work, we define the importance value of $j$-th module $imp(f_j)$ based on the agreement coefficients at the final iteration $\{c_{ij}^{(T)}\}_{i=1}^N$. Because the agreement coefficient $c_{ij}^{(T)}$ is accumulated across $T$ iterations, the importance value mixes the routing information across all iterations, which is absent in Top-K approaches.

[ref] Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. ICLR 2017.

Q2: Can you explain the relationship between y from Equation 4 and the outputs from the modules?

Sorry for any confusion. Equation 4 is designed to compute the cross-entropy loss between model's outputs and labels of sample. The variable $y$ corresponds to the label for input sample $x$. The variable $o$ represents the logits obtained from the classifier component applied to the aggregated outputs of modules $\{v_j^{(T)}\}_{j=1}^M$. This can be expressed as $o=\text{Classifier}(\text{Agg}(\{v_j^{(T)}\}_{j=1}^M))$, where $\text{Agg}$ denotes the aggregation operation, and in our experiments, we utilize the average operation.

Q5: 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?

The Top-K method employs the noisy top-$k$ gating mechanism to select K modules based on the outputs of router. With only two modules, we specify the value of K as 1. To compare with ensemble approach, we directly adjust the value of K to 2 and remove the noisy gating mechanism. Notably, We find that despite the ensemble variant achieves similar test accuracy on the seen digital compositions (approximately 96%), its performance decreases for the unseen digital compositions (approximately 43%). We attribute this discrepancy to the absence of modular inductive bias for ensemble approach, preventing it from effectively selecting and composing modules to deal with unseen digital compositions.

Q6: 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?

We provide a figure (see Fig 8 in Appendix D) to explain what occurs in every iteration. At the initial stage (t=0), both modules $f_1$ and $f_2$ are activated in response to the input image. As the iteration proceeds, module $f_1$ gradually deactivates, while module $f_2$ remains unchange. In the iterative dynamics, the input-output agreement can be treated as if it was a loglikelihood and is added to the previous logit $a_{ij}^{(t-1)}$, as expressed in Line 3 of Algorithm 1. Guided by the accumulated loglikelihood $a_{ij}^{(t)}$, the router models a conditional distribution to classify which specailist modules can acquire the input vector $s_i$ (as expressed in Line 4 of Algorithm 1), although there does not exist ground-truth labels for the classification task.

Q8: 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?

Thanks to this valuable suggestion. We construct a simple continual learning scenario where the min-max digital task and parity code task come sequentially. In this scenario, the first task comprises sufficient training data, while the second task only contains limited training data. Solving this problem with SMN is straightforward by applying different specialist modules based on the task ID, as parameter-isolation methods in task incremental learning [ref]. However, we explore a more difficult setting to demonstrate the scalabity of agreement router, where the task ID is unavailable in the module selection process. More specifically, We train SMN with 2 modules on the first task, and then adapt it to the second task by adding 1-2 new modules. During test, the agreement router cannot acquire the task information to select specialist modules. The results are as follow:

We find that introducing new modules for the parity code can improve the performance, meanwhile only causing a little performance decrease in the min-max task. We attribute this to the scalablity of agreement router, as we observe that agreement router adaptively deactivate the newly added modules for the old task. We also find that fine-tuning $W_a$ on the new task will hurt the performance in the old task, as the learned routing information on old task loses.

[ref] Progressive Neural Networks. Arxiv 2016.

Thank you for taking the time to review our work. We are glad that you deemed our method to be novel and unique. We appreciate the valuable questions and concerns, and try to address them in the following.

W1&Q1&Q4: How the weight $W_a$ of agreement router (in Equation 1) and $W_c$ in module parameters (in Equation 12) are adapted when new modules are introduced? And, how does this adaptation impact the performance?

In the SMN framework, we freeze both the pre-trained weight $W_a$ and $W_c$ when introducing new modules. We explain the reason of this setting and show how it impact the performance.

• Freeze $W_c$ in module parameters. In our SMN framework, most parameters of new modules (e.g., $W$ and $W_c$ in Equation 12) are shared with pre-trained modules. In this case, fine-tuning the pre-trained weight $W_c$ will destory the learned knowledge within existing modules, leading to a decline in performance, as illustrated in the table below.

  Variants	add=0	add=2	add=4
  Freeze $W_c$	87.9 (4.1)	88.7 (2.0)	88.7 (1.7)
  Finetune $W_c$	84.9 (3.8)	85.6 (2.2)	85.0 (3.7)

• Freeze $W_a$ of agreement router. With the consideration of scalability, we freeze $W_a$ when new modules are introduced, because finetuning $W_a$ will lead to additional computation load and loss of learned knowledge (see this result in the response to Question 8).

  Variants	add=0	add=2	add=4
  Freeze $W_a$	87.9 (4.1)	88.7 (2.0)	88.7 (1.7)
  Finetune $W_a$	91.5 (3.8)	92.6 (2.2)	89.0 (3.7)

However, when we only focus on the current task, fine-tuning $W_a$ is a good choice to enhance the performance. In the table above (last row), we adapt the pre-trained SMN to the parity code task by finetuning both $W_a$ and newly added modules. We observe that fine-tuning $W_a$ 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.

W2&Q3: How does the SMN work when there are multiple layers of modules?

To verify the efficacy of SMN in presence of multiple module layers, we stack two layers of modules in the module processor, comprising 8 and 2 specialist modules for each layer, respectively. In the multi-layer SMN, each module layer owns a unique agreement router for selecting different specialist modules. We evaluate the generalization ability of the multi-layer SMN on the new compositions, following the same setting of Table 1 in the main paper. The results are presented below. We find that stacking more layers of module can further improves the out-of-distribution generalization ability of SMN.

W3&Q7: A detailed computational cost analysis is missing, which is crucial for understanding the trade-offs between cost and accuracy of proposed method vs prior methods.

Thank you for pointing out this. For analyzing computational cost, we report the FLOPs, training and inference time (seconds), and accuracy for our SMN and other modular variants in the min-max game experiment.

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.