More Experts Than Galaxies: Conditionally-Overlapping Experts with Biologically-Inspired Fixed Routing
Authors: Sagi Shaier
Francisco Pereira
Katharina von der Wense
Lawrence Hunter
Matt Jones
Venue: ICLR 2025 Poster (2025)
Keywords: Deep learning
Mixture of Experts
Modularity
Sparsity
Conditional Computation
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We propose Conditionally Overlapping Mixture of ExperTs (COMET), a general deep learning method that induces a modular, sparse architecture with an exponential number of overlapping experts

Abstract: The evolution of biological neural systems has led to both modularity and sparse coding, which enables energy efficiency and robustness across the diversity of tasks in the lifespan. In contrast, standard neural networks rely on dense, non-specialized architectures, where all model parameters are simultaneously updated to learn multiple tasks, leading to interference. Current sparse neural network approaches aim to alleviate this issue but are hindered by limitations such as 1) trainable gating functions that cause representation collapse, 2) disjoint experts that result in redundant computation and slow learning, and 3) reliance on explicit input or task IDs that limit flexibility and scalability.
In this paper we propose Conditionally Overlapping Mixture of ExperTs (COMET), a general deep learning method that addresses these challenges by inducing a modular, sparse architecture with an exponential number of overlapping experts. COMET replaces the trainable gating function used in Sparse Mixture of Experts with a fixed, biologically inspired random projection applied to individual input representations. This design causes the degree of expert overlap to depend on input similarity, so that similar inputs tend to share more parameters. This results in faster learning per update step and improved out-of-sample generalization. 
We demonstrate the effectiveness of COMET on a range of tasks, including image classification, language modeling, and regression, using several popular deep learning architectures.

Published as a conference paper at ICLR 2025
MORE EXPERTS THAN GALAXIES: CONDITIONALLY- OVERLAPPING EXPERTS WITH BIOLOGICALLY-
INSPIRED FIXED ROUTING
Sagi Shaier
Francisco Pereira
Department of Computer Science University of Colorado Boulder sagi shai erecolrrclreeee edu
Machine Learning Core
National Institute of Mental Health
franci sco. perei ra gov
Katharina von der Wense Department of Computer Science University of Colorado Boulder
Lawrence E Hunter
Department of Pediatrics
University of Chicago
Institute of Computer Science
lehunt Soaunter@ucicaag..o.
Johannes Gutenberg University Mainz
katharin na kann cetoolrraoo. edu
Matt Jones
Department of Psychology and Neuroscience University of Colorado Boulder
Google DeepMind mcj0col crado. edu
ABSTRACT
The evolution of biological neural systems has led to both modularity and sparse coding, which enables energy efficiency and robustness across the diversity of tasks in the lifespan. contrast, standard neural networks rely on dense, non- specialized architectures, where all model parameters simultaneously updated to learn multiple tasks, leading to interference. Current sparse neural network approaches aim alleviate this issue but are hindered by limitations such as trainable gating functions that cause representation collapse, 2) disjoint experts that result in redundant computation and slow learning, and reliance on ex- plicit input or task IDs that limit flexibility and scalability. In this paper we pro- pose Conditionally Overlapping Mixture of ExperTss (COMET), general deep learning method that addresses these challenges by inducing modular, sparse ar- chitecture with an exponential number of overlapping experts. COMET replaces the trainable gating function used in Sparse Mixture of Experts with fixed, bi- ologically inspired random projection applied to individual input representations. This design causes the degree of expert overlap to depend input similarity, that similar inputs tend to share more parameters. This results in faster learning per update step and improved out-of-sample generalization. We demonstrate the effectiveness of COMET on a range of tasks, including image classification, lan- guage modeling. and regression, using several popular deep learning architectures. Code can be found here: https: http::///gttubb.../ com/ www.shaaaairraae git
INTRODUCTION
In recent years, there has been trend towards developing increasingly larger models (OpenAl 2023a by Fedus et al. 2022: Shuster et al. 2022 Chowdhery et al. 2022),, drrien by the understand- ing that a neural network's learning capacity depends its number of parameters (Shazeer al. 2017). This approach has yielded impressive results in various fields, including computer vision (Doooviskkyy et al. 2021 Kirillov al. 2023 and language modeling (OpenAl) 2023b Chowdhery

---

Published as a conference paper at ICLR 2025
However, with such large size come difficulties, including increased training costs and et al. 2022).
growing requirements for large amounts of memory and storage.
One approach to mitigating some of these challenges is sparsity, where a subset of the model's parameters is selectively utilized in the computational graph. This concept of sparsity has been
widely explored in machine learning (Jacobs al.] 1991; Jordan & Jacobs
1993: LeCun et al.
1989 [Zhou et al. 2019 Hoefler et al. [021] [Shazeer et al. [017]] Bair et al.) 2024).
Researchers
have observed significant benefits of sparsity, including reduced inference costs (Han et al. 2016
Shazeer et al. 2017), improved generalization capabilities (LeCun et al.) 1989 Jacobs & Bur Burkholz
2024 Frankle Carbin 2019 Paul al. 2023), enhanced learning efficiency (LeCun al. 1989
accelerated learning speed (LeCun et al. 1989; Mittal et al. 2022). less interference and forgetting
(Raia 2020), forward knowledge transfer (Andle et al. 2023 Raia 2020: Yildirim et al.
2023 and
compositionality (Pfeiffer et al. 2024).
Early work on sparsity in neural networks focused on simple methods such as Dropout (Srivastava)
et al. 2014) and LI regularization (Tibshirani, [996] Ng. (004).. These and subsequent works ex-
plore sparsity at a fine level of granularity, including single parameters (Mallya et al.. 2018: Mallya
& LLZzebnik 2018), individual neurons (Xu et al., 2024). CNN filters (Chen et al. 2019). Other
research has explored sparsity at the level of whole networks or sub-networks within the mixture of
experts (MoE) framework (Shazeer al. 2017 Bengio 2013). These methods generally utilize
routing or gating function (Rosenbaum et al. 2017 2019 Shazeer et al. 2017 Pfeiffer et al. 2024),
which decides which parameters sub-networks of the model to activate based on the input.
However, existing sparse methods have limitations, which we would summarize as five key con-
cerns: Firstly, most approaches rely on trainable gating functions (Shazeer et al. 2017
Mostafa &
Wang 2019
Rasmussen & Ghahramani] [200T: Li et al., 2023; Rahaman et al.] 2021
Chen et al.
2019 Shazeer et al.) 2018 Zhou et al. [2022] Gururangan et al. [2022] Lin et al.) [021] Ba & Frey
2013 Bengio et al.
2016 Fedus et al. [022 Mallya et al.) 2018 Mallya & Lazebnik 2018 Fer-
nando et al. [017]] Keshari et al.J 2018p. This design choice is problematic for several reasons,
including forgetfulness in continual learning (Preiffer et al.) 2024 Raia 2020). representation col-
lapse (i.e., degenerate experts; Chen et al.)
PPeiffer et 20.44) complex training procedures (Rosenbaum et al. 2019), and other issues (Rosenbaum et al. 2017]] Pfeitfer et al. [0244] Rosenbaum
ct al. 2019). Moreover, using non-trainable routing functions can be more effective (Mittal et a
2022 Mugeeth al. 2022). Secondly, many state-of-the art systems employ architectures based on disjoint experts that do not share parameters (Shazeer et al. 2017] Pfeiffer et al.] 2024). This design choice can lead to redundancies and may limit generalization; overlap can also be beneficial (French) 1993 Maini et al. 2023). Thirdly, even when experts overlap, it is unclear whether models can effectively learn map similar inputs the same experts, potentially resulting in redundancies
(Chen et al. 2023) or interference (Pfeiffer et al., (022)). Fourthly, many existing methods require
input or task
IDs determine which mask apply (Mallya et al. 2018: Yang et al., 2020 Masse
et al. 2018
Maini et al.] 2023 Pes et al. 2024; Mittal et al. 2022] Mugeeth et al. 2022 Kang
ltt al. 2024
Wortsman et al.] 2020). which can be restrictive, as meta- intommnnfonntion about inputs is
rarely available in real-world applications (Aljundi et al. 2019 Ye & Bors 2022 Wang et al. (022).
Lastly, the number of experts current systems is limited, often ranging from a few to a couple of thousand (Shazeer et al., 2017: Jiang et al. 2024), which may not be sufficient for complex tasks (Rasmussen & Ghahramani 20011)
In this paper we introduce Conditionally Overlapping Mixture of ExperTs (COMET), a general deep learning method that induces a modular, sparse architecture in neural networks, with a number of important properties. First, COMET uses a non-trainable gating function, eliminating the need for iterative pruning or continuous sparsification. Instead, we employ a fixed random projection fol- lowed by k-winner-take-all cap operation, inspired by the brain's efficient use of limited number of active cells via lateral inhibition. As in the brain, these mechanisms combine to produce sparse representations with overlap that depends on input similarity (Bruhin & Davies 2022). Second, COMET does not require fixed specialization of sub-networks or advance knowledge of the active neurons required for each task, enabling more flexibility and adaptability. Third, the number of possible experts in COMET is exponential in the model size, exceeding the limit of few thousand in recent work, to effectively tackle more complex tasks. Fourth, these experts overlap based on unsupervised information from input similarities. This yields faster learning and improved general- ization. It does this without increasing the number of trainable parameters, or requiring input or task IDs to determine which mask to apply.
2

---

Published as a conference paper at ICLR 2025
COMET integrates concepts from diverse research areas into concise framework: fixed random projection and t-winnerrtake-alll from neuroscience, routing functions from modular neural net- works, expert-based approaches from the MoE literature, the notion of implicit experts from dy- namic neural networks, the integration of sparsity and modularity from conditional computation, input-dependent masking from various deep learning areas, and the importance of active parameter overlap from continual learning. The present paper focuses on learning and out-of-sample gen- eralization in single tasks, but we conjecture COMET's input-dependent sparsity will also yield advantages for settings involving multiple tasks, including transfer learning. continual learning, and robustness to catastrophic forgetting. These will be tested in future work, although we report a preliminary test of transfer learning in appendix AA
We validate our approach through experiments on seven diverse tasks, including image classifi- cation, language modeling, and regression, demonstrating that our method is applicable to many popular model architectures such as vision transformers, MLP-mixers, GPTs, and standard MLPs, and consistently provides improved performance.
2 RELATED WORK
Sparsity Sparsity in deep learning refers to systems where not all parameters are active or partic- ipating in the computational graph. Fixed sparsity, such as L.1 regularization, intends to minimize the number of active parameters through optimizing continuous loss function. Variable sparsity includes various methods, such as randomness as in Dropout (Srivastava et al., (2014), and input- dependent, such as in sparse MoE (Shazeer et al. 2017). Notably, continuously optimizing deep networks to be sparse is an NP-hard problem (Jacobs & Burkhhol, 2024), and pre-defined sparse architectures can be restrictive. To circumvent that, propose sparsification using biologically motivated approach of random projection followed by a cap operation, which activates the strongest cells in the network, similar the sensory system of fruit flies (Bruhin & Davies 2022).
Modularity Modularity is related to sparsity, where information is conditionally routed to a subset of the network's 's parameters (Preiffer et al. 2024 Rosenbaum al.] (017). Modular methods can be split into two types. Firstly, those with trainable routing functions, including the sparse MoE (Shazeer et al., 2017), adaptively determine the active parameters during training. Although widely used, this approach can lead to representation collapse, forgetfulness, and redundant computation (Pfeiffer al., 2024). Secondly, fixed routing function methods have been shown to be more ef- fective (Mittal et al. 2022 Mugeeth et al., 2022). However, these methods typically require prior knowledge of module specialization (Pfeifter al., 2024; Muqeeth et al., 2022 Mittal al. 2022), which is often not available in practice. In contrast, our work demonstrates that module specializa- tion can be achieved in a fully unsupervised manner using fixed random projections. Unlike previous studies that employed fixed random projections routing functions, either with non-overlapping ex- perts (Roller et al. 2021] Chen et al. 2023) or with reduced performance (Bruhin & Davies 2022), our approach focuses
the more challenging setting of overlapping experts and achieves large performance gains across diverse tasks. Additionally, our expert overlap correlates with input simi- larity, making it more likely that parameters are shared between similar inputs. This, in turn, enables positive knowledge transfer between items, resulting in faster learning and improved generalization. Moreover, as the experts overlap, our approach benefits from input-dependent sparsity without an explicitly modular architecture (like in MoE).
Conditional Computation Conditional computation integrates sparsity and modularity, in that pa- rameters are dropped in a learned and optimized manner, rather than randomly and independently (Bengio 2013). A notable example is the sparse MoE framework (Shazeer et al., 2017). Our ap- proach diverges from previous work by uniquely integrating several desirable properties, including fixed routing, overlapping expert assignments based on input similarities, mask determination with- out requiring meta-information. and scalability to exponentially large number of experts.
Dynamic Neural Networks Dynamic neural networks (Han et al. 2021) share similarities with conditional computation. Unlike traditional MoE approaches, dynamic neural networks do not have explicitly defined experts. Instead, they dynamically select units, layers, or components from the main model for each input (Han et al. (021]. Our approach has parallels with some dynamic neural network methods, such as Piggyback (Mallya et al. 2018) and PackNet (Mallya & Lazebnik| 2018), which also do not define explicit experts. However, method differs two key ways. Firstly,
3

---

Published as a conference paper at ICLR 2025
does not require meta-information to determine the expert mask. Secondly, it ensures that any given input always maps to the same expert throughout both training and inference. This contrasts with architectures that use trainable gating functions, where the same input activate different experts or representations at different stages of learning, which can be seen as a "moving target".
The example-tie tied dropout scheme Maini al. (2023) is another dynamic neural network in which neurons are partitioned into a "generalization pool that is active for all inputs and a "memorization" pool for which small subset active for each input. The assignment of memorization neurons to each input example is random and fixed throughout training. Our approach differs from example-tied dropout that similar inputs share more active neurons, which facilitates generalization.
The model superposition of [Chuugg et al. (2019) applies task-dependent orthogonal transformation to the activation vector each layer. Assuming the input distribution has effective dimensionality less than that of the activation space (i.e., the layer width), these transformations map different tasks into different subspaces allowing them to be learned comi-independently Our input-dependent masking also restricts the representation of each example subspace, but it does not require task IDs and it leads similar inputs to have more overlapping representations.
MIXTURE OF EXPERTS AND INPUT-DEPENDENS MASKING
A standard MoE architecture involves disjoint set of experts and a gate that combines their pre- dictions (Jacobs et al. 1991). For example in the sparse MoE framework proposed by Shazzer et al.) (0177) each MoE module consists of n expert networks, Ej..., En, and gating network, G, that outputs sparse ne-dimensional vector of mixture weights. The gating and expert networks are all trainable, each with its own set of parameters. The prediction for an input x is Σ F..))..
Separately, several recent works have proposed versions of nependdeenneet masking (Maini et al., 2023 Mallya al., 2018 Mallya & Lazebnik 2018] Yang al.] 2002; Hung et al.) 2019). The gen- eral framework vnvolves a network with n neurons and a masking function m: X' + {0. 1 (where X is the input space). processing an example æ, the network's activations are multiplied elemen- twise with m(x). Thus the prediction for is (m( (3)) where Fin(a) is the function computed by the sub-network corresponding to the mask m(x).
Combining these two lines of work, we propose to view the sub-networks defined by input- dependent masking as overlapping experts. Any two experts will typically share many active neu- rons, and hence weights. This is in contrast to standard MoE where the experts learn disjoint sets of parameters. In the overlapping MoE framework, every subset of the full network is (potential) expert. The sub-network (m)(z) is an expert for æ. and it is also a partial expert for any other x' to the degree that m(x) and m(x') overlap, as determined by the inner product m(x)) m(().. ). In the overlapping MoE framework, the gating network G is replaced by the masking function m.
We further propose that similar inputs should map to similar (i.e., more overlapping) experts. This will facilitate generalization because what is learned about one input will selectively generalize to similar inputs. The next section cxplains how COMET achieves this property using biologically inspired fixed random projections (V, in eq. (4) and k-winner-take take-lll capping (Ck, in eq. (5))
CONDITIONALLY OVERLAPPING MIXTURE OF EXPERTS (COMET)
Our proposed COMET method applies to any backbone NN, augmenting with second NN called a routing network that computes input-dependent masks for all layers of the backbone network.
We first describe the COMET architecture for the where the backbone network is MLP. Let the backbone MLP have L layers, with layer i having N₂ neurons and learnable parameters comprising weight matrix W, E NNN.NN, and bias b, N Then the forward pass of the unmodified MLP is defined by
0 = W + b
(1)
0; = saar)
(2)
for 1 SES L, where a is the pre-activation at layer e, f is the elementwise activation function, Z0p is the input to the network, and as is its output.
4

---

Published as a conference paper at ICLR 2025
COMET Layer
Puly Cornaded
Layer
Lager
Conneced
Layer
Pardon
Mal Connaded
Input
Layer
Figure 1:
Illustration of a 2-layer MLP with embedded COMET layers. Note that COMET layers do not contain predefined experts, but instead dynamically selects subset of the backbone MLP's parameters to activate, effectively creating implicit experts. The sparsity level determines the pro- portion of parameters activate. Real value teal, zeros in white, ones grey.
COMET's routing network is a second MLP with the same shape, defined by random weight matri- ces Vz (for simplicity we omit bias parameters). We sample Vi from the same distribution used for initializing We (U(-N=1 N/= in our experiments). We denote this network's pre-activations and activations as C₁ and Zy (analogous to a₂ and øt in the backhone network), with input 2 = Zn- The computation of the routing network is similar to that of the backbone MLP, except that each layer it computes a binary vector m, that is then used to mask the activations in both networks. The mask is computed using a k-winner winnerrttaeealll take-all capping function C+:
1 Kj U 2 w/ s
[Cx(o)]:
(3)
0
otherwise
We allow fixed proportion Pt. of neurons each layer survive the mask, that ke Nr. Then the forward pass of the routing network is defined by
C₁ ViZe-1
(4)
m, Ckee (ce)
(5)
21 = my g(cs)
(6)
where indicates elementwise multiplication and g the routing network's activation function (we use the identity g(c) in the present experiments). The layerwise masks m, computed by the
routing network are then applied the backbone network, so that eqs. and are replaced by
a₂ W
32 m. f(ar))
Note that the network's 's output az is computed before m.s would be applied, avoiding undesirable masking of the model's prediction.
This input-dependent masking results in a maximum number of experts that is exponential in the model size at each layer, specifically (%r) Therefore practical settings every input will have its own expert. One consideration for this calculation might be interference among experts. Previ- work has studied how multiple models can be superposed within one network (Cheung al. 2019), and Ellhage et al. (2022) show a layer with N₂ neurons can hold O(e~) representations that pairwise orthogonal within certain finite tolerance, thus minimizing interference. However, more important for the present work is that overlap between models a desired property because promotes generalization between similar inputs.
COMET layers differ from sparse MoE (Shazeer et al. 2017) layers in two major ways:
1. Architecturally: Whereas layer in standard layered MoE architecture consists of experts and a gating network, a COMET layer contains a random, non-train non-trainable matrix and s-winner-take-wee cap operation. Instead of pre-defined experts, COMET layer modifies the computation the MLP to activate only a subset of its parameters contingent on the implicit expert.
input; this subset can be
5

---

Published as a conference paper at ICLR 2025
2. In the way the information is passed: Sparse MoE is applied in layers with new gating network at each layer, which takes as imput the backbone activation at the previous layer. Thus the gating and backbone networks each layer take the same input. In COMET architecture the routing network operates independently inputs to the given example maps to the same implicit expert throughout both training and inference.
the backbone network, so the distinct (except for the first layer of the network). This ensures that
Several other important differences between existing approaches and COMET are: 1) COMET's gating network does not require training; 2) COMET does not require fixed specialization of each network module, or advanced knowledge the combination of modules required for each task; 3) the experts in COMET overlap based unsupervised information from input similarities; 4) COMET does not require input or task IDs to determine which mask to apply: 5) the number of possible experts in COMET is exponential in the model size.
EXPERIMENTS
5.1 SYNTHETIC DATA EXPERIMENTS
In this section, describe experiments to verify key properties of a COMET network. First, verify that the combination of the fixed routing function and cap operator maps similar inputs to similar masks and show how this sharpens the model's generalization. Second, we verify that the network makes an effective use of the available neurons.
5.1.1 FIXED INPUT-DEPENDENT ROUTING NETWORK
Maskc Average Costre Simbarity
0.k 0.2
Agliratiorss Naarguu Coadte Semturky
K. 02
P.K. 0.3
P.K. 0.3
0.
Puk:
P.k
k.
P.k.
0.6
kk 0.7
k: 0.7
k: 0.8
0.4
P.k
0.2
0.0
Cosine Smillntty
Costne
(a)
(b)
Figure 2: Routing properties our gating function, which combines fixed random projections with a cap operator. (a) We compare the similarity of input pairs to the similarity of their corresponding binary masks (gates) for different sparsity levels. This plot shows that similar inputs tend to have similar masks. (b) We compare the similarity of input pairs to the similarity of their corresponding masked activation vectors the backbone network. This plot reveals that similar inputs are mapped to similar activations, and that this relationship is sharper for sparser networks (note Pk 1 is vanilla MLP). These properties facilitate forward knowledge transfer, even without supervision.
Our goal is to develop fixed routing function that maps similar inputs similar (i.e., overlapping) experts, thereby facilitating knowledge transfer between items and leading to faster learning and improved generalization. One way approximate generalization between individual training and test items is with the neural tangent kernel (NTK: et al. 2018). Let = (W1, b,...., Wr, br.) denote the flattened concatenation of all model parameters, let taain and test be arbitrary training and testing items, and let arrin and a be the corresponding model predictions for some fixed setting 0. Generalization from x train to test can be defined as the change in prediction a from including a train the training set. Formally, under vanilla GD optimizer on loss L, and in test is the limit of small learning rate a, the contribution of x train to change in a
1
Da
test - K K(x (artraiin x test V train ctrain
(9)
a-30

---

Published as a conference paper at ICLR 2025
where K (attrinn x(titt is the N2 x Ny matrix-valued NTK
K train
x: rest)
test 00
(
(8at.min
(10)
00
The RHS of eq.
(00) sums over elements of 0, and the contribution from We is
Σ dak aWe,s
(Dat aWess
trainn
Daya Dat
best
Da Dain
Σ
Σ 1 2-j
test
train
(1I)
Thus the contribution of We to generalization from train to test is proportional to the inner product (a az=1). This inner product will be positively related to input similarity (a.train x gestt) even in an unmodified MLP, but our question is how the relationship changes under COMET.
To answer this question, we conducted an experiment using 500 random pairs of input vectors. Each input had length 100 with components sampled iid from M(, 25), with μ a random integer between O and 100. We calculated the cosine similarity (i.e., normalized inner product) between each pair of inputs. We then randomly initialized 10 COMET networks for every pair, each comprising a backbone network and a routing network which were both MLPs containing 10 hidden layers (L 11) with 512 neurons per hidden layer. We passed both inputs through each of the 10 COMET networks, using eqs. to (8) with varying degrees of sparsity Pk-
To analyze the behavior of our fixed gating function, we performed two complementary analyses. First, we computed the cosine similarity between the masks obtained for the two inputs in each pair, concatenated across layers as (m1..... m.-1). Note that cosine similarity between binary vectors equals their degree of overlap, i.e. the proportion of active neurons for one input that are also active for the other. Second, we measured the cosine similarity between the two inputs' representations in the backbone network after applying the gating function in eq. (8) again concatenating across layers as (.),,..,, -I) To obtain a more robust estimate, we averaged the cosine similarities across the 10 COMET networks for each input pair, yielding the results in fig. 2
This experiment reveals that when input distributions are more similar the overlap between their binary masks increases (fig. 2b). This in turn strengthens the relationship between input similarity and activation similarity in the backbone network relative to the baseline MLP with px = (fig. b Drawing on the NTK analysis above, we conclude that COMET's routing function leads the model to generalize using a narrower effective kernel. A narrower kernel should not be expected to yield universal improvement, but should be beneficial when the base model has excess capacity for the task. The experiments in the next subsections support this prediction, in that we see an advantage for COMET particularly with larger models.
5.1.2 EXPERT UTILIZATION
MAUTO orren
LSE
Nours1
120
200
220
Lega
Laper3
P.K
Un
Figure 3: Illustration of neuron activity across COMET layers in a 4-layer MLP. We visualize the utilization of neurons in two randomly initialized networks with varying sparsity levels, using the CIFARIO dataset. The plots show that our network effectively utilizes all its parameters, with no "dead neurons" and no signs of representation collapse, even at very high sparsity levels.
7

---

Published as a conference paper at ICLR 2025
One of the challenges in training sparse architectures is representation collapse, where small sub- set experts or neurons becomes dominant, leading to -nder-utilization of others. This issue is particularly concerning when training gating networks, it can result in "dead" experts or neurons that are never activated. Given that our gating network is fixed, is essential to investigate whether this behavior occurs, it would be permanent.
To address this, we conducted an experiment where we generated 1000 randomly initialized 4-layer COMET MLPs with varying neuron counts per layer N₂ U (100, 1000), and varying sparsity level 1 Ph- ~ (0.05, 1). and passed the CIFARIO dataset through each one. We then analyzed the utilization of each neuron, measured by how often it is activated with the given masks, and found that only 7.6e 4% of neurons across the ensemble had 0% utilization, and fewer than 2% of the models had any such neurons, mostly models with high sparsity level. Figure 3 presents utilization plots for two representative networks, illustrating that our fixed gating network avoids representation collapse and "dead neurons." Moreover, when we passed a different dataset (such as CIFAR 100) through these models, we discovered that many previously inactive neurons were now being utilized. Thus, neurons that are infrequently utilized appear to be reserved for unseen data, highlighting the network's 's adaptability and capacity for generalization.
The finding is consistent with our previous analysis in section [.1.1] which showed that the input- dependent gating design inherently activates similar parameters for similar inputs, facilitating for- ward knowledge transfer. Our results suggest that our approach effectively mitigates the risk of representation collapse and promotes healthy utilization of neurons in the network, all without rely-
ing on supplementary mechanisms, such as specialized loss terms (Shazeer et al. 2017).
5.2 IMAGE CLASSIFICATION
We extend our investigation by integrating the COMET method into a diverse range of popular architectures, including Vision Transformers (ViTs), MLP-Mixers, and standard MLPs.
5.2.1 STANDARD MLP - CIFARIO
We apply the COMET method to a standard MLP with 4 layers, varying the number of neurons in each layer and the sparsity levels. To evaluate its performance, we compare it to 10 related methods:
Standard Model: A standard MLP model with the same number of neurons and no sparsity.
Smaller Model: A smaller model with reduced number of neurons, specifically Pv where Ne is the width of the standard model.
Dropout Model: A standard model with dropout rate equal to Pk-
An MLP with a trainable routing function. The cap operation is applied directly to the backbone network by replacing eq. (5) with my = Cke (az), so that the routing function selects the highest & values and masks the remaining ones.
Topk Model:
MoE Trainable: A MoE model with [1/pk] experts, each having N&N neurons in each layer. The routing network is a trainable MLP with one hidden layer and a sparse [1/pk] -dimensional output.
MoE Non-trainable: Same as MoE Trainable, with fixed routing function.
Layer-wise Routing: An MLP where each backbone hidden layer representation is projected using random matrix, which is then used to develop the binary mask for the next layer. This is done by replacing eq. with C₂ = Vi
Bernoulli Masking: An MLP where each training example is associated with fixed binary mask drawn from a Bernoulli distribution, with probability equal to px- Thus the relationship between inputs and their masks is arbitrary, rather than being mediated by the routing network in COMET.
Example-tied Dropout: Example-tied dropout (Maini et al. 2023), where each example in the training data is associated with fixed binary mask drawn from a Bernoulli distribution, with prob- ability equal to and fixed number of "generalization neurons" are active for all examples.
Standard model L1: A standard ML.P model, but using I.1 regularization to induce sparsity.
8

---

Published as a conference paper at ICLR 2025
hourons-100
Neurins-S00
heurons- 1000
Reuronn--30 3000
6.4
6:2
6.41
0:2
0.3
0:0
2.00
034
6144
634
0:4
1:20
0.20
0.0
Valderom
0.2
8.0
0.20
220
COMET
JoyecknscJooutnng
dropout
Figure 4: Illustration of 4-yyer MLP networks trained on CIFAR1O, showcasing the impact of varying network capacity and sparsity levels. As we increase the number of neurons and decrease sparsity (moving from top left to bottom right), we observe shift in the best-performing model. Initially, the standard model outperforms the COMET model when network capacity low. How- ever, network capacity grows, the COMET model emerges the top performer.
We evaluate these models on the CIFARIO dataset Krizhevsky (2009), with results shown in fig, Overall, the optimal model architecture depends the capacity of the network. When the number of neurons is limited and the network has a low capacity to learn the task (i.e., low Pt). the standard model that utilizes all neurons outperforms most models. However, as network capacity increases with more neurons, the COMET model emerges as the top performer. This suggests that the ben- efits of selective neuron activation become more pronounced capacity increases. Notice in the high-capacity regime the Smaller Model matches the Standard Model, indicating that simply adding more neurons does not improve performance while adding neurons subject to COMET's structured sparsity does.
5.2.2 CONTEMPORARY ARCHITECTURES
We further extend the COMET method to contemporary architectures in the Vision domain, includ-
ing ViT Dosovitskiy et al. (021) and MIP-Mixer Tolstikhin et al. (2021). To do this, we apply the COMET random projection followed by the cap operation in the MLP layers of each with px 0.5.
We evaluate the performance of these models on four widely-used image classification datasets:
SVHN Netzer al. (20TT) CIFAR Krizhevsky (009). CIFAR100 Krizhevsky (2009) and Tiny
ImageNet Le & Yang (201)). Our results can be seen in figs.
10
and
similar trend emerges these architectures: as network capacity increases. the optimal model architecture shifts. In smaller networks, where the number of neurons in the MLP layer is limited, the standard model performs roughly similarly to the COMET model. However, even in these net- works, incorporating COMET layers yields notable performance improvements. As we scale up the network by adding more neurons, COMET displays superior performance across all five model ar- chitectures and four datasets. It achieves faster convergence and significantly higher accuracy, with gains of up to 9% in VIT Large on CIFAR100, 100. Moreover, we observe that the performance gap be- tween the COMET- COMET-based models and their standard counterparts widens the model size increases, with larger models exhibiting both better performance and faster learning rates. This reinforces our finding that selective neuron activation becomes increasingly beneficial network capacity grows.

---

Published as a conference paper at ICLR 2025
CIHARI0O
Tiny ImageNet
50-
4.
40:
Vallidation Accuracy
-32
30-
Validation Accuracy
26
20-
176
200
202
200
202
220
203
202
220
100
250
50-
30
30:
20-
200
203
400
ASC
Epoch
Epoch
Figure 5: ViTs and MLP-Mixers on CIFARI00. Figure 6: ViTs and MLP-Mixers on Tiny Ima-
[Highest accuracy] trainable param.)
geNet. [Highest accuracy] (# trainable param.)
wikitext
codeparrot
CONST
14
CONT
CHAR7 Med
CORET (TAA
35
Val adtion Perp
35-
35
Validation Perpl
12
12
12
30-
30
10
10
10
25
25
20
4
100
200
1.00
200
100
200
1000
250
500
100 200 300
Step
Step
Figure 7: GPTs trained on WikiText. [Lowest Figure 8: GPTs trained on CodeParrot. [Lowest
perplexity] (# trainable param.)
perplexity] (# trainable param.)
5.3 LANGUAGE MODELING AND REGRESSION
We apply COMET language modeling on Wikitext (Merity et al.) 2016) and CodeParrot (Tunstall
et al. 2022) with varying GPT model sizes, with results in figs. 7 7nndd g We again observe that as network capacity increases, the COMET model outperforms the standard model, with larger models exhibiting not only greater performance difference but also faster learning rates, highlighting the benefits of selective neuron activation in language modeling tasks. To further validate our results, also evaluated COMET on the SARCOS regression dataset. Our findings in appendix Ahow that our conclusions generalize to this setting as well.
CONCLUSIONS
In this work, we propose a sparse neural network method, COMET, that induces a modular, sparse architecture with an exponential number of overlapping experts and alleviates key limitations of existing modular approaches, including trainable gating functions that often lead to representation collapse, non-overlapping experts that hinder knowledge transfer, and the need for explicit input IDs. By leveraging biologically-inspiree fixed random projection and k-winner-take-al capping operation, COMET determines expert overlap based on input similarity in a fully unsupervised manner. enabling faster learning and improved generalization through enhanced forward transfer. should be noted that key features of COMET may occur naturally in standard networks. First, similar inputs already tend to have similar hidden representations (see the Pke = 1 curve in fig. (bb). Second, there may be cases where networks self-organize to process different inputs in different subspaces of activation space (see Elhage et al.) 2022 for related demonstration). Nevertheless COMET magnifies these effects and our experiments show it improves performance over standard networks. Through extensive experiments various tasks, including image classification, language modeling, and regression, we demonstrate that COMET achieves improved performance. especially for larger models, demonstrating that method is applicable to many popular model architectures.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
This research was supported by the Intramural Research Program of the National Institute of Mental Health (NIMH), under project ZIC-MH002968 (PI: Dr. Francisco Pereira), and grant 2020906 from the National Science Foundation (PI: Matt Jones). We would also like to thank Michael Mozer for his invaluable insights and guidance.
REFERENCES
Rahaf Aljundi, Klaas Kelchtermans, and Tinne Tuytelaars. Task-free continual learning, 2019, URL https: //apxiv wttp:/v...rrr.b//..eeese 03596
Josh Andle, Ali Payani, and Salimeh Yasaei-Sekeh. Investigating the impact of weight sharing decisions on knowledge transfer in continual learning, 2023. URL httts://arxiv.orrs
abs/231 09506
Jimmy Ba and Brendan Frey. Adaptive dropout for training deep neural networks. In
C.J. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K.Q. Weinberger (eds.), Ad- vances in Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013. URL. /ttpss:/pproceed/ ngs. www.ree.....m newrriisssss.nee.nee.see.......e file/ 95
pdff
Anna Bair, Hongxu Yin, Maying Shen, Pavlo Molchanov, and Jose Alvarez. Adaptive sharpness- aware pruning for robust sparse networks, 2024. URL httph: htttpp://arxxv..rgg/bb//////
14306
Emmanuel Bengio, Pierre-Lue Bacon, Joelle Pineau, and Doina Precup. Conditional computation in neural networks for faster models, 2016. URL URL
htt...///1111.... 06297
Yoshua Bengio. Deep learning of representations: Looking forward, 2013.
URL
https:
wwwrxiv..rr/abb//1131 0445
Nina Dekoninck Bruhin and Bryn Davies. Bioinspired random projections for robust, sparse classi- fication, 2022. URL
09222
Tianlong Chen, Zhenyu Zhang, Ajay Jaiswal, Shiwei Liu, and Zhangyang Wang. Sparse moe as the new dropout: Scaling dense and self-slimmable transformers, 2023. URL https://arxxv.
org/abs/ 2303.
6160
Zhourong Chen, Yang Li, Samy Bengio, and Si Si. You look twice: Gaternet for dynamic filter selection in cnns, 2019. URLhttps://arre wwwp//org/aaaaeee wwt./orgr/aba.......// 1120
Brian Cheung, Alexander Terekhov, Yubei Chen, Pulkit Agrawal, and Bruno Olshausen. Superpo- sition of many models into one. Advances in neural information processing systems, 32, 2019.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Lev- skaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Samkaranarayana Pillai, Marie Pellat, Aitor Lewkowyez, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang. Bren- nan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellsterm, Douglas Eck, Jeff Dean, Slav Petrox, and Noah Fiedel. Palm: Scaling language modeling with pathways, 2022. URL //axxiv.////x....///a//ia wrg abs///a///// 22004 02311
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image worth 16x16 words: Transformers for image recognition at scale, 2021. URL htt.s: http///axvvvv//// www.wiv..rggggaa waw.rr//baa/a// /20aaaaaa21100 11929

---

Published as a conference paper at ICLR 2025
Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfield-Dodke Robert Lasenby, Dawn Drain, Carol Chen, et al. Toy models of superposi- tion. arXiv preprint arXi::2209. 10652, 2022.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022. URL
03961
Chrisantha Fernando, Dylan Banarse, Charles Blundell, Yori Zwols, David Ha, Andrei A. Rusu, Alexander Pritzel, and Daan Wierstra. Pathnet: Evolution channels gradient descent in super neural networks, 2017.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks, 2019. URL
Robert M. French. Using semi-distributed representations to overcome catastrophic forgetting in connectionist networks, 1993. URL htttss:/cch..aaa http:/aww....r...pppaaa///////pa
Suchin Gururangan, Mike Lewis, Ari Holtzman, Noah A. Smith, and Luke Zettlemoyer. DEMix lay- ers: Disentangling domains for modular language modeling. In Marine Carpuat, Cathe-Ca de Marneffe, and Ivan Vladimir Meza Ruiz (eds.), Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 5557-5576, Seattle, United States, July 2022. Association for Computational org/
Linguistics. doi: /2266///w//222.aaaeee 2022. latacl-mainn 407
URL
Song Han, Huizi Mao, and William J. Dally. Deep compression: Compressing deep neural networks with pruming, trained quantization and huffman coding, 2016. URL httptt:/axxi.v..oo
abs/1510.000 00149
Yizeng Han, Gao Huang, Shiji Song, Le Yang, Honghui Wang, and Yulin Wang. Dynamic neural networks: A survey, 2021. URL htttts::///aiiiiie org/ www.aaas/... 04 906
Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks, 2021. URL
https://arxix: org/ abs /tw.////22222 2002.1112222 00554
HuggingFace. Training a causal language model from scratch. 2022. Accessed: [2024].
co/
Steven C. Y. Hung. Cheng-Hao Tu, Cheng-En Wu, Chien-Hung Chen, Yi-Ming Chan, and Chu- Song Chen. Compacting, picking and growing for unforgetting continual learning, 2019, URL
https://arxiv. wttp:/iwwvv.///////////ooo
06562
Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey Hinton. Adaptive mixtures of local experts. Neural Computation, 3(1):79-87 1991. doi: 1tt1////cco..99.......
Tom Jacobs and Rebekka Burkholz. Mask in the mirror: Implicit sparsification, 2024. URL hotps:
www.xi...rg/as////.pee 09966
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- cralization in neural networks. Advances in neural information processing systems, 31, 2018.
Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gi- anna Lengyel, Guillaume Bour, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, Marie- Anne Lachaux, Pierre Stock, Sandcep Subramanian, Sophia Yang. Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée I Lacroix, and William El Sayed. Mixtral of experts, 2024. URL https: / wrw.i.vvvvmmm.
htt///rxvv....gg/e. http:/.rg/abaaeeeaeeee abtp//1ss/2211... /401 04088
Matt Jones, Peter Chang, and Kevin Murphy. Bayesian online natural gradient (bong), 2024, URL https://arxxtt http://awwwvv.ora///////o/a 19681
12

---

Published as a conference paper at ICLR 2025
M.I. Jordan and R.A. Jacobs. Hierarchical mixtures of experts and the algorithm. In Proceedings of 1993 International Conference on Neural Networks (IJCNN-93-Na) -a-Naya Japan), volume 2, pp. 1339-1344 vol.2, 1993. doi: 10.1 10.110/ICCNNN1993..6677
Haeyong Kang, Jaehong Yoon, Sung Ju Hwang, and Chang D. Yoo. Continual learning: Forget- free winning subnetworks for video representations, 2024. URLhttps://azii..oo wrg.iv..rrraaaa abs/
2312. 973
Rohit Keshani, Richa Singh, and Mayank Vatsa. Guided dropout, 2018. URL http:://paaai
org/abs/181 1812. 03965
Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg, Wan- Yen Lo, Piotr Dollár, and Ross Girshick. Segment anything, 2023. URL URL htttttt///axiiiiii wwtpabs..or..os///oeeooo....ss 2304 02643
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical Re-
port, University of Toronto, 2009. URL https://www cs wor.ott.eedouuuttttt..... "kriz/
fearning--eatures--22 2009-
Ya Le and Xuan S. Yang. Tiny imagenet visual recognition challenge. In CS 23/N, volume 7, pp-
3, 2015.
Yann LeCun, John Denker, and Sara Solla. Optimal brain damage. In D. Touretzky (ed.), Advances in Neural Information Processing Systems, volume 2. Morgan-Kaufman 1989. URL htttp:://proceeeeeettt ngs. wtw.rr/ cc cc/paper fil files/ #1les/paper/1969// file file///982baaccr1000020001111112 0t177776668-PPpe.Ppe
Bo Li, Yifei Shen, Jingkang Yang. Yezhen Wang, Jiawei Ren, Tong Che, Jun Zhang, and Ziwei Liu. Sparse mixture-of-expprrs are domain generalizable learners, 2023. URL
org/abs/2206. 04046
Zehui Lin, Liwei Wu, Mingxuan Wang, and Lei Li. Learning language specific sub-network for multilingual machine translation. In Chengqing Zong, Fei Xia, Wenjie Li, and Roberto Navigli (eds.), Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 293-305, Online, August 2021. Association for Computational Linguistics. doi: 10. 06653/v1/202// w/w.//.aclllooggg....... long............... URL BtsLtttss//aaaatthll..ooo org 2021 acl-1lacllloom ac1
Pratyush Maini, Michael C. Mozer, Hanie Sedghi. Zachary C. Lipton, J. Zico Kolter, and Chiyuan Zhang. Can neural network memorization be localized?, 2023. /tthhtss://axxi
abs/ abs/2307 2307 09542
Arun Mallya and Svetlana Lazebnik. Packnet: Adding multiple tasks to a single network by iterative pruning, 2018. URL BtLLtttss////xiv.coom org/ http://b//111...10
05769 /Ft
Arun Mallya, Dillon Davis, and Svetlana Lazebnik. Piggyback: Adapting a single network to mul- tiple tasks by learning to mask weights, 2018. URL https://aara wrg./bgg//aaa.a.
06519
Nicolas Y. Masse, Gregory D. Grant, and David J. Freedman. Alleviating catastrophic forgetting using context-dependent gating and synaptic stabilization. Proceedings the National Academy of Sciences, 115(44), October 2018 ISSN 1091-6490. doi: 10. 10.1073/ppas.11 http7///naa.1008839009 1803839115 15. URL
http::/dxx doi. www.oooorgr1...0 1013/ Tat.aaa..1183999 1803839115
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016. httpattts://arxt. org abs 100...1609000000 07843)
Sarthak Mittal, Yoshua Bengio, and Guillaume Lajoie. modular architecture enough?, 2022.
URL wttpthttt////iiiiie http:/orrgassss...e.ee 2206. 22711022663233333 02713
Hesham Mostafa and Xin Wang. Parameter efficient training of deep convolutional neural networks by dynamic sparse reparameterization, 2019. URL https://arxiv.org/
05967
13

---

Published as a conference paper at ICLR 2025
Mohammed Muqeeth, Haokun Liu, and Colin Raffel. Models with conditional computation learn suboptimal solutions, 2022. URL
iccinn20222odelsspp
Yuval Netzer, Tao Wang. Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Read- ing digits in natural images with unsupervised feature learning. In NIPS Workshop on Deep Learning and Unsupervised Feature Learning, 2011. URL http::/ufldl..../ stanfoord
edu. http://usenuubbrr//iss//0000 http:///o....ooooooooooouuorrrrr bousenumbers.pde
Andrew Y Ng. Feature selection, 1 I vs. 12 regularization, and rotational invariance. In Proceedings of the twenty-first international conference on Machine learning, pp. 78, 2004.
OpenAl. Chatgpt: Optimizing language models for dialogue, Jan 2023a.
URL https://
openai wtt..lccccocooo..o ccom/bloog/chattgt.//
OpenAl. Gpt-4 technical report, 2023b.
Mansheej Paul, Feng Chen, Brett W. Larsen, Jonathan Frankle, Surya Ganguli, and Gintare Karolina Dziugaite. Unmasking the lottery ticket hypothesis: What's encoded a winning ticket's mask? In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL UtLL:tttp:///peeeeviiee Eorum? -dt.indddxxSsaaddd uk2
Lorenzo Pes, Rick Luiken, Federico Corradi, and Charlotte Frenkel. Active continual learning in time-to-first-ppike neural networks, 2024. URL /tLL/sst//////ri..orr
abs/2404 19419 9
Jonas Pfeiffer, Sebastian Ruder, Ivan Vulić, and Edoardo Maria Ponti. Modular deep learning, 2024. URL
org/ abs////s//222/// 2302 11529
Nasim Rahaman, Muhammad Walced Gondal, Shruti Joshi, Peter Gehler, Yoshua Bengio, Francesco Locatello, and Bernhard Schölkopf. Dynamic inference with neural interpreters, 2021. URL
https: //axivv./ org/abs/ 22222110111 06399
Hadsell R;Rao D;Rusu AA;Pasannu Raia. Embracing change: Continual learning in deep neural networks, 2020. URL, httttttt///t/mmmtt ncbi. nlm. nih. gov/33158755)
Carl Rasmussen and Zoubin Ghahramani. Infinite mixtures of gaussian process experts. In T. Di- etterich, S. Becker, and Z. Ghahramani (eds.), Advances in Neural Information Processing Sys- tems, volume 14. MIT Press, 2001. URL neur.pss...e.pppp.pp... cc/ppper pdf
Carl Edward Rasmussen and Christopher K. I. Williams. Processss Machine Learning. MIT Press, 2006. URL /tt.attte//gguuuut..prrrrr.....e
pd
Stephen Roller, Sainbayar Sukhbaatar, A Arthur Szlam, and Jason Weston. Hash layers for large sparse models, 2021. URL
04426
Clemens Rosenbaum, Tim Klinger, and Matthew Riemer. Routing networks: Adaptive selection of non-linear functions for multi-task learning, 2017. URL
01239
Clemens Rosenbaum, Ignacio Cases, Matthew Riemer, and Tim Klinger. Routing networks and the challenges of modular and compositional computation, 2019. URL https://arxiv..rrg/
abs/1904.122 12774
Noam Shazeer, *Azalia Mirhoseini, *Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hin- ton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of- experts layer. In International Conference on Learning Representations, 2017. URL https: //openeeveww wtt.//eetu//////
ckMDg1g
Noam Shazeer, Kayvon Fatahalian, William R. Mark, and Ravi Teja Mullapudi. Hydranets: Special- ized dynamic architectures for efficient inference. In 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8080-8089, 2018. doi: 10. 1090/CPPP W0w.V....01..00 20.8.00843.
14

---

risioontrannsformrrr-....a Published as a conference paper at ICLR 2025
Kurt Shuster, Jing Xu, Mojtaba Komeili, Da Ju, Eric Michael Smith, Stephen Roller, Megan Ung, Moya Chen, Kushal Arora, Joshua Lane, Morteza Behrooz, William Ngan, Spencer Poff. Naman Goyal, Arthur Szlam, Y-Lan Boureau, Melanie Kambadur, and Jason Weston. Blenderbot 3: deployed conversational agent that continually learns to responsibly engage, 2022. URL https: //arxiw..rgg/ab//2220 03188
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting. Journal of Machine Learning Research, 15(56):1929-199) 1958, 2014. URL http: http://jjalr...r/apppeeeeeee pttpppppspppssss///
srivastaval4a. html
Robert Tibshirani.
Regression shrinkage and selection via the lasso, Journal of the Royal Statistical Society Series B: Statistical Methodology, 58(1):267-288) 1996.
Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Un- terthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, and Alexey Dosovitskiy. Mlp-mixer: An all-mip architecture for vision, 2021. URL https:
htspgg/bb/220000 01601
L. Tumstall, L. von Werra, and T. Wolf. Natural Language Processing with Transformers: Building Language Applications with Hugging Face. O'Reilly Media, 2022. ISBN 9781098103248. URL
https://books: google ccm/books? ? CuezaAACAA
Zhenyi Wang. Li Shen, Le Fang, Qiuling Suo, Tichang Duan, and Mingchen Gao. Improving
fask-free continual learning by distributionally robust memory evolution, 2022. URL https:
Mitchell Wortsman, Vivek Ramanujan, Rosanne Liu, Aniruddha Kembhavi, Mohammad Rastegari, Jason Yosinski, and Ali Farhadi. Supermasks in superposition, 2020. URL URL https://axxiv::
org/ /t///bb//22 2006. 14769
Haoyun Xu, Runzhe Zhan, Derek F. Wong, and Lidia S. Chao. Let's focus neuron: Neuron-level supervised fine-tuning for large language model, 2024. URL htttss:/
/bs/
2403. 11621
Li Yang, Zhezhi He, Junshan Zhang, and Deliang Fan. Ksm: Fast multiple task adaption via kemel- wise soft mask learning, 2020. URL wttttsss//aaxi..ttttt org /ttp:/bbb/00.../// 05668
Fei Ye and Adrian G. Bors. Task-free continual learning via online discrepancy distance learning. 2022. URL /ORl://pps//aaaxvi.o////////2/22
Murat Onur Yildirim, Elif Ceren Gok Yildirim, Ghada Sokar, Decebal Constantin Mocanu, and Joaquin Vansehoren. Continual learning with dynamic sparse training: Exploring algorithms for effective model updates, 2023. URL
htg/abbs23008 /200/ 14831
Kentaro Yoshioka. and related models on cifar-10. 2024.
vision-tansffomerr-cffff......
Training vision transformers (vit)
Hattie Zhou, Janice Lan, Rosanne Liu, and Jason Yosinski. Deconstructing lottery tickets: Zeros, signs, and the supermask. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bue, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32, Cur- ran Associates, Inc., 2019. URL UR//appss//ppoeeedi.......o cc/pppprrrcpppppprpppppp files/
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew Dai. Zhifeng Chen. Quoe Le, and James Laudon. ofxture-of-experts with expert choice routing, 2022. URL https: //arxivv wttp://gppap.........o 2202.09366
org/abs/////aooooooooo...//.ooar abs
15

---

Published as a conference paper at ICLR 2025
A EXPERIMENT DETAILS
In this section, provide detailed description of the experiments conducted to evaluate the perfor- mance of COMET. The following subsections outline the experimental setup, including the datasets used, model architectures, and hyperparameters. We also present additional experimental results and analyses.
A.1 MAIN RESULTS: STANDARD MLP CIFARIO
In this subsection, we present the main results of our experiments on the CIFARIO dataset using standard l-yer MLP architecture. We first describe the model configuration and training setup, followed by a description of the hyperparameter settings used to vary model capacity and sparsity levels.
Model Configuration and Training We employ a standard 4-yyeer MLP architecture, utilizing the SGD optimizer with learning rate of le-4. To ensure robustness, we train each model over random seeds for 100 epochs. We systematically explore the effects varying model capacity and sparsity levels by modifying the number of neurons in each layer and the sparsity ratio.
Model Capacity Variations We consider four different model capacities by setting the number of neurons in each layer to 100, 500, 1000, 3000.
Sparsity Level Variations For each model capacity, investigate the impact of three different sparsity levels: 0.1, 0.5, and 0.9.
A.2 ADDITIONAL RESULTS: STANDARD MLP CIFARIO
This subsection presents additional experimental results on the CIFARIO dataset using a standard 4- layer MLP architecture, but with higher learning rate of Ie-3. We describe the model configuration and hyperparameter settings used, and report the results of varying model capacity and sparsity levels.
Model Configuration and Training To further assess the robustness of the COMET method, we also conduct experiments using higher learning rate of 1e-3. In these experiments, we again employ a standard 4-yerer MIP,, SGD optimizer, and systematically vary model capacity and sparsity levels by adjusting the number of neurons in each layer and the sparsity ratio, respectively.
We consider an additional model capacity to a total of five different model capacities by setting the number of neurons in each layer to 100, 500, 1000, 3000, or 9000.
Model Capacity Variations
Sparsity Level Variations For each model capacity, investigate the impact of three different sparsity levels: 0.1, 0.5, and 0.9.
As illustrated in fig. 9 consistent trend emerges as systematically vary the number of rons and sparsity levels. Moving from top left to bottom right, we observe shift in the optimal model configuration. Initially, when network capacity is limited, the standard model outperforms the COMET model. However, as network capacity increases, the COMET becomes the top performer, surpassing every other model. This trend reinforces our key finding: selective neuron activation be- comes increasingly beneficial as network capacity increases, enabling faster learning and improved generalization through enhanced forward transfer.
A.3 CONTEMPORARY ARCHITECTURES
This subsection presents experimental results contemporary architectures, including ViT and MLP-Mixer models. We deseribe the hyperparameter settings used for each architecture, and detail the modifications made analyze the effect of our sparsity method, COMET.
16

---

Published as a conference paper at ICLR 2025
Neuronsw188 100
Neurons-500
Neurons= 1000
Neuro1s-5000
Nourons-5000
0:33
0.00
0.00
Welcerom
0.45
235
0.00
COMET
bamoull dropout
pouting
dropout.mode
Figure 9: Illustration of 4-yerrr MLP networks trained on CIFARIO with higher learning rate (le- 3), showcasing the impact of varying network capacity and sparsity levels. As we systematically increase the number of neurons and decrease sparsity (moving from top left to bottom right), we observe a shift in the best-perforning model. Initially, the standard model (without sparsity, blue) outperforms the COMET when network capacity is low. However, network capacity grows, the COMET model emerges as the top performer.
Hyperparameter Settings We built upon the tuned hyperparameters from Yoshioka (2024) and made the following modifications to analyze the effect of our sparsity method, COMET. Specifi- cally, we systematically increased model capacity by adding more neurons to the MLP layers of each model. Additionally, we used the tanh activation function on the backbone MLP layers. See appendix A.7 for further analysis activation functions.
Vision Transformer Hyperparameters The standard ViT uses set of hyperparameters, which modified follows. The number of classes is set to 10 for CIFARI0 and SVHN, 100 for CIFARI00, and 200 for Tiny Imagenet. The model depth is 6, with attention heads. We increased the MLP dimension from 512 3072 for the ViT medium model and to 6144 for the ViT large model. The dropout rate is 0.1, and the patch size is The embedding dropout rate also 0.1. We trained the models for 200 epochs with learning of le-4.
MLP-Mixer Hyperparameters The MLP-Mixer model uses a different set of hyperparameters. The patch size is 4. We increased the dimension from 512 3072 for the MLP-Mixer medium model. The model depth is 6, and the number of classes is set to 10 for CIFARIO and SVHN, 100 for CIFARIO0 100, and 200 for Tiny Imagenet. We trained the models for 500 epochs with a learning rate of 1e-3.
Training Setup All models were trained using the Adam optimizer with cosine learning rate schedule on single A100 GPU. To ensure robustness, we train each model over 3 random seeds.
A.4 IMAGE CLASSIFICATION
We provide additional experimental results on the image classification task, evaluated on the CIFAR- 10 and SVHN datasets, illustrated in Figure 11 and Figure 10 respectively.
17

---

Published as a conference paper at ICLR 2025
SVHN
CIFARID
60
60.
60
Valiidation Accuracy
60-
24.41.24
202
23
200
202
Validation
200
203
20
202
220
100
250
80
60-
60
60
20-
200
25C
Epoch
Epoch
Figure 10: VITs and MiP-Mixers on SVHN.
Figure 11: VITs and MLP-Mixers on CIFARIO.
[Highest accuracy] (# trainable param.)
[Highest accuracy] (# trainable param.)
A.5 LANGUAGE MODELING
We extend our evaluation of COMET to the task of language modeling, examining its performance on various GPT model variants.
A.5.1 MAIN RESULTS
This subsection presents our main results on the language modeling task, detailing the performance of COMET on three different GPT model variants. We describe the hyperparameter settings and training setup used for each model, and report the results of our experiments.
GPT Model Variants We train three variants of the GPT model, each with different set of parameters. The standard GPT-2 model has 12 layers, 768 hidden units, attention heads, and 117M parameters. The GPT-2-Medium model has 24 layers, 1024 hidden units, 16 attention heads, and 345M parameters. The GPT-2-Large model has 36 layers, 128 hidden units, 20 attention heads, and 774M parameters.
Hyperparameter Settings We built upon the tuned hyperparameters from HuggingFace (2022) Our optimizer of choice was Ad with a learning rate of 5e-4, weight decay 0.1, and 1. 1,000 warmup steps. We also used gradient accumulation with 8 steps, which resulted in effective batch size of 256, calculated by multiplying the per-device train batch size (32) by the gradient accumula- tion steps (8). We used tanh activation function the backbone MLP layers and a cosine learning rate schedule with warmup. We also enabled mixed precision training to accelcrate computations.
Training Settings Each model was trained from scratch single A A100 GPU. Due computa- tional constraints and time limitations, we restricted training to either 3 epochs or maximum of 24 hours. To ensure robustness, we train each model over 3 random seeds.
A.5.2 ADDITIONAL RESULTS
We conducted additional experiments investigate how the choice of hyperparameters influences the performance of our method, COMET, when applied to the MLP layers of GPT models. This analysis aims provide deeper understanding of the robustness and adaptability of COMET under various hyperparameter settings.
The following Figures mark COMET models with spec true, due the fact that COMET's input- dependent gating mechanism leads to the formation of experts that are selectively specialized for specific inputs. The standard models are marked with sale_false.
Tokenizer Effect We note that. in general, the GPT models evaluated tend learn faster when using the tokenizer provided by HugggingFaace (2022). However, due to its widespread adoption, we opt the standard GPT-2 tokenizer for the remainder of our experiments:
18

---

Published as a conference paper at ICLR 2025
tokenizer AutoTokenizer from-pretrained (" ("gpt2")
To validate our main findings, we first assess the performance of the different GPT-2 model sizes on WikiText and CodeParrot using the standard GPT-2 tokenizer and 50% sparsity level. Our results are presented figs. to
Consistent with our main results, we find that the COMET-based model learns faster, even with a smaller model size, when using the standard GPT-2 tokenizer. This suggests that the observed pattern is robust and not specific to the tokenizer used.
valication_perplexity
spec, True_got2 .vikitest spec_ False_gott_______e
200
150
50
Step
10k
20k
Figure 12: Validation perplexity of GPT-2 on Wikitext dataset using the standard GPT-2 tokenizer.
Deliiation_perleest spec. Falso, gpt2
spec, 2pt2
200
100
50
Step
2k
4k
6k
8k
10k
12k
Figure 13: Validation perplexity of GPT-2 Medium Wikitext dataset using the standard GPT-2 tokenizer.
Learning Rate Effect To investigate the impact learning rate on COMET method, we scale the learning rate by factor of 10, from le-4 to le-3, The results of this experiment are presented in
figs. 18 to 23
Our findings show that. across all experiments, the COMET models consistently learn faster than their standard counterparts. However, at the 50% sparsity level, we observe that the smaller models often perform slightly worse than the standard models. This result is consistent with previous findings, which suggest that adding sparsity to models with limited capacity can negatively impact performance.
Furthermore, we identify an interesting trend as the model size increases, as seen in figs. 19 and 23 Specifically, when using larger learning rate and larger model sizes, the standard models
20
19

---

Published as a conference paper at ICLR 2025
validation perplexity
wikitest spec www.eegpt--largee.iieeeeeee wikktest
200
150
100
Steg
2k
4k
6k
Bk
Figure 14: Validation perplexity of GPT-2 Large on Wikitext dataset using the standard GPT-2 tokenizer.
pelppetion_perplexity
spec, True, Tttpeeppt_____eeeeeeeeeeeeeee codeparrot -spec_ got2 codeparrot
50
40
30
20
Step
500
1k 1.5k 2k 2.5k
3k
Figure 15: Validation perplexity of GPT-2 on CodeParrot dataset using the standard GPT-2 tokenizer.
tend overfit experience exploding gradients, resulting significant increase in validation perplexity. In contrast, adding sparsity using COMET not only enables faster learning but also mitigates overfitting and gradient explosion, allowing for stable training with a larger learning rate.
Batch Size Effect To further investigate the robustness of the COMET method, we examine the effect of reducing the batch size by factor of 4, achieved by decreasing the gradient accumulation step from to Our results are presented in figs. 24 to 29
Our findings demonstrate that COMET remains effective even with reduced batch size, making a viable option for users with limited computational resources. We observe that on the Wikitext dataset, the smaller GPT model with COMET learns faster and achieves comparable performance to the standard model at the end of training. In contrast, on the CodeParrot dataset, the smaller model with COMET outperforms the standard model. Moreover, to smaller batch sizes, where standard models may struggle with overfitting or exploding gradients. By incorporating sparsity, COMET enables more stable training and better performance, even in resource-constrained environments.
note that COMET's benefits extend
Mixed Precision Effect
We further investigate the robustness of the COMET method by evaluat- ing its performance under mixed precision training. Specifically, assess the impact of switching from FP16 to FP32 precision on the COMET-based models. Our results are presented in figs. 30
32
20

---

Published as a conference paper at ICLR 2025
validation perplexity
spec_ True peue_ppte-meeimmmmo codepprrot spec_False wtt.:aaalltpeeeeeeeeeeee stteeeedlmmocccorreeeoooo
40
20
Step
500
1k
1.5k
2k
2.5k
3k
Figure 16: Validation perplexity of GPT-2 Medium on CodeParrot dataset using the standard GPT-2 tokenizer.
validation validation_perplexityy
sple_False_eeee hat___alee__eeeeee httpe__gtt-largecodeeeeee codeparrot www.:r/oppptloreeeeeeeeeeeo
35
25
20
10
Step
Figure 17: Validation perplexity of GPT-2 Large on CodeParrot dataset using the standard GPT-2 tokenizer.
In summary, our experiments demonstrate that the COMET method consistently enables faster learn- ing across various model sizes, datasets, and training settings. By incorporating COMET, ob- serve improved performance and robustness, even when switching from FP16 to FP32 precision. This is particularly valuable, as FP32 precision often preferred in certain applications where nu- merical stability is crucial. Moreover, having models that can effectively learn in both FP16 and FP32 precision regimes provides greater flexibility and adaptability, allowing for more efficient de- ployment on wide range of hardware platforms.
A.6 REGRESSION
This subsection presents our evaluation of COMET on a regression task using the SARCOS dataset. We deseribe the experimental setup, including the dataset, model architecture, and hyperparameter settings, and report the results of our experiments.
Dataset To conclude evaluation, apply the COMET method to regression task using the SARCOS dataset. This dataset is derived from an inverse dynamics problem involving a 7- joint anthropomorphic robot arm, where the goal is to predict the 7 joint torques based on 21- dimensional input space consisting of joint positions, velocities, and accelerations. We focus on single output dimension, following the approach of Rasmussen & Williams 2006). The dataset
21

---

Published as a conference paper at ICLR 2025
validation perplexity
eelidationppeppexitty
spec_ False_sott__aaaeea_______ wikitest spec Troe_gpt2_wkitee
80
60
40
Step
5k
10k
15k
20k
Figure 18: Validation perplexity of GPT-2 on Wikitext dataset using the standard GPT-2 tokenizer and learning rate of le-3.
validation_ perpllixttepeextttyy
spec False_ medium wikitest spec. Truc, gpt2 medium wikitext medlum_wwiktertt
200
100
5k
15k
20k
Figure 19:
Validation perplexity of GPT-2 Medium on Wikitext dataset using the standard GPT-2 tokenizer and learning rate of 1e-3.
Building on the work of Jones et al. (2024). we employ a 4-layer MLP network, but experiment with varying the number neurons each layer and the level of sparsity.
publicly available at https: /ttppss/gguusiaappoocsssss http://siaapppooessooooo///////a http://.//aaggpmpaaaa..m
We use SGD optimizer with a learning rate of le-2. To ensure robustness, we train each model over random seeds for 50 epochs. We systematically explore the effects of varying model capacity and sparsity levels by modifying the number of neurons in each layer and the sparsity ratio.
Model Capacity Variations We consider five different model capacities by setting the number of neurons each layer to 100, 500, 1000.
Sparsity Level Variations For each model capacity, we investigate the impact of three different sparsity levels: 0.1, 0.5, and 0.9.
The results, presented in fig. 33 demonstrate the effectiveness of our approach in this domain. Consistent with our previous findings, again observe that the COMET model outper- forms the standard model as network capacity increases, confirming the benefits selective neuron activation in regression tasks.
Results
22

---

Published as a conference paper at ICLR 2025
validation, perplexity
swwppppelsseeesseeeeeepp wwwpsee_ppttlorrgeeiiiiiiiee wikitest
spec. True True_ppt2-large _wikitet
2500
2000
1500
1000
500
Step
5k
10k
15k
20k
Figure 20: Validation perplexity of GPT-2 Large on Wikitext dataset using the standard GPT-2 tokenizer and learning rate of le-3.
validation_perplexiii
False, 8pt2 codeparrot spec. True, got2 codeparrot
Figure 21: Validation perplexity GPT-2 on CodeParrot dataset using the standard GPT-2 tokenizer and learning rate of le-3.
A.7 ACTIVATIONS ANALYSIS
In this section, we analyze the impact of different activation functions the performance of COMET-based models. We present an experimental evaluation of various activation functions on 4-layer MLP network trained on the CIFARIO dataset, and discuss the results.
We now investigate the impact of utilizing various activation functions within the backbone network on the performance of a COMET- based model during training.
Evaluating the Impact of Activation Functions on COMET
Experimental Setup We evaluate the effect of different activation functions on 4-layer MLP network trained on the CIFARIO dataset with the following settings: SGD + Momentum optimizer, a learning rate of 1e-3, sparsity level of 0.5, and 3000 neurons in each layer.
Results Our results presented in fig. We find that the COMET model outperforms the standard model when most activation functions are applied on the backbone network. However, we observe that COMET does not perform well the standard model for certain non-monotonic activation functions, specifically GELU, Mish, and SiLU, under the settings we chose (sparsity level 0.5 and 3000 neurons in each layer). Due time constraints, do not delve deeper into the
23

---

Published as a conference paper at ICLR 2025
validation, valilxtion_peppexiityy
spec, www__ppt-mmeiummmomm
12
10
8
Step
1k
2k
3k
4k
5k
6k
Figure 22: Validation perplexity of GPT-2 Medium on CodeParrot dataset using the standard GPT-2 tokenizer and learning rate of le-3.
validation validation_perplextty
codepparot spec_ Twwe__gpttllrgeeeeemmm coodeparot
60
40
Step
3k
Validation perplexity of GPT-2 Large CodeParrot dataset using the standard GPT-2 tokenizer and a learning rate of 1e-3.
Figure 23:
reasons behind this phenomenon but suggest that future work could investigate why non-monotonic functions may hinder the effectiveness of COMET.
A.8 COMET: ROUTING NETWORK ARCHITECTURE
We have implemented COMET using a routing network having the same architecture as the back- bone MLP network, but this is not necessary. COMET only requires the routing network generate mask at each layer having the same shape (or a shape that can be broadcast to) the shape of the corresponding backbone layer. Future work will investigate alternative routing network architectures to determine if they can improve the performance and efficiency of COMET-based models.
A.9 RUNNING TIME AND MEMORY REQUIREMENTS
While COMET does not introduce additional trainable parameters, does require extra putation to calculate the masks. We evaluate the training times and GPU memory usage of COMET compared to standard vision models.
Tables and 2 show that while COMET introduces some additional computation, the overhead is moderate. For example, the added GPU usage ranges from 0.3 GB (4%) to 7 GB (54%), where the
24

---

Published as a conference paper at ICLR 2025
Validation Perplexity
spec_ Ralse_got2, wikitest
spec 2ot__pp22_kkieee
400
300
200
100
Step
20k
40k
60k
80k
100k
Figure 24: Validation perplexity of GPT-2 on Wikitext dataset using the standard GPT-2 tokenizer and reduced batch size of 64 ( of original).
Validation Perplexity
spec. got2-medium spec. True, gpt2--pddlun
300
200
100
Step
10k
20k
30k
Figure 25: Validation perplexity of GPT-2 Medium on Wikitext dataset using the standard GPT-2 tokenizer and reduced batch size of 64 & of original).
larger percentage differences occur in cases where the ViT network excessively overparameterizze (e.g., when the number of neurons in the original network is multiplied by 12, from 512 to 6144).
worth noting that for simplification, our implementation has not been optimized:
1. We did not use sparse matrix multiplication methods, which could avoid unnecessary com- putations through the zero masks. This would reduce both the training time (by preventing gradient calculation of the zero gates) and overall GPU usage and inference time (due to less computation). This optimization could lead to significant improvements, potentially even compared the base models.
2. As mentioned in appendix A.8 the same dimensions as the weight matrix W; at each layer. This means that each random projection matrix has the same shape as the corresponding FC layer's matrix. Notably, however, the only requirement is that they share the same input and output dimensions, leaving room for future exploration of alternative configurations.
focus specific case where the fixed matrix C, has
Overall, these results suggest that COMET's additional computational overhead is manageable, with our simple implementation, and that its benefits can be achieved without significantly sacrific- ing efficiency.
25

---

Published as a conference paper at ICLR 2025
Validation Perplexity
Fepc__eeleeppptteeeee wikieett
spec_ wow__spt-llrgg_eviee wwikieet
500
300
200
100
Step
10k
20k
30k
Figure 26: Validation perplexity of GPT-2 Large on Wikitext dataset using the standard GPT-2 tokenizer and reduced batch size of 64 (1 of original).
Validation Perplexity
spec, True wttpeppptteaeeeeee codeparrot spec. False, gots codeparrot
25
1k
2k
3k
4k
5k
6k 7k
Figure 27: Validation perplexity GPT-2 on CodeParrot dataset using the standard GPT-2 tokenizer and reduced batch size of 64 ( of original).
A. TRANSFER LEARNING
To further demonstrate COMET's ability to transfer learned knowledge, we conducted an additional experiment. We trained MLP COMET and standard MLP models CIFAR-10 and then fine-tuned them on SVHN. The results, presented in fig. 33] show that COMET exhibits superior transfer learn- ing capabilities. Specifically, "COMET with transfer" outperforms "standard model with transfer"
Dataset / Model
SVHN
CIFARIO
CIFAR100
Tiny ImageNet
COMET ViT
6.6
Standard VIT
6.1
6.1
6.1
6.1
COMET ViT-Med
12.3
12.3
12.3
12.3
Standard ViT-Med
9.2
9.2
9.2
9.2
COMET ViT-Large
19.8
19.8
19.8
19.8
Standard ViT-Large
12.8
12.8
12.8
12.8
COMET MLP-Mixer
6.7
6.7
6.7
6.7
Standard MLP-Mixer
6.4
6.4
6.4
6.4
COMET MEP-Mixer-Med
25.5
25.5
25.5
25.5
Standard MLP-MMier-MMd
23.2
23.2
23.2
23.2
Table 1: GPU Utilization in GBs. All models were trained on single A100 GPU.
26

---

Published as a conference paper at ICLR 2025
Validation Perplexity
spec_ watsee_ggtt-meeiinn wwwermmdioeeeoooerreeeee codeparrot
spec mee__ppt--meiimmm.mmm
20
15
Step
1k 2k 3k
4k
5k
6k
7k
Figure 28: Validation perplexity of GPT-2 Medium on CodeParrot dataset using the standard GPT-2 tokenizer and reduced batch size of 64 (1 of original).
Validation Perplexity
Faise, large. codeparrot
True, large, codeparrot
22
20
16
Step
500
Figure 29: Validation perplexity of GPT-2 Large on CodeParrot dataset using the standard GPT-2 tokenizer reduced batch size of ( of original).
in both initial accuracy and learning speed. Moreover, "COMET with transfer" also surpasses the overall performance of both "COMET without transfer" and "standard model without transfer" on SVHN. This suggests that COMET's advantage is not limited to out-of-sample generalization in the same distribution but also extends to transfer across tasks.
Dataset / Model
SVHN
CIFARIO
CIFAR100
Tiny ImageNet
COMET ViT
0.8
1.2
Standard VIT
0.9
0.7
0.7
1.2
COMET ViT-Med
2.0
1.5
1.5
2.4
Standard ViT-Med
1.3
1.1
1.1
1.3
COMET ViT-Large
3.3
2.2
2.0
3.3
Standard ViT-Large
1.8
1.4
1.3
1.7
COMET MLP-Mixer
1.8
1.5
1.4
2.3
Standard MLP-Mixer
1.6
1.4
1.3
2.3
COMET MLP-Mixer-Med
10.2
5.4
9.3
Standard MLP-Mixer-Mo
7.6
4.5
8.1
Table 2: Training times in hours. All models were trained single A100 GPU.
27

---

Published as a conference paper at ICLR 2025
validation eelipxtionpprplexitt
spec__alsee false__ot2__aiee______ wikitest
spec True_gpt2_wkktee
120
80
40
Step
5k
10k
15k
20k
Figure 30: Validation perplexity of GPT-2 on Wikitext dataset using the standard GPT-2 tokenizer and FP32 precision.
validation_
spec, Tww._rrrrtta... (ro._ggt2-meellumm .vikitewt spec_ Fwc__aalle__pllleee malse_gpt2-medlum wikitext
300
250
200
100
Step
500
1k
1.5k
2k
Figure 31: Validation perplexity of GPT-2 Medium Wikitext dataset using the standard GPT-2
tokenizer and FP32 precision.
Interestingly, we observe larger gap performance between "standard model with transfer" and "standard model without transfer" compared to the gap between "COMET model with transfer" and "COMET model without transfer". This might suggest that the standard model benefits more from transfer learning non-iid data. However, since overall performance is generally taken as the primary metric for evaluating transfer learning, COMET's superior performance on SVHN is the most relevant outcome. This finding promising for future work that will more thoroughly investigate COMET in multi-task settings including transfer learning and continual learning.
28

---

Published as a conference paper at ICLR 2025
validation, velipation_perplexityy
spec False_ codepaarrtt spec www._spt2_ccoopeeeeeeeeee
15
500
1k
1.5k
Figure 32: Validation perplexity GPT-2 on CodeParrot dataset using the standard GPT-2 tokenizer and FP32 precision.
SARCOS
Neurons= =100
Neurons=500
Neurons 1000
350
Standard
300
Standard
300
Standard
(0.02M)
10.51M]
250
(2.03M)
250
0.7
250
200
150
100
100
100
50-
50
10
20
30
50
0
10
40
50
10
30
40
50
120
140
(0. (0.02M))
19.51M)
120
12.03M)
COMFT
Validation Loss
0.5
(0.02M)
80-
10.51M]
100
(2.03M)
80
80-
60-
40.
20-
10
20
50
10
40
50
30
40
50
140
120-
(2M)
(0.51M)
120
(2.03M)
COMET (0.02M)
COMET
COMET
100
0.9
100
k
60
40-
10
20
0
10
50
40
Epochs
Figure 33: Illustration of 4-layers MLP trained on SARCOS.
29

---

Published as a conference paper at ICLR 2025
gelu
sigmoid
Velidation Accuracy
Volidation Accuracy
Velidation Accuracy
model
spectfiity
spocticily mode
spectficty model
spocticlyy modol
26
60
24
Epoch
Epoch
Epoch
Epach
clu
tanh
sclu
softplus
Validation Accuracy
Accuracy
Ualidattio Accuracy
Validation
standard prdeel speclisky (astel
Menccrd mode
exrdert model
19
donterd mode!
spec Faily, modee
spacll maxtel
spociiicty model
Epoch
Epoch
Epoch
Epoch
Figure 34: Comparison of the performance of COMET-trained MLP networks on the CIFAR1O dataset using different activation functions.
0.7
0.6
5VHN Validation Accuracy
0.5
0.4
COMET With Transter
0.2
COMET Without Transfer
Standard Modal Transter
Standard
5
10
15
20
25
30
Figure 35: Non-lID Transfer: Evaluating COMET and the standard model on non-iid transfer learn- ing. particular, we train both COMET and the standard model on CIFAR followed by finetuning on SVHN (with transfer), or just train both models on SVHN (without transfer).
30


Reviews:
[{"review_id": "TyBUL3aCZL", "reviewer": "Reviewer_fx4N", "summary": "The paper introduces Conditionally Overlapping Mixture of ExperTs (COMET).\nCOMET uses biologically inspired, fixed random projections to generate binary masks that define subnetworks know as 'experts'.\nThe mask generation process is input-dependent, causing similar inputs to activate overlapping sets of experts.\nThe authors test the models on a range of benchmark tasks, finding that COMET performs well, particularly for large model sizes.", "questions": "How does this model perform on a continual learning benchmark, such as permuted MNIST or split-CIFAR-100? \n\nWhat are the additional costs to using COMET in terms of training time and memory usage?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well written and the core idea is explained clearly. \n\nThe authors demonstrate key properties of the COMET model, notably showing that similar inputs tend to activate overlapping experts, facilitated by the fixed gating mechanism.\n\nThe model is tested on a wide selection of benchmark tasks including computer vision, language modelling, and regression. \n\nThe authors demonstrate the benefit of using COMET, particularly at large model sizes.\n\nThe use of COMET requires no additional trainable parameters which is quite advantageous.", "weaknesses": "In other works these gating functions can help alleviate catastrophic forgetting for tasks that are presented sequentially, but this is something that has not been tested in this paper. \n\nGiven that previous work has similarly employed networks to determine gating, I am not entirely convinced that the novelty here is sufficient. However, I acknowledge that unlike prior approaches, which relied on trainable gates, this method uses fixed random projections.\n\nThere will be additional computational costs to using COMET but there is not an in-depth analysis of these costs in the paper. An analysis of training/inference times and GPU memory usage between COMET and the standard models would strengthen the submission.", "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": "4x0VN0IaiS", "reviewer": "Reviewer_ETwf", "summary": "This paper introduces COMET (Conditionally Overlapping Mixture of ExperTs), a new method for creating sparse neural networks. The authors show using fixed, biologically-inspired routing can create more efficient and effective neural networks, particularly for larger models, while avoiding common problems in sparse architectures like representation collapse and poor knowledge transfer. The key insight is that COMET creates input-dependent sparsity without needing to learn the routing mechanism. COMET uses a fixed, biologically-inspired random projection combined with a k-winner-take-all operation to route inputs through the network, rather than using trainable gating functions.", "questions": "Is there any reason to believe this phenomenon does not already occur in large networks? \\cite{elhage2022toy} describe a situation where neural networks encode the phenomenon observed in \\cite{cheung2019superposition} during the course of training. Are there advantages of explicitly creating the superposition?\n\nFigure 4 shows a strange result where the smaller_model performs consistently as well as the standard_model even for fairly low p_k values. It's unclear to me why there would be a benefit for COMET if at the neurons=3000, the smaller_network at pk=.1 will perform as well as the standard_model.  What is being gained here?\n\n@article{elhage2022toy,\n  title={Toy models of superposition},\n  author={Elhage, Nelson and Hume, Tristan and Olsson, Catherine and Schiefer, Nicholas and Henighan, Tom and Kravec, Shauna and Hatfield-Dodds, Zac and Lasenby, Robert and Drain, Dawn and Chen, Carol and others},\n  journal={arXiv preprint arXiv:2209.10652},\n  year={2022}\n}", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Unlike other methods, the proposed COMET method has no trainable gating functions (unlike standard Mixture of Experts) and avoids representation collapse.\n\nDoes not require explicit input/task IDs or pre-defined expert specialization.\n\nWorks across multiple architectures (MLPs, ViTs, GPT, MLP-Mixers).\n\nThe work is particularly similar to \\cite{cheung2019superposition}, especially with the use of a random projection matrix V to handle the decision to mask. The justifications for using random projections in \\cite{cheung2019superposition} seem to align well with the described capacity benefits of the COMET method in larger networks as compared to smaller networks. In particular, with a larger number of neurons, the probability of interference between masks rapidly decreases.\n\n@article{cheung2019superposition,\n  title={Superposition of many models into one},\n  author={Cheung, Brian and Terekhov, Alexander and Chen, Yubei and Agrawal, Pulkit and Olshausen, Bruno},\n  journal={Advances in neural information processing systems},\n  volume={32},\n  year={2019}\n}", "weaknesses": "The arguments for the number of experts is based on the possible permutations of masks that can be created which gives an unrealistically large number of possible experts. But this does not account for interference issues and establishing a bound more grounded in reality would be very helpful. The theory work in \\cite{cheung2019superposition} should help better define these bounds.\n\nThere's a claim of \"improved generalization through enhanced forward transfer\", but it's unclear what experiments in this paper demonstrates better transfer learning.", "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": "Gh6RJ97Vq6", "reviewer": "Reviewer_rVeg", "summary": "The paper, titled Conditionally Overlapping Mixture of Experts (COMET), proposes a method aimed at overcoming limitations in existing sparse neural network architectures. COMET introduces a modular, sparse structure with a biologically inspired fixed routing approach that eliminates the need for task IDs and trainable gating functions, commonly associated with representation collapse and redundancy. Instead, the authors implement a k-winner-take-all cap operation, enabling experts to overlap based on input similarity. This approach aims to improve generalization and facilitate faster learning, validated across various tasks and architectures, including MLPs, Vision Transformers, and GPT-based models.", "questions": "1. Which components of COMET are inspired by the concept of biological random projection?\n\n2. How should the hyperparameter $k$ in Equation (3) be determined?", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. COMET presents a novel routing method that replaces trainable gating functions with fixed, biologically inspired routing, which is rare in modular neural network approaches.\n\n2. The proposed method is tested across diverse architectures and tasks, such as image classification, language modeling, and regression, suggesting versatility.", "weaknesses": "1. Section 3 lacks crucial methodological details necessary for a complete understanding of the proposed COMET approach. For instance, it is unclear which specific design elements in COMET were directly influenced by the concept of biological random projections. \n\n2. The paper’s writing lacks clarity, making it difficult to fully understand the design of COMET. I recommend including a preliminary section that outlines the foundational Mixture of Experts (MoE) framework, followed by a clear discussion on how COMET’s design diverges from and improves upon existing methods.\n\n3. While COMET is designed to improve modularity and interpretability, the authors do not demonstrate how the model’s interpretability has improved. More extensive interpretability metrics or qualitative evaluations would support the claimed benefits of COMET.\n\n4. The absence of code and essential implementation details significantly hampers reproducibility and raises concerns about the robustness of the results.", "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)