Spatial Mixture-of-Experts
Authors: Nikoli Dryden
Torsten Hoefler
Venue: NeurIPS 2022 Accept (2022)
Keywords: deep learning
mixture of experts
regression
weather prediction

Abstract: Many data have an underlying dependence on spatial location; it may be weather on the Earth, a simulation on a mesh, or a registered image. Yet this feature is rarely taken advantage of, and violates common assumptions made by many neural network layers, such as translation equivariance. Further, many works that do incorporate locality fail to capture fine-grained structure. To address this, we introduce the Spatial Mixture-of-Experts (SMoE) layer, a sparsely-gated layer that learns spatial structure in the input domain and routes experts at a fine-grained level to utilize it. We also develop new techniques to train SMoEs, including a self-supervised routing loss and damping expert errors. Finally, we show strong results for SMoEs on numerous tasks, and set new state-of-the-art results for medium-range weather prediction and post-processing ensemble weather forecasts.

Spatial Mixture-of-Experts
Nikoli Dryden
Torsten Hoefler ETH Zürich htor@inf.
ETH Zürich ndrydenGethz.ch ch
.ethz. wwwe@ifffetthh...mm
ch
Abstract
Many data have an underlying dependence on spatial location; it may be weather on the Earth, simulation on a mesh, or registered image. Yet this feature is rarely taken advantage of, and violates common assumptions made by many neural network layers, such as translation equivariance. Further, many works that do incorporate locality fail to capture fine-grained structure. To address this, we introduce the Spatial Mixture-of Experrss (SMoE) layer. a sparsely-gater layer that learns spatial structure in the input domain and routes experts at fine- fine-grained level to utilize it. We also develop new techniques to train SMOEs, including a self-supervised routing loss and damping expert errors. Finally, we show strong results for SMOEs on numerous tasks, and set new state-of-the-art results for medium-range weather prediction and post-processing ensemble weather forecasts.
Introduction
Many datasets exhibit an underlying, location-basse structure, where the value at point depends on where that point is. For example, weather predictions [78] depend on their location on Earth; many scientific simulations are parameterized on an underlying mesh [9]; data (e.g., faces) may be specifically aligned [91]; or it may approximately hold, as in natural images with centered objects. Further, tasks on such data are often dense regression tasks, such as predicting weather several days in the future. Numerous architectures have been successfully applied such tasks. Convolutional neural networks (CNNs) [80] and transformers [74] show promise for medium-rane -aange weather prediction. Locally cocally-connected networks (LCNs), which use independent, rather than shared, filters at each point, have been applied to weather post-processing [38], face recognition [46, 77, 91], and other tasks [18, 71], specifically to learn local features. Low-rank local connectivity (LRLCN) [30] relaxes the translation equivariance of CNNs while requiring fewer parameters than LCNs. Other approaches, such CoordConv [65], add an additional inductive bias by providing explicit input coordinates.
However, for tasks on data with location- structure, prior approaches suffer from various limitations. Convolution assumes that data is translation equivariant, which does not hold for such tasks [46, 91]. Approaches like LCNs require many parameters. Many architectures have been designed for classification tasks and fail to perform well on regression because they opererate at too coarse granularity and are unable to capture key details. This limits their applicability to important tasks, such as medium-range weather prediction or climate simulations. Indeed, simple heat diffusion task with occation-dependen diffusivities (see 13.1), many approaches do not learn the location dependence at all and instead converge to "average" diffusivity.
To address this, we introduce novel neural network layer, the Spatial Mixture-of- -xpertsssssss (SMoE) layer ($2). An SMOE uses a learned gating function to sparsely select and route from a shared set of experts to each spatial location (e.g., pixel) in an input sample. This enables experts to specialize to the unique characteristics of different "regions" within the data and easy interpretability by analyzing the gate. SMOEs require the assumption that all samples in the dataset have a similar underlying spatial structure, but this often holds (at least approximately) for many datasets, such weather, where each example is on the same latitude/longitude grid. For the SMOE gating function, we introduce tensor routing, simple and cheap routing function that effectively learns this spatial 36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

CoundErry Shot with
Convolution:
Low-rank locally connected:
Mixture-off-xpper
location
Spatial
Ritrr
Bash fitres
canteenn aodi
levmed
Inputs
Experts
Comvolation
Mixture-of offiuure-f-Experts
Gale
stant
Inductive blas
Lnweeank hoadly
compute
Mlld LCN
connected
Cle onton
Locally-connected
Spatial oftture-of-Expent
Spatial Mixture-ol-Exxeee
Rites
Input
locally connected
rouiin
leam locaioe-
locaiion
depontion
functions
Fulpecoereciee
Regrerdion
Parameters
Figure 1: Comparison of SMoEs to other networks. Left: Qualitative comparison of parameter counts and
inductive biases. Center:
Pattern of filter application among models. Right: Routing in MoEs and SMOEs.
structure ()..1). We also find that end-to-end training approaches, as in standard mixture-of-experts (MoEs) [82, 86], do not capture location-dependern classification loss ($2.2) to train the gate, a self-supervised loss which provides a better learning signal the gate about incorrect routings. We also introduce expert error damping ($2.3), which limits expert updates from misroutings, and further improves performance. Both methods rely on extracting information from the error signal (gradient of the layer output) which would otherwise be uninformative. Figure provides an overview of SMoEs and a qualitative comparison to important related work. With these methods, SMoEs are able learn fine-grain, vecation-dependee structure which other architectures miss, and consequentially deliver much better performance.
well. Instead, we introduce separate routing
SMoEs are related to and inspired by prior work in MoEs (e.g., [82, 86, 102]), but there are key differences. Existing MoEs use coarse marse-grained routing at the sample [99]. token [86], or patch [82] level, and experts produce coarse output (e.g., entire samples or channels) while SMOEs route a fine-grained, per-pixel level. Fine-grained routing is important for enabling experts to specialize to fine-scale information, which may be crucial [74, 84J. Such MoEs also typically aim to increase model capacity while keeping compute costs constant, whereas the goal of SMOES is to capture spatial dependencies. Other work has aimed to incorporate spatial dependence, such as LRLCNs [30], which learn content- and location-dependen weights to combine set of basis filters. However, we found that these prior methods failed to capture the fine-grained features necessary for good performance on dense regression tasks ($3.1).
We conduct experiments on several benchmark datasets with SMoEs and conduct extensive ablation studies of SMOE design decisions ($3). Using a simple heat diffusion dataset, we showcase the limitations of other models and the power of SMOES in controlled situation. We then apply SMOES to medium-range weather prediction and outperform the state-of-the-art [53, 80] on Weath- erBench [78]; and set a new state-of-the-art for post-processing ensemble weather forecasts on the EN dataset [8]. Finally, we show that SMOEs can also be applied to image classification tasks, where we match or outperform LRLCNs while using fewer parameters.
Our code is available at https: //ggt//btttttttt wwcchhuu..oomeeeeeee coon/spc//smos smoe.
1.1 Related Work
Mixture-of-Experte and conditional computation. While MoEs have long been used [15, 49, 52], since the advent of deep learning, there has been much work in applying MoEs, conditional computation, and dynamic routing to DNNs [1, 4, 5, 7, 10-12, 19-21, 23, 25, 27, 29, 33, 34, 39, 51, 59, 60. 70, 75, 76, 82, 83, 86, 99, 102, 107]. Often, the goal is to increase model capacity without a corresponding increase in compute costs and models use coarse-graine, graiined routing. Notably, MoEs route at the sample (e.g., [25, 29, 70, 83]), token (c.g.. [33, 59, 86]), patch (c.g., [82]) level; are typically trained with extensive auxiliary losses; and often target classification tasks. This coarse routing, in particular, means there is limited opportunity for experts to specialize. MoEs specifically for vision tasks also typically operate at a sample-level (e.g.. [1, 4, 39, 99, 102]) and use experts to specialize filter channel choices. In contrast, our SMOES induce an inductive bias via fine- grained, recation-dependee
routing for problems with such underlying structure; typically do not need auxiliary losses beyond the routing classification loss; and work well for dense regression tasks.
Local spatial structure. Locally connected networks (i.e., convolution filters without weight sharing) were historically successful in vision tasks such as facial recognition [18, 46, 77, 91]. Further

---

Experts &
Expert Error Damping ($2.3)
E 10==1ourr output
Damp error signals for incorrect routings
channels
Incorrect routings
MSE loss
L N
Xx
v)²
Input
top sparse routing ($2)
H1
Select experts per pixel
input
Apply experts at locations
Error signal
W
channels
dL
2
Y)
dx
N
Routing Classification Loss ($2.2)
Compute Loss
Construct labels
Identify incorrect routings with error signal magnitude
1 1/2 0 1
Gate G, tensor routing
D |E| (62.1)
Binary cross- entropy
0 1/2
0 0
Forward
Backward
1 0
0
Overview SMOE architecture and training using a mean-square errоr loss and an example input. Each location in an input is routed to E' experts from a set & based on gate G (&82 and 2.1). @ The gate trained using routing classification loss, which identifies incorrect routings based on the error signal ($2.2). ① The error signal to experts is damped for locations they were incorrectly routed to ($2.3).
Figure 2:
work incorporated periodic or tiled weight sharing [37, 71, 105]. However, more recent work often exclusively uses convolution, which has been found perform better [72] while requiring significantly fewer parameters. CoordConv [65] explicitly provides pixel coordinates input. Because of their structure, a single convolution CoordConv layer cannot learn location -eaationndeeedeent sparse activations as SMoEs do, and multi-layer networks give limited improvements in practice. Low-rank locally connected networks [30] learn to combine small number of basis filters location- and input- dependent manner. These combiners lack sparsity and apply all basis filters at each point, and are softmax- normallzed, which can result few filters dominating. limiting diversity. Separate work has used attention to provide position-dependerr scaling integrate context [35, 50, 64, 93, 95, 100]. This has culminated in vision transformers [26] and variants [28, 62, 67, 96, 97], which use specialized attention mechanisms to integrate spatial information. Other work uses Markov [61, 66] or conditional random fields [16, 43, 106], or graph-based methods [17, 98, 104], to learn long-range dependencies. Additional work has studied incorporating equivariance properties into neural networks (e.g., [13]).
Sparsity. Our use of sparse routing resembles work on sparsity in DNNs in general [45]. Many approaches learn fine-grained mask to identify which weights to keep or prune (e.g., [63, 90]), while other approaches sparsify activations (e.g., [3, 47, 56, 68, 69]). These works use sparsity to improve runtime performance, typically at inference time.
Weather prediction. Medium-range weather prediction, or forecasting roughly three seven days in advance [6]. is of broad societal import [57]. There has been much interest in applying deep learning community benchmark. Deep learning has also been successfully applied to the related tasks of ensemble post-processing [38, 79] and now- cow-casting (forecasting a few hours in advance) [2. 31, 81, 87-89]. Standard CNNs are currently of-t-off-hhe-nn art on WeatherBench [80] (although graph neural networks show promise similar data [53]) and are competitive on post- post-processing tasks [8]. SMOEs can take advantage of the extensive lne-grained, location dependent structure in weather data for improved performance.
this task [85], and WeatherBench [78] serves
Spatial Mixture of Experts
We now introduce the Spatial Mfxture-of-Essee -f-peerrtss (SMoE) layer. An SMOE uses learned gating function ($2.1) 1) to select specific experts from shared set to apply at each point (e.g., pixel) in an input. The gate learns the underlying spatial structure of the data, and routes specific experts to different "regions", allowing them to specialize. An SMOE is predicated on the assumption that the spatial structure similar across all samples in a dataset. We also assume data are on Cartesian mesh (i.c., grid), although this could be generalized. To train SMOEs, introduce self-supervised

---

routing classification loss ($2.2) and expert error damping ($2.3), which we found essential for achieving the best performance. Figure provides an overview of SMOEs and their training.
We first define the SMOE layer in full generality, then discuss the particular implementation we use. (See $3.1 for ablations.) Let x: E RCxBxx be input sample (we assume 2D data for simplicity). The SMOE layer consists of set & of experts, of which E' S [&] will be selected at each point, and a gating function G : RCXH + IRIE The gating function has exactly E nonzeros at each spatial point (HxW), which correspond to the selected experts at that point. Each expert applied at the points in æ where it has been selected, and may include additional context from x (e.g.. surrounding points if is convolution), and produces an output of size F. The outputs of each expert at point are then concatenated to form the output channel dimension (of size E-F) of the SMOE. More precisely, let gather, (-) select only the input entries where is nonzero. Then an SMoE is:
& is
y g(t (G(z) (e.):.....:ee (x))),
where is element wiiee multiplication, [e: stacks tensors, and y E RRFxHHw When E < [E], the gating function induces significant sparsity; this allows an SMOE to compute experts only at the points where they are to be applied, avoiding most computation.
This formulation yields a weighted SMOE, where the expert outputs are scaled by the gating function. An alternative, which is slightly more efficient and can be more readily interpretable, is an unweighted SMOE, in which case experts are not scaled, and the gather only uses G(r)) to select experts.
In this work, we focus simple, yet powerful, SMOE layer using convolution filters experts:
y/ = G(x);
Σ Wi. * Tex
cE(C] where * is cross-correlation, i E [[E]], and we have elided the gather for simplicity.
2.1 Gating Functions
The gating function G in an SMoE is critical for learning the underlying location-dependencee of the data. While G can be any function, good gate should be relatively cheap. We follow prior work on MoEs (e.g.. [82, 86)) and use top-E routing to select experts and sparsify the gate: G(z) = topg (g(z)), where g(x) E R|E)xXxW' is a learnable gating layer and topg selects the E largest entries along the first (expert) dimension and sets the remaining entries to 0. However, found that normalizing the gating function using softmax, as is typically done, was not necessary.
There are many options for the gating layer. Many MoE works have used MI.Ps [82, 86), but we found this did not perform well. Convolution or CoordConv [65] also did not perform well ($3.1), 1). Instead, we find that simple tensor routing gating layer worked best. With this, g(x) D E |g |xHXW where D is a learnable tensor that directly encodes the underlying location dependence and routing structure without depending on the input. Further, uses one parameter per expert per location and requires no computation beyond optimizer updates, making it an efficient choice.
We initialize DD using uniform distribution over -3] 3 El/EF), which corresponds a Kaiming uniform initialization with fan-in //E| [41]. However, in many practical cases, some data about locations may be known (e.g., whether it is land or sea when doing weather prediction). In such cases, D can be initialized based on this data to assign groups of experts in advance. This allows the gating function to benefit from prior knowledge while still being able to adjust the routing. Additionally, network may contain many SMOE layers, in which case we can share DD between layers that have the same spatial dimensions and number of experts, which reduces the overhead of the gate.
Finally, many MoE formulations use a number of auxiliary losses to ensure gates learn good routing policies and avoid mode collapse (e.g., [10, 82, 86]). When training with the routing classification loss we propose, we did not find additional auxiliary losses to be necessary.
2.2 Training and the Routing Classification Loss
Training MoEs is typically done end-to-end, with the experts and gate learning simultaneously, and sparse gradients based on the routing. We found this did not lead to good performance with SMOEs and tensor routing, particularly on regression tasks, as the gate did not learn well: it rarely changed its routing decisions from initialization. We hypothesize that this is due to "mismatch" in

---

the gradients on regression tasks, where they are informative for experts but not the gate, because regression aims to make a continuous prediction over both positive and negative values, whereas selecting an expert requires threshold (see ŞB). To address this, we train the gate with separate, self-supervised loss function, the routing classification (RC) loss. The key idea to consider routing as dense, multi-label classification task: selecting the "correct" experts at each point (ef. semantic segmentation). The RC loss does exactly this, and trains the gate by constructing appropriate labels. This also helps avoid mode collapse, where only small number of experts are selected (see $3.2).
In order to construct these labels, we need to determine whether the gate selected the correct set of experts at each point. However, such information is not directly available. Instead, we use the error signal into the SMoE layer (i.e., the gradient of the layer output w.r.t. the loss) as a proxy, and say that the routing was incorrect when the error signal has a large magnitude for the expert at that point, as this will imply a correspondingly large gradient update. Further, in the case of a mean-square error loss I (commonly used for regression), the error signal can directly encode the prediction error. Let X be the predictions, Y the true values, and N the number of elements in X. Then the error signal is:
dL d 1 -
2
N (X YY-
dx dx N Hence, the error signal is simply the (scaled, signed) error of the predictions. While this exact relation ceases to hold as backpropagation proceeds, the intuition behind the error signal magnitude remains.
We now define the routing classification loss. Given the error signal e into an SMOE layer, and an error quantile q (a hyperparameter), we say that selecting an expert at a point was incorrect if at that expert and point is greater than the qth quantile of e, and correct otherwise. We use quantiles as they are independent of E's scale, which may be hard determine and change during training. We then construct the labels for each point as follows: Unselected experts start with label 0. A correctly selected expert has label 1. Finally, if an expert was incorrectly selected, its label is and we add 1/LE1 F) to the label value of each unselected expert (note |E| - F is the number of unselected experts). This corresponds to a uniform prior that the correct expert could be any unselected expert. With these labels, the RC loss for a gate is then the binary cross-entropy loss of the gate output.
2.3 Expert Error Damping
While the RC loss enables an SMOE gate to trained directly based on whether it routed correctly, experts that were incorrectly routed to still perform gradient updates based on this routing. This results in experts updating to improve their performance for locations they may not be applicd at in future training iterations after routing changes. To mitigate this, we propose expert error damping, where the portion of an error signal that corresponds to incorrect routings is damped to limit incorrect updates. We find that this can improve SMoE performance and reduce convergence time.
Expert error damping similar the RC loss, and we classify incorrect routings in the same way. We then scale the eтror signal into the experts by a constant factor at each point where the routing was incorrect. This will limit the magnitude of the update made in response to the misrouting.
2.4 Practical Implementation
We now discuss the implementation of an SMOE layer, primarily focusing on the simple SMoE with convolution filters we use. In an ideal implementation, the flop cost of an SMOE is the cost of the gate plus the cost of applying the selected experts. When using tensor routing gate, there are no flops in the gate. The flops from applying the selected convolutional experts is equivalent to a standard convolutional layer with number of filters equal to the number of selected experts. Hence, SMOEs are quite efficient flop-wise. However, recent work has shown that data movement is also critical to performance [48]. Because of this, we do not expect even well-optimized implementations to match the runtime of a convolution layer due to the sparse routing, limited locality in accessing expert filters, and other operations, although work on benefits. Despite this, during inference additional optimizations may be available because the tensor routing gate does not depend on the input and computations could be reordered to maximize locality.
sparsity [22] should offer
Unfortunately, efficiently implementing the irregular access patterns in an SMOE is challenging in standard Python frameworks, and likely requires custom compute kernels. Instead we opt for a naive implementation in Py PyTooc [73] where we apply all experts at all points and then use gather and scatter operations to implement sparse routing, Experts are concatenated in sorted routing score
5

---

True region mask diflusion stenciis SMoF routing map expert weights
Model
#Params
Epochs % within 1%
Convolution
146
102
91.3=
CoordConv [65]
56
91.1=0.2
CondConv [102]
200
120
91,2*0.3
LRLCN [30]
516
27
91.8*6.3
LRLCN-ND [30]
12315
91.5
LCN
36764
110
100.0 +
FC
16.7M
250
14.2
ViT [26]
200k
67
93.5
V-MoE [82]
470k
93.8
SMoE
27+12288
8
100.0
Figure 3: Heat diffiusion results. Left: SMOE and baseline performance. Center: Dataset region map and diffu- sion stencils. Right: MoE routing map and experts. SMoEs learn correct stencils and location-dependence.
order (i.e., as given by G(z))); while this can result channel orders changing during training, we find it has little impact (see $3.2). We find this implementation sufficient even for large-scale tasks.
3 Experiments
We now present experimental results using SMOEs and a variety of baseline models, as well as ablations of different SMOE components. First, we describe a simple location-dependent heat diffusion dataset, which we use to study the ability of different architectures to learn location dependence in a controlled environment ($3.1). We then present results on the WeatherBench [78] medium-range weather forecasting challenge ($3.2) and the ENS-10 [8] dataset for post-processing ensemble weather forecasts ($3.3), where SMoEs set new state-of-the-arr performance results. Lastly, we show results on several image classification tasks, illustrating the breadth of SMoE applicability even in situations without strict location-dependence ($3.4).
All results were run using PyTorch [73] version 1.11 on a large cluster with 16 GB V100 GPUs. We summarize training details throughout this section, and provide full details in $A. We use existing Py TyTchh implementations when available, and implemented methods ourselves otherwise. Unless noted, all SMOE results used 3x3 convolution kernel experts, unweighted tensor routing gates, RC loss, expert error damping, еrror quantile q = 0.7, and damping factor 0.1. We report the mean and standard deviation (mean tand) over ten runs for all experiments, except for WeatherBench and ImageNet, which used three runs. In total, we used about 30k GPU hours for experiments.
3.1 Location- Heat Diffusion
To study location-depenndence in controlled manner, generated heat diffusion dataset with location-dependent diffusivities. The task is to predict what the heat will be at the next timestep. given the current heat at each point. Because we know the exact diffusivities and their distribution in the data, it is easy to identify how well a model has learned this task. We first describe the dataset and its generation in detail, then discuss results and ablation studies on SMOE components.
The dataset consists of a region map, where each location is assigned a type, and each type corresponds to different heat diffusivity. The region map and diffusivity assignment are then fixed for the dataset. We generate the region map randomly using a simple preferential attachment method. This defines the location-dependenee that wish learn. To generate samples, first randomly distribute drops of heat across the domain, then apply five-point diffusion stencils at each point, using the diffusivity of the region type for each point. For simplicity, we use zero boundary conditions. This process is iterated to generate many timesteps from the given starting state. The dataset then consists of the generated timesteps for many different starting states.
Dataset.
a model is able to learn the diffusion stencils and the location -octtiondiffusivity correspondence, exactly predict the next timestep. Further, the diffusivity stencils are also simple, and exactly correspond to a 3x3 convolution kernel with no nonlinearity.
The particular dataset we use consists of 100,000 64x64 samples, with 1,000 initial states evolved for 100 timesteps each. Adding more samples did not significantly change results. There are three region types, with diffusivities 0.25, 0.025, and 0.0025. Figure 3 (center) shows the region map and

---

Table I: Effect of RC loss and expert error damping on the heat diffusion dataset.
Table 2: Effect of different SMOE gate functions the heat diffusion dataset. Gate
RC loss
Damping
Epochs
% within 1%
#arams
% within 1%
X
X
14
91,5
Fully- Colly-connected
50M
85,2
V
16
91.6
3x3 convolution
27
91.3+0.3
X
96.7 +
3x3 CoordConv [65]
91.5+0.4
V
8
100.0
3x3 LCN
110592
96.3
3x3 CoordConvx3
255
91.649.a
Tensor routing
12288
100.0#0
diffusion stencils. We report results using "% within 1%", the percentage of locations in sample that are within 1% relative error of the true value, as this is more interpretable than a -quare error.
Results. Figure 3 (left) shows results on the heat diffusion dataset for SMOEs and a number of baselines: CNNs, LCNS, fully-connected (FC) layers, CoordConv [65], CondConv [102], LRL-- CNs [30], vision transformers (ViT) [26], and vision MoEs (V-MoE) [82]. For the LCN and FC layers, we use only single layer additional ones showed no benefit. Convolution, CoordConv, CondConv, and LRLCN are the best network found among a set with up to three layers, 3x3 kernels, 12 filters per layer, batchnorm, and ReLU activations. ViT and V-MoË use one transformer block with a patch size of 4x4, an embedding dimension of 128, and four heads. LRLCNs use three basis filters and an input-dependent combiner. We also tried unshared combining weights with no input dependence (LRLCN-ND). V-MoEs select one expert from set of three. Our SMOEs use a single layer with three experts and select E = 1 expert per point. All models were trained with batch size 32, Adam [54] with a learning rate of 0.001 (decayed by 10x after no validation improvement for 15 epochs), and early stopping after no improvement for 30 epochs. Additional hyperparameter tuning did not significantly improve results. We report SMoE parameters as expert gate parameters.
SMoEs achieve perfect performance this dataset. Further, by examining the learned routing and experts (Fig. 3, right), we can see that it has indeed correctly learned the diffusion stencils and location-dependdence LCNs also achieve this, but require 3x more parameters, and require 110 epochs to converge (versus 8 for SMoEs). Fully-connected layers failed to learn the data well, likely due to the challenge of optimizing so many parameters. Other methods all converge between 91 and 94% within 1%. Examining their predictions and weights, we observe that they do not appear to have learned the decation-dependence of the diffusivities. and instead converged predicting with an "average" diffusivity across the domain. We also tried larger (deeper and/or wider) convolutional networks, but performance did not improve. MoE methods (CondConv and V-MoE) also fail in this manner, their coarse-grained experts are unable to specialize. Further, the LRLCN-ND fails in this manner, despite its architecture being similar to an SMoE when there is one output channel (a location-dependerr softmax-weightt combination of three basis kernels). We believe the LRLCN-ND exhibits a similar gradient "mismatch" as discussed earlier ($2.2).
We now discuss number of different ablations of the SMOE architecture and design.
What if the "right" expert configuration is not known? While in the above experiments, were able to select the SMOE expert configuration (number of experts, number of selected experts, expert filter size) so that it is both necessary and sufficient to learn the task, in many situations this information may not be available. We considered alternative SMOE configurations varying each of these parameters: ? using six experts; @ experts with 5x5 kernels; O and selecting two experts per location from six total. For case O, we summed the two SMOE output channels together.
In all three cases, the SMOE achieved 100.0**% within 1% on the heat diffusion task. In O, we found that they learned duplicate diffusion stencils and still routed them appropriately. @ learned the five-point stencil plus a boundary of near-zero values, thus being nearly identical to the 3x3 kernel. Finally, learned diffusion stencils that summed together to produce the correct diffusivity. Thus, we can see that SMOEs are robust and adapt well to these sorts of architecture choices.
RC loss and expert error damping. Table shows results for training SMOEs with and without our routing classification loss ($2.2) and expert error damping ($2.3). Without the RC loss, SMOE performance at par with other baselines in Fig. 3, but once is added, performance improves sig- nificantly as the gating function now learns the location-dependency in the data. Adding expert error damping further improves performance and convergence by limiting the impact of gate misroutings
7

---

Table 3: WeatherBench [78] results wlethtueeweiighted RMSE).
Table 4: ImageNet [24] validation accuracy.
7500 [m²
T850 [K]
Model
Top-1 %
Top-5 %
Model
3 days 5 days
3 days
5 days
ResNet-50 [42, 94]
80.8320.06 04 25.39=0.00
Rasp and Thuerey [80] 316*2.1 56325.
1.80
2.8420.03
LRLCN [30]
80,90 0z 95,41 25.41
SMOE after first layer
80.8549.05
95,40=20:01
LRI.CN [30]
SMoE after first layer
ViT (2x2) [26]
- 2x wide
4382*
305+1.9
310 55522*
290
54921.
638431
556+2.2
2.2420.0:
1.77+0.01
1.7620.03
1.73
2.88 os
2.80++0.0s
2.7820.01
2.7920..0
3x3 convs-SMoB
Wide ResNct-50-2 [103]
Last layer SMOE
81.33 o3 95.52
81.76 1. os 95.7440.0z
80.91 4 04 95.42-20.01
Last layer SMOE
298+2.6
553+3.2
1.73+0.02
2.78+0004
3x3 conV
278+2.0
530+1.
1.69+0.01
2.65++0.02
gate prior
270+1.9
525
1.66+0.02
2.60+0.01
rand fixed gate init 328
572
1.89+0.0s
2.96++0.00
R&T [80] (pretrained)
267=1-6
500
1,66
2.43
SMOE (pretrained)
253
488+
1,57+0.02
2.34+0002
extra ERA5
232
440+1.2
1.46+0.02
2.19
1.40
198
382+2.0
1.420
2.06 00.02
on expert learning. However, damping on its own offers little benefit, as it does not improve gate learning. These results show that these refinements are critical for good performance.
Gating function. Table 2 shows the performance of different gating functions ($2.1) on the SMOE. We consider six options: A single fully-connected layer (as is commonly used in MoEs [82, 86]); a single 3x3 convolution, CoordConv [65]. or LCN layer; gate with three CoordConv layers with batchnorm and ReLU; and our tensor routing gate. When training, we also considered auxiliary losses and other methods for improving performance (see $C) and report the best result. Our tensor routing offers the best performance. An LCN performs second-best bestt, likely because it also uses separate parameters per location, but uses 9x as many parameters and requires significant computation. Other methods do not appear able to effectively capture ceatiion-deendennc
Other ablations. We conduct a number of additional ablation studies in &C, including using auxiliary losses and routing noise during training, routing normalizations, and expert functions.
3.2 Medium-Range Weather Prediction
We now discuss results on the WeatherBench [78] medium-range weather forecasting benchmark. This benchmark uses the ERA5 reanalysis dataset [44], with hourly global weather data for 1979- 2018. We use the data subset suggested by Rasp et al. [78] at 5.625 resolution (32x64 grid points) and train on data from 1979-2015, validate on 2016, and report test results for 2017-2018. We otherwise follow the training methodology of Rasp and Thuerey [80]. The target quantities to predict are geopotential 500 hPa (7500) and temperature at 850 hPa (T850) with a three- and five-day lead time. We report results using weiihtd root-mean- -quare error (RMSE).
As baseline, we use the ResNet architecture [42] introduced by Rasp and Thuerey [80]. which currently reports the best results on WeatherBench. This architecture consists of 19 residual blocks each with two [3x3 convolution -> LeakyReLU - batchnorm - dropout] layers, plus initial 7x7 convolution layer. All convolutions but the last have 128 filters. We consider three additional baselines. The first is identical to the above, but with twice many filters (256) in each convolution. Second, we replace 3x3 convolutions with LRLCN [30] layers. Finally, we use four-layer ViT [26] with patch size 2x2, hidden dimension 1024, and eight heads (the best performing configuration).
We adapt the Rasp and Thuerey ResNet to use SMOEs with three configurations: adding an SMoE layer after the first convolution; adding an SMOE layer after the final convolution; and replacing all 3x3 convolutions with SMOE layers. Each SMOE selects the same number of experts as the original layer had filters, and has twice many experts (i.e., [ 256, E 128). We also share the tensor routing gate across all SMOE layers with the same number of experts, so its overhead is minimal.
Because the weather data is on a fixed grid with underlying decation-dependene (the Earth), we expect SMOES to convey some benefit by specializing to the characteristics different regions. In Table 3, we observe that this is indeed the case. Adding SMOES improves results in all situations, with the most significant improvement coming through replacing all 3x3 convolutions with SMOFs. This showcases the advantage of incorporating appropriate location-dependentt biases. Wider ResNets
8

---

Table 5: Results for prediction correction on the ENS-10 [8] dataset for ensemble post-processing.
7,500 [m² 3]
T850 [K]
T2m [K]
Metrie Model
5-ENS
10-ENS
5-ENS
10-ENS
5-ENS
10-ENS
EECRPS CRPS
EMOS
79.12 79.1220.12
78.8020.2
0.72122.01
0.72022.00 0.706
0.71 0.71120000 20005
U-Net 76.54=0:20
76.1840.12
0.685
0.670
0.657=9.01
0.644490.0
SMoE 68.9420-14
67.43
0.61220.01
0.590 2:0.02
0.601
20.02
0.594
EMOS 29.21
29.0240.13
0.24740:00
0,245440.2
0.244420000
0.24140002 20.02
U-Net 27.78
27.55
0.230
0,229 1:0.03
0.225
20.00
0.220
SMOE
23.79-2.20 20.20 23.1020.16
0.207 20.03
0.197 29.03
0.199 20.01
0.190 40002
offer limited improvement (in line with results reported by Rasp and Thuerey [80]). The location-
dependent filters of LRLCNs improve over ResNets, but fail to match SMOEs. We were unable to achieve good performance with ViTs, but did observe that they are highly sensitive to patch size.
Incorporating prior knowledge into gates. While the exact nature of the location-dependence of this data is unknown, we do have broad prior some aspects of it, such as whether a point is land or sea. This information can be incorporated into an SMOE by initializing the tensor routing gate to bias routing to different experts. To this end, we use the land-sea mask from ERA5 to initialize the gate to route land locations to half the experts and sea locations to the other half. Note this does not fix the routing, as the gate is able adjust as it learns. Further, the land-sea mask already included in the input data, so all models already had access to this information.
Results with this are in the "+ gate prior" line of Table 3, and perform best. This configuration sets new state-of-ther tfrte--rt for WeatherBench when not using additional data. Indeed, it nearly matches the performance of ResNet with 150 years of additional pretraining data from climate simulations [80].
We also tried a configuration where the gate was initialized randomly and fixed rather than learned ("rand fixed gate init"). This performs worse than our baseline, as the network cannot adapt its routing choices, and each expert sees even fewer points in each sample than a standard network, resulting in
less learning. Thus, learning the routing function is critically important to good performance.
Additional data. Following Rasp and Thuerey [80], we use 150 years of data from the MPI-ESSM- HR climate model from the CMIP6 archive [32] to pretrain our best SMOE configuration, which was then fine-tuned on ERA5 data as above. This significantly outperforms both our SMOEs without pretraining and Rasp and Thuerey's 's pretrained ResNet. We incorporated more data to further push the performance by adding ERA5 data from the most recent back extension 1979),1979,, increasing the dataset size by about 50%. This shows improved results; however, we suspect performance is saturating due to the coarse spatial resolution of the data. We therefore trained a final configuration with higher resolution (1.4") data. Using this, our SMOEs significantly outperform the ofte-of-theeart on WeatheBBench; indeed, its performance on T850 is very close to that of the operational Integrated Forecast System [78]. Our results are also competitive with those of Keisler [53], although they are not directly comparable (due to, e.g., different data resolutions).
Mode collapse. Many MoEs suffer from expert or mode collapse (c.g.. [10, 82, 86]). where only a small number of experts are selected. This is typically avoided with routing noise and/or auxiliary "load balance" losses. On the heat diffusion dataset, we found these losses to offer no benefit (SC). We also did not observe mode collapse SMOEs on WeatherBench. With the RC loss, we directly train the gate, updating routing weights toward other experts after mistakes, and so avoid such issues.
Expert selection order. During training, the order experts are concatenated may change (due to changes in relative routing scores, or selecting different experts), which will impact the order of channels seen by subsequent layers. When training on WeatherBench, we found this not to have a significant impact: expert order stabilizes early, allowing layers to operate on stable representations. Further, most "swapping" occurs among low-confidence experts, so is limited to subset of channels.
3.3 Post- Post-Processing Ensemble Weather Forecasts
Numerical weather prediction systems typically utilize ensembles of simulations in order to quantify uncertainty and improve forecast quality [14]. However, such ensembles typically exhibit systematic biases [92], and correcting them improves forecast skill [14, 85, 101], a task for which deep learning has shown promise [38, 79]. We use the ENS-10 dataset [8], which consists of twenty years (1998-

---

2017) of global reforecast [40] data at 0.5* spatial resolution. We follow the benchmarking setup of Ashkboos al. [8], and correct predictions for Z500, T850, and meter temperature (T2m) at 48 hour lead-time using both five and ten ensemble members. We report results using continuous ranked probability score (CRPS) and extreme event weighted CRPS (EECRPS).
We adapt the U-Net model from the ENS-I0 baselines, as it delivers good performance and operates on global data (other methods use patches). Similar to our approach for WeatherBench, we replace each 3x3 convolution with an SMOE with four times as many experts as the original layer, and select the same number of experts as the original layer had filters. We share tensor routing gates between all layers with the same spatial dimensions and number of experts, with the exception that the encoder and decoder trunks also use separate gates. As baselines, we use the original U-Net architecture and Ensemble Model Output Statistics (EMOS) [36], a standard post-processing method.
We observe in Table 5 that, similar WeatherBench, SMOEs offer significant improvements in forecast skill across all situations, and set a new state-of-the-art for prediction correction on the ENS-10 10 dataset. This also demonstrates that SMOEs are able to scale to the very large spatial domain used by the ENS-10 data and still learn the appropriate location dependence.
3.4 Image Classification
Lastly, we present results on several image classification tasks; we focus here on ImageNet- [24] and discuss results on additional datasets in $D. While these datasets do not have a strict location- dependent structure, relaxing the strict translation equivariance of convolutions can bring benefits, and enables a direct comparison with Elsayed et al. [30]. We follow their experimental methodology and train using the recipe of Vryniotis [94]. We either insert an SMoE layer after the first or last convolutional layer of ResNet-50 [42] or replace all 3x3 convolutions with SMOE layers. Our SMOEs have twice as many experts as the original convolution layer and select half of them, to keep output dimensions constant. Gating layers are shared among all equally-sized blocks. For comparison, we also train ResNet-50 with all 3x3 convolutions replaced by LRLCN [30] layers; and a Wide ResNet-50-2 [103], which has comparable parameters to SMOEs.
Table 4 shows that SMOEs outperform LRLCNs when we replace all 3x3 convolutions, while using 56% of the parameters. However, a wide ResNet performs best overall. Nevertheless, this shows that ImageNet classification does indeed benefit from relaxing translation equivariance.
4 Discussion
We presented the Spatial Mixture-of-Experts layer, a novel layer that learns underlying location dependencies in data and then uses fine-grained routing to specialize experts to different areas. We also introduce routing classification loss and expert error damping, which enable SMOEs to perform well on dense regression tasks. Prior work shows limited effectiveness on these tasks: Either it does not capture
(e.g., convolutions) or operates at a coarse-grained level (c.g.. standard MoEs). By overcoming these challenges, we show a new capability for neural networks, and set new otate-of-the-arts for medium-rant renie weather prediction and ensemble post-proe post-processing.
Many other problems broad societal import have a similar spatial structure, particularly in scientific domains [9], and we expect SMOEs to be applicable to them. However, tasks such facial recognition and surveillance have also historically shown benefit from such improvements [91] and SMOEs should therefore be used with care.
SMoEs show that learning Docation-dependenn is a powerful inductive bias for certain types of data, and there are many avenues for further study. Two key areas of particular interest are to develop improved implementations for fine- grained, sparse routing; and to generalize SMOEs from operating on grids to general graphs, which would enable them to be applied to many additional tasks.
A Acknowledgements and Disclosure of Funding
We thank the members of SPCL at ETH Zürich Eecs, and Peter Dueben and Mat Chantry of ECMWF, for helpful discussions; and the anonymous reviewers for their suggestions and feedback. This work has received funding from the European High- Pegf-Peeformancc Computing Joint Undertaking (JU) under grant agreement No. 955513 (MAEL STROM), and from Huawei. N.D. received support from the ETH Postdoctoral Fellowship. We thank the Swiss National Supercomputing Center (CSCS) and Livermore Computing for computing infrastructure.
10

---

References
[1] Alhabib Abbas and Yiannis Andreopoulos, 2020. Biased mixtures of experts: Enabling computer vision inference under data transfer limitations. IEEE Transactions on Image Processing 29 (2020). arXiv:2008.0966 [cs.LG]
[2] Shreya Agrawal, Luke Barrington, Carla Bromberg, John Burge, Cenk Gazen, and Jason Hickey. 2019. Machine learning for precipitation nowcasting from radar images. (2019). 19X1v:1912.12132 [cs.CV]
[3] Subutai Ahmad and Luiz Scheinkman. 2019. How be dense? The benefits of using highly sparse representations. (2019). arXiv:1903.1125 [cs.LG]
[4] Karim Ahmed, Mohammad Haris Baig, Lorenzo Torresani. 2016. Network of experts for large-scale image categorization. In European Conference on Computer Vision (ECCV). arXiv:1604.061 19 [cs.CV]
[5] Amjad Almahairi, Nicolas Ballas, Tim Cooijmans, Yin Zheng, Hugo Larochelle, and Aaron Courville. 2016. Dynamic capacity networks. In International Conference on Machine Learning (ICML).
arXiv:1511.07838 [es.L.G]
[6] American Meteorological Society, 2022. Medium-range forecast. https::/gloosaary.amm www.eaarr..mettocoorr.g.o
Makki/Meeium-raaa _forecast
[7] Mikel Artetxe, Shruti Bhosale, Naman Goyal, Todor Mihaylav, Myle Ott, Sam Shleifer, Xi Victoria Lin, Jingfei Du, Srinivasan Iyer, Ramakanth Pasunuru, al. 2021. Efficient Large Scale Language Modeling with Mixtures of Experts. arXiv:2112.1066 [cs.CL.]
[8] Saleh Ashkboos. Langwen Huang, Nikoli Dryden, Tal Ben-Nun, Peter Dueben, Lukas Gianinazzi, Luea Kummer, and Torsten Hoefler. 2022. ENS-10: Dataset For Post-Processing Ensemble Weather Forecast. Proceedings of the Neural Information Processing Systems (NeurIPS) Track on Datasets and Benchmarks. arXiv:2206. 14786 [es.LG]
[9] Nathan Baker, Frank Alexander, Timo Bremer, Aric Hagberg. Yannis Kevrekidis, Habib Najm, Manish Parashar, Abani Patra, James Sethian, Stefan Wild, 2019. Workshop report on basic research needs for scientific machine learning: Core technologies for artificial intelligence. Technical Report. United States Department of Energy, Office of Science.
[10] Emmanuel Bengio, Pierre-Lue Bacon, Joelle Pineau. and Doina Precup. 2016. Conditional computation in neural networks for faster models. In International Conference on Learning Representations Workshops (ICLR-W) W). 1511v:151..00 066297 [es.LG]
[11] Yoshua Bengio. 2013. Deep learning of representations: Looking forward. In International conference on statistical language and speech processing. arXiv: [cs.LG]
[12] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. 2013. Estimating propagating gradients through stochastic neurons for conditional computation. (2013). arXiv:1308.344 3432 (es.LG)
[13] Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. 2021. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. (2021). arXiv:2104.13475 [cs.LG]
[14] Roberto Buizza and David Richardson. 2017. 25 years of ensemble forecasting at ECMWF. ECMWF Newsletter (2017). Issue 153. https:/ //ww/ ecmwf. iatt//ode/88199
[15] Ke Chen, Lei Xu, and Huisheng Chi. 1999. Improved learning algorithms for mixture of experts in multiclass classification. Neural networks 12, (1999).
[16] Liang-Chich Chen, George Papandreou, Iasonas Kokkinos, Kevin Murphy, and Alan I. Yuille. 2017. DeepLab: Semantic image segmentation with deep convolutional nets, atrous convolution, and fully connected CRFs. JEEE transactions on pattern analysis and machine intelligence 40, 4 (2017). arXiv:1606.00915 [cs.CV]
[17] Yunpeng Chen, Mareus Rohrbach, Zhicheng Yan, Yan Shuicheng, Jiashi Feng, and Yannis Kalantidis. 2019. Graph- based global reasoning networks. In Conference Computer Vision and Pattern Recognition (CVPR). arXiv:1811.1 12111 [cs.CV]
[18) Yu-hsin Chen, Ignacio Lopez Moreno, Tara Sainath, Mirkó Visontai, Raziel Alvarez, and Carolina Parada. 2015. docally-connecee and convolutional neural networks for small footprint speaker recognition. In INTERSPEECH.
11

---

[19] Zhourong Chen, Yang Li. Samy Bengio, and Si Si. 2019. You look twice: GaterNet for dy- namic filter selection in CNNs. In Conference on Computer Vision and Pattern Recognition (CVPR). arXiv:1811.1 1205 [cs.CV]
[20] Kyunghyun Cho and Yoshua Bengio. 2014. Exponentially increasing the capacity-to-computttion ratio for conditional computation in deep learnin. (2014). arXiv:1406.736. [stat.ML]
[21] Aidan Clark, Diego de las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, et al. 2022. Unified Scaling Laws for Routed Language Models. arXiv:2202.01 169 [cs.CL]
[22] Shail Dave, Riyadh Baghdadi, Tony Nowatzki, Sasikanth Avancha, Aviral Shrivastava, and Baoxin Li. 2021. Hardware acceleration of sparse and irregular tensor computations of MI. models: A survey and insights. Proc. IEEE 109, 10 (2021). arXiv:2007.0066 [cs.AR]
[23] Andrew Davis and Itamar Arel. 2013. Low-rank approximations for conditional feedforward computation in deep neural networks. In International Conference on Learning Representations Workshops (ICLR-W). -- arXiv:1312.446
[cs.LG]
[24] Jia Deng. Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. 2009. ImageNet: A large-scale hierarchical image database. In Conference on computer vision and pattern recognition (CVPR).
[25] Ludovic Denoyer and Patrick Gallinari. 2014. arXiv:1410.0510 [cs.LG]
Deep sequential neural network. (2014).
[26] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. 2021. An image is worth 16x16 words: Transformers for image recognition scale. In International Conference on Learning Representations (ICLR). arXiv:2010.1 1929 [cs.CV]
[27] Nan Du, Yanping Huang, Andrew M Dai, Simon Tong. Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. 2021. GLaM: Efficient Scaling of Language Models with Mixture-of-Expert arXiv:21
[cs.CL]
[28] Stéphane d' Ascoli, Hugo Touvron, Matthew I. Leavitt, Ari S Morcos, Giulio Biroli, and Levent Sagun. 2021. VIT: Improving vision transformers with soft convolutional inductive biases. In International Conference on Machine Learning (ICMI.) Xiv:21033.0699 [cs.CV]
[29] David Eigen, Mare' Aurelio Ranzato, and Ilya Sutskever. 2013. Learning factored representations in a deep mixture of experts. (2013). arXiv: 1312.4314 [cs.LG]
[30] Gamaleldin Elsayed. Prajit Ramachandran, Jonathon Shlens, and Simon Kornblith. 2020. Revisiting spatial invariance with low-rank local connectivity. In International Conference on Machine Learning (ICMI.). arX25:2002.02959 [cs.CV]
[31] Lasse Espeholt, Shreya Agrawal, Casper Sønderby, Manoj Kumar, Jonathan Heek, Carla Bromberg, Cenk Gazen, Jason Hickey, Aaron Bell, and Nal Kalchbrenner. 2021. Skillful Twelve Hour Precipitation Forecasts using Large Context Neural Networks. (2021). 1:Xiv:2111107440 [cs.LG]
[32] Veronika Eyring, Sandrine Bony, Gerald A Meehl, Catherine A Senior, Bjorn Stevens, Ronald Stouffer, and Karl E Taylor. 2016. Overview of the Coupled Model Intercomparison Project Phase 6 (CMIP6) experimental design and organization. Geoscientific Model Development 9, 5 (2016).
[33] William Fedus, Barret Zoph, and Noam Shazcer. 2021. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv:2101 01.096611111111 [cs.LG]
[34] Chrisantha Fernando, Dylan Banarse, Charles Blundell, Yori Zwols, David Ha. Andrei A Rusu, Alexander Pritzel, and Daan Wierstra. 2017. PathNet: Evolution channels gradient descent in super neural networks. arXiv preprint arXiv: 070110007770300 (2017). arXiv: 1 701108734 [cs.NE]
[35] Hiroshi Fukui, Tsubasa Hirakawa, Takayoshi Yamashita. and Hironobu Fujiyoshi. 2019. Attention branch network: Learning of attention mechanism for visual explanation. In Conference on computer vision and pattern recognition (CVPR), 1rXiv:1812. 10025 [cs.CV]
[36] Tilmann Gneiting, Adrian E Raftery, Anton H Westveld III, and Tom Goldman. 2005. Calibrated probabilistic forecasting using ensemble model output statistics and minimum CRPS estimation. Monthly Weather Review 133, 5 (2005).
12

---

[37] Karo Gregor and Yann LeCun. 2010. Emergence of complex- like cells in temporal product network with local receptive fields. arXiv preprini arXiv:1006.044: (2010). arXiv:1006.044: [cs.NE]
[38] Peter Grönquist, Chengyuan Yao, Tal Ben-Nun, Nikoli Dryden, Peter Duchen, Shigang Li, and Torsten Hoefler. 2021. Deep learning for post-processing ensemble weather forecasts. Philosophical Transactions of the Royal Society A 379, 2194 (2021). arXiv:2005.0874. [cs.LG]
[39] Sam Gross, Mare' Aurelio Ranzato, and Arthur Szlam. 2017. Hard mixtures of experts for large scale weakly supervised vision. In Conference on Computer Vision and Pattern Recognition (CVPR). arXiv:1704.06366 [cs.CV]
[40] Thomas M. Hamill, Jeffrey S. Whitaker, and Steven Mullen. 2006. Reforecasts: An Important Dataset for Improving Weather Predictions. Bulletin of the American Meteorological Society 87, 1 (2006).
[41] Kaiming He, Xiangyu Zhang, Shanqing Ren, and Jian Sun. 2015. Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification. In International conference on computer vision (ICCV). McXXiv:5022008855 [cs
[42] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun, 2016. Deep residual learning for image recog- nition. In Conference on computer vision and pattern recognition (CVPR). arXiv:1512.03388 [cs.CV]
[43] Xuming He, Richard S Zemel, and Miguel A Carreira-Perpinnn.. 2004. Multiscale conditional random fields for image labeling. In Conference Computer Vision and Puttern Recognition (CVPR).
[44] Hans Hersbach, Bill Bell, Paul Berrisford, Shoji Hirahara, András Horányi, Joaquín Muñoz-Sabater, Julien Nicolas, Carole Peubey, Raluca Radu, Dinand Schepers, et al. 2020. The ERA5 global reanalysis. Quarterly Journal of the Royal Meteorological Society 146, 730 (2020).
[45] Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. 2021. Sparsity in Deep Learning: Pruning and growth for efficient inference and training in neural networks. Journal of Muchine Learning Research 22, 241 (2021). arXiv:2102.00554 [cs.LG]
[46] Gary B Huang, Honglak Lee, and Erik Learnd-Miller 2012. Learning hierarchical representations for face verification with convolutional deep belief networks. In Conference on Computer Vision and Puttern Recognition (CVPR).
[47] Kevin Lee Hunter, Lawrence Spracklen, and Subutai Ahmad. 2021. Two Sparsities Are Better Than One: Unlocking the Performance Benefits of Sparse-Spane Networks. (2021). arXiv:2112.13880 [cs.LG]
[48] Andrei Ivanov, Nikoli Dryden, Ben-Nun, Shigang Li, and Torsten Hoefler. 2021. Data movement all you need: A case study on optimizing transformers. In Conference on Machine Learning and Systems (MLSys). arXiv:2007. 00072 [es.LG]
[49] Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton. 1991. Adaptive mixtures of local experts. Newral computation 3, (1991).
[50] Saumya Jetley, Nicholas A Lord, Namhoon Lee, and Philip HS Torr. 2018. Learn to pay attention. In International Conference on Learning Representations (ICLR). 1XXiv:180400239) [cs.CV]
[51] Xu Jia, Bert De Brabandere, Tinne Tuytelaars, and Luc V Gool. [n.d.]. Dynamic filter networks. Advances in neural information processing systems (NeurIPS) ([n. d.J). arXiv:1605.0667 [es.LG]
[52] Michael Jordan and Robert A Jacobs. 1994. Hierarchical mixtures of experts and the EM algorithm. Neural computation 6, 2 (1994).
[53] Ryan Keisler. 2022. arXiv:2202.07575 [physics.ao-ph]
Forecasting Global Weather with Graph Neural Networks. (2022).
[54] Diederik P Kingma and Jimmy Ba. 2015. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR). a1Xiv:1412.6980 [es.I.G]
[55] Alex Krizhevsky, 2009. Learning multiple layers of features from tiny images. (2009).
[56] Mark Kurtz, Justin Kopinsky. Rati Gelashvili, Alexander Matveev, John Carr, Michael Goin, William Leiserson, Sage Moore, Nir Shavit, and Dan Alistarh. 2020. Inducing and exploiting activation sparsity for fast inference on deep neural networks. In International Conference on Machine Learning (ICML).
[57] Jeffrey K Lazo, Rebecca E Morss, and Julie L Demuth. 2009. 300 billion served: Sources, perceptions. uses, and values of weather forecasts. Bulletin of the American Meteorological Society 90, (2009).
13

---

[58] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. 1998. Gradient-based learning applied document recognition. Proc. IEEE 86, 11 (1998).
[59] Dmilry Lepikhin, HyoukJoong Lee. Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. 2021. GShard: Scaling giant models with con- ditional computation and automatic sharding. In International Conference on Learning Representations. arXiv:2006.1666 16668 [cs.CL]
[60] Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Lake Zettlemoyer. 2021. BASE layers: Simplifying training of large, sparse models. In International Conference on Machine Learning (ICML). 1.iv:2103.1677 [cs.CL]
[61] Chuan Li and Michael Wand. [n.d.]. Combining Markov random fields and convolutional neural networks for image synthesis. In Conference on Computer Vision and Pattern Recognition (CVPR). arXiv:1601.04550 04589 [cs.CV]
[62] Jingyun Liang, Jiezhang Cao, Guolei Sun, Kai Zhang, Luc Van Gool, and Radu Timofte. 2021. SwinlR: Image restoration using swin transformer. In International Conference on Computer Vision (ICCV). arXiv:2108.100 10257 [eess.IV]
[63] Tao Lin, Sebastian Stich, Luis Barba, Daniil Dmitriev, and Martin Jaggi. 2020. Dynamic model pruning with feedback. In International Conference on Learning Representations (ICLR).
arXiv:2006.02553 [cs.LG]
[64] Drew Linsley, Dan Shiehler, Sven Eberhardt, and Thomas Serre. 2019. Learning what and where to attend. In International Conference on Learning Representations (ICLR). arXiv: 100.08899 [cs.CV]
[65] Rosanne I. Liu, Joel Lehman, Piero Molino, Felipe Petroski Such, Eric Frank, Alex Sergeev, and Jason Yosinski. 2018. An intriguing failing of convolutional neural networks and the CoordConv solution. In Advances in Neural Information Processing Systems (NeurIPS). arXiv:1807.00 007.00327 [cs.CV]
[66] Ziwei Liu, Xiaoxiao Li, Ping Luo, Chen Change Loy, and Xiaoou Tang. 2017. Deep learning Markov random field for semantic segmentation. JEEE transactions on pattern analysis and machine intelligence [cs.CV]
40,8 8 (2017).
[67] Ze Liu, Yutong Lin, Yue Cao. Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. 2021. Swin transformer: Hierarchical vision transformer using shifted windows. In International Conference on Computer Vision (ICCV). 2rXiv:2103.... 14030 [cs.CV]
[68] E Majani, Ruth Erlanson, and Yaser Abu-Mostafa. 1988, On the K-winners-take-al network. In Advances in neural information processing systems (NeurIPS).
[69] Alireza Makhzani and Brendan J Frey. 2015. Winner-take-all autoencoders. In Advances in neural information processing systems (NeurIPS). arXiv:1409.2777 2252 [cs.LG]
[70] Mason McGill and Pietro Perona. 2017. Deciding how to decide: Dynamic routing artificial neural networks. In International Conference on Machine Learning (ICML).
[st
[71] Jiquan Ngiam, Zhenghao Chen, Daniel Chia, Pang Koh. Quoe Le, and Andrew Ng. 2010. Tiled convolutional neural networks. Advances in Neural Information Processing Systems (NeurlPS) (2010).
[72] Roman Novak, Lechao Xiao, Jachoon Lee, Yasaman Bahri, Greg Yang. Jiri Hron, Daniel A Abolafia, Jeffrey Pennington, and Jascha Sohl-Dickstein. 2019. Bayesian deep convolutional networks with many channels are gaussian processes. In International Conference on Learning Representations (ICLR).
arXiv:1810.05118 MT.)
[73] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary Vitioo Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. 2019. Py PoTrch: An imperative style, high-performance deep learning library. Advances in neural information processing systems (NeurIPS). arXiv:1912.0170:: [cs.LG]
[74] Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, et al. [n.d.]. NourCastNet: A global
high-resolution weather model using adaptive fourier neural operators. ([n. d.J).
1211
14

---

[75] Svetlana Paylitskaya, Christian Hubschneider, Michael Weber, Ruby Moritz, Fabian Huger, Peter Schlicht, and Marius Zollner. 2020. Using mixture of expert models gain insights into semantic segmentation. In Conference on Computer Vision and Pattern Recognition Workshops (CVPRW).
[76] Prajit Ramachandran and Quoe V Le. 2018. Diversity and depth in per-example routing models. In International Conference on Learning Representations (ICLR).
[77] Marc' Aurelio Ranzato, Joshua Susskind, Volodymyr Mnih, and Geoffrey Hinton, 2011. On deep generative models with applications to recognition. In Conference Computer Vision and Pattern Recognition (CVPR).
[78] Stephan Rasp, Peter D Ducben, Sebastian Scher, Jonathan A Weyn, Soukayna Mouatadid, and Nils Thuerey. 2020. WcatherBench: a benchmark data set for data-driven weather forecasting. Journal of Advances in Modeling Earth Systems 12, 11 (2020). arXiv:2002.0046) iphysics.ao-phl
[79] Stephan Rasp and Sebastian Lerch. 2018. Neural networks for postprocessing ensemble weather forecasts. Monthly Weather Review 146, 11 (2018). arXiv:180S.00 0005009091 [Lt
[80] Stephan Rasp and Nils Thuerey. 2021. Data-driven medium-range weather prediction with ResNet pretrained on climate simulations: new model for WeatherBench. Journal of Advances in Modeling Earth Systems 13, (2021). arXiv:2008.086 0866606888666/0111 (physics.ao-phe
[81] Suman Ravuri, Karel Lenc, Matthew Willson, Dmitry Kangin, Remi Lam, Piotr Mirowski, Megan Fitzsimons, Maria Athanassiadou, Sheleem Kashem, Sam Madge, et al. 2021. Skilful precipitation nowcasting using deep generative models of radar. Nature 597, 7878 (2021). arXiv:2104.0095. [cs.LG]
[82] Carlos Riquelme, Joan Puigcerver, Basil Mustafa. Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. 2021. Scaling Vision with Sparse Mixture of Experts. In Advances in Neural Information Processing Systems (NeurIPS), arXiv:2106.0555 05997 [cs.CV]
[83] Clemens Rosenbaum, Tim Klinger, and Matthew Riemer. 2018. Routing networks: Adaptive selection of non-linear functions for multi-task learning. In International Conference on Learning Representations (ICLR). arXiv:1711 111::1711.01299 1.02339 [cs.LG]
[84] Christoph Schär, Oliver Fuhrer, Andrea Arteaga, Nikolina Ban, Christophe Charpilloz, Salvatore Di Giro- lamo, Laureline Hentgen, Torsten Hoefler, Xavier Lapillonne. David Leutwyler, et al. 2020. Kilometer- scale climate models: Prospects and challenges. Bulletin of the American Meteorological Society 101, 5 (2020).
[85] MG Schultz, Clara Betancourt, Bing Gong, Felix Kleinert, Michael Langguth, LH Leufen, Amirpasha Mozaffari, and Scarlet Stadtler. 2021. Can deep learning beat numerical weather prediction? Philosophical Transactions of the Royal Society 379, 2194 (2021).
[86] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. 2017. Outragcously large neural networks: The sparsely-gated mixture-of- of-experts layer. In International Conference on Learning Representations (ICLR). T.XXv:1701..66 [cs.LG]
[87] Xingjian Shi, Zhourong Chen, Hao Wang, Yan Yeung, Wai-Kin Wong, and Wang-chun Woo. 2015. Convolutional LSTM network: A machine learning approach for precipitation nowcasting. Advances in neural information processing systems (NeurIPS). arXiv: M..Xv:15066042444 [cs.CV]
[88] Xingjian Shi, Zhihan Gao, Leonard Lausen, Hao Wang. Dit- Yan Yeung. Wai-kin Wong, and Wang- Woo. 2017. Deep learning for precipitation nowcasting: A benchmark and a new model. Advances in neural information processing systems (NewriPS). ar arxi 7044535888 [cs.CV]
[89] Casper Kaae Sønderby, Lasse Espeholt, Jonathan Heek, Mostafa Dehghani, Avital Oliver, Tim Salimans, Shreya Agrawal, Jason Hickey, and Nal Kalchhrenner. 2020. MetNet: A neural weather model for precipitation forecasting. (2020). 20X.:20003.12144 [cs.LG]
[90] Suraj Srinivas, Akshayvarun Subramanya, and R Venkatesh Babu. 2017. Training sparse neu- ral networks. In Conference on computer vision and pattern recognition workshops (CVPRW). arXiv:1611.00 1.6694 [cs.CV]
[91] Yaniv Taigman, Ming Yang, Mare' Aurelio Ranzato, and Lior Wolf. 2014. DeepFace: Closing the gap to human-level performance in face verification. In Conference on Computer Vision and Pattern Recognition (CVPR).
[92] Zoltan Toth and Eugenia Kalnay. 1993, Ensemble forecasting at NMC: The generation of perturbations. Bulletin of the American Meteorological Society 74, 12 (1993).
15

---

[93] Ashish Vaswani, Prajit Ramachandran, Aravind Srinivas, Niki Parmar, Blake Hechtman, and Jonathon Shlens. 2021. Scaling local sell-attention for parameter efficient visual backbones. In Conference on Computer Vision and Puttern Recognition (CVPR). arXiv:2103.1273: [cs.CV]
[94] Vasilis Vryniotis. TorchNision's
Models Us- hetps://ppttorch.orr/// how-to- train- of www.aattteessssstteeeeeeeeee-s-e using- corchvision-latesee pawtttttssttpimiieeeette
2021.
How
to
Train
ofate-of-the-Art
ing
Latest
Primitives.
[95] Fei Wang. Mengqing Jiang. Chen Qian, Shuo Yang. Cheng Li, Honggang Zhang, Xiaogang Wang, and Xiaoou Tang. 2017. Residual attention network for image classification. In Conference on computer vision and pattern recognition (CVPR). MXXi:11704006000 [es.CV]
[96] Wenhai Wang. Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao. 2021. Pyramid vision transformer: A versatile backbone for dense prediction without convolutions. In International Conference Computer Vision (ICCV). arXiv:2102.121 12122 [cs.CV]
[97] Wenhai Wang, Enze Xie, Xiang Li, Deng- Ping Fan, Kaitao Song. Ding Liang. Tong Lu, Ping Luo, and Ling Shao. 2022. PVT v2: Improved baselines with Pyramid Vision Transformer. Computational Visual Media (2022). arX1v:110.113997 [cs.CV]
[98] Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He. 2018. Non-local neural networks. In Conference on computer vision and pattern recognition (CVPR). arXiv:1711.0997 [cs.CV]
199] Xin Wang, Fisher Yu, Lisa Dunlap, Yi-An Ma, Ruth Wang, Azalia Mirhoseini, Trevor Darrell, and Joseph E Gonzalez. 2020. Deep mixture of experts shallow embedding. Uncertainty in Artificial Intelligence. arXiv: 10ti:1806.01533 [cs
[100] Sanghyun Woo, Jongchan Park, Joon- Young Lee, and In Kweon. 2018. CBAM: Convolutional block attention module. European conference on computer vision (ECCV). arXiv:1807 www.18807065222 06521 [cs.CV]
[101] World Metcorological Organization. 2021. Guidelines on Ensemble Prediction System Postprocessing. 2021 edition. Technical Report. World Metcorological Organization.
[102] Brandon Yang, Gabriel Bender. Quoc Le, and Jiquan Ngiam. 2019. CondConv: Conditionally
parameterized convolutions for efficient inference. In Advances in Neural Information Processing Systems [cs.CV]
(NeurIPS). arXiv:1904.044
[103] Sergey Zagoruyko and Nikos Komodakis. 2016. Wide residual networks. In British Machine Vision Conference (BMVC). 1rXiv:1605.07111 [cs.CV]
[104) Songyang Zhang, Xuming He, and Shipeng Yan. 2019. LatentGNN: Learning efficient non- local relations for visual recognition. In International Conference on Machine Learning (ICML). 1xXiv:195.11164 [cs.CV]
[105] Kaili Zhao, Wen-Sheng Chu, and Honggang Zhang. 2016. Deep region and multi-label learning for facial action unit detection. In Conference on computer vision and pottern recognition (CVPR).
[106] Shuai Zheng, Sadeep Jayasumana, Bernardino Comera-Paredes Vibhav Vineet, Zhizhong Su, Dalong Du. Chang Huang, and Philip HS Tor. 2015. Conditional random fields as recurrent neural networks. In International Conference Computer Vision (ICCV).
[107] Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus, 2022. Designing effective sparse expert models, 2xXi:22202.08900 [cs.CL]
16

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] We discuss the key assumption
that all data must have similar location-dependenee for SMoEs to be applicable in $1. We discuss other assumptions in $2 and implementation details in $2.4.
(c) Did you discuss any potential negative societal impacts of your work? [Yes] We discuss this in $4.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes] 2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main exper- imental results (either in the supplemental material or URL)? [Yes] We include code and instructions in the supplementary material.
(b) Did you specify all the training details (e.g.. data splits, hyperparameters, how they were chosen)? [Yes] 53 discusses general training details and SA provides additional details.
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes] All results in $3 were run multiple times, and we report the mean and standard deviation.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] This is discussed in $3.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [Yes] Most datasets we use are publicly available. Our heat diffusion dataset is synthetic and can be generated by our included
code. The supplementary material includes the license for our code.
(c) Did you include any new assets either in the supplemental material or as URL? [Yes] We include code to implement SMOEs and run experiments.
(d) Did you discuss whether and how consent was obtained from people whose data you'l re using/curating? [No] Some image classification datasets we present results on, notably ImageNet, contain images of people. We have not obtained consent from them. However, these are standard, widely-used benchmark datasets for image classification, and we do not include any results that disclose personal information. We do not introduce or curate any data that contains personal information.
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [No] We do not introduce or curate any new data that contains personal information offensive content. We do use ImageNet standard benchmark, which may contain personally identifiable information (pictures of people), but do not report any results that include personally identifiable information or offensive content.
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
17


Reviews:
[{"review_id": "kNw-kfP4bc", "reviewer": "Reviewer_cLKX", "summary": "This paper proposes a new way of tailoring the processing of subparts of the input data based on their localisation, in the case where all samples from a given dataset share the same spatial structure. This can be of use in cases where treating the localisation as a mere input feature is not enough, but appeals for more drastic measures breaking the spatial invariance properties of some classic networks such as CNN.\n\nThe method presented relies on learning different experts, each being specialized in processing certain parts of the input data. At each localisation, the choice of experts to use is done by looking at a \"routing table\", scoring each expert for each localisation, regardless of the input sample. There is only one routing table for the whole dataset: the routing is independent of the input sample. It is learnt via a newly introduced method treating the routing as a classification task, with labels derived from the loss of the task at hand.\n\nThe method is shown to compare favorably to previous baselines on a task created for this paper where the \"true parameters\" of the model generating the sythetic data are known, so as to perform an in-depth ablation study highlighting the impact of each design decision. On a weather-prediction benchmark using real-world data, the method is shown to be state-of the art, and can also bring a slight improvement to a ResNet 50 on Imagenet.", "questions": "* On the heat diffusion dataset, did you try to compare yourselves to some baseline CNN architectures having the same number of parameters  as your SMoE model ? (In Figure 3, I only see a line writing 146 for the CNN vs 27+12288 for SMoE)\n* Is it possible to compare SMoE to the other methods on WeatherBench ?\n* The quantile of exactly which distribution is used to create the labels in the training of your routing tensor ? \n* In the practical implementation part, you said that all experts at all points are evaluated. If I am understanding right, this means that in the last layer of your model, instead of using the $\\texttt{scatter}$ and $\\texttt{gather}$ functions of Pytorch, you could directly evaluate the loss for all experts. Would it be useful to make use of this signal to train your routing system ?\n\n* (there is a small typo line 311: \"Each SMOE ~~has~~ selects\")", "limitations": "Yes, the authors adressed the limitations and potential negative societal impact in their discussion paragraph.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "$\\textbf{Strengths}$:\n\nThe paper is clearly motivated and well written overall. The method introduced to train both the routing system and the experts at the same time is original and kept simple, the relevance of each component is thoroughly tested, and the robustness of the model to changes in variables unknown by the user is challenged in a controlled setting on synthetic data. The method is shown to work in practice on real-world data.\n\n$\\textbf{Weaknesses}$:\n\n* The explanation of the procedure used to create the labels for the classification task to train the routing tensor is a bit unclear to me: the quantile of exactly which distribution is used to create the labels ? \n* In WeatherBench (Table 3), SMoE is only compared to one method whereas other models leveraging the position where tested and shown to be on par with SMoE performances on the heat diffusion data (Figure 3.a) and simple Image datasets (Table D.1 and D.2). It would be interesting to know how those previously introduced methods fare on a case where it seems breaking the spatial invariance makes sense.", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "RD7qFn3UNBZ", "reviewer": "Reviewer_wHKK", "summary": "Authors proposed a spatial mixture-of-expert layer for CNNs to take into account of spatial dependence in data or features.", "questions": "Computational costs should be analyzed and compared.", "limitations": "Okay", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: to incorporate local dependence or locality in data and features is always beneficial. Yes, most CNNs do not explicitly consider this or have mechanisms to embed this feature into the learning process.\n\nWeaknesses: authors create an impression that the proposed scheme is the only one that consider locality in data. There are always many mechanisms to incorporate local dependences in CNNs such as graphs, CRF or MRF, in segmentation and recognition. The proposed scheme seems only marginally improve performances if any, while added computational costs may outweigh the benefit.\n\nTitle is misleading - word \"Layer\" should be added.", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "E0GnP6wATpL", "reviewer": "Reviewer_vuZ5", "summary": "This paper presents a location-based mixture of experts for use in applications where modeling depends on the image grid location, such as weather forecasting.  Experts are selected using a map of size (E,H,W), where E is the size of experts pool, so that one set of top-k experts are learned for each grid location.  In order to encourage the system to select appropriate experts, an auxiliary loss is added that classifies whether the selected experts are \"correct\" or \"incorrect\" based on quantile of the gradient magnitude; experts with large error gradient are assigned a uniform classification target among the other experts, encouraging the system to find a new one.  The largest gradients are also scaled down in their bprop application, with the intuition that they will be applied to experts that may no longer be selected.  The system is verified on a toy dataset constructed from synthesized heat diffusion, and evaluated on WeatherBench and Imagenet, obtaining significant improvements in weather forecasting.\n", "questions": "Additional questions:\n\nl.194:  The claim that sparse selection results in the same number of FLOPS as a standard conv layer with same number of filters, is not quite right.  For each location, the additional patch context must also be included.  In the worst case, each neighboring location selects different experts in a 3x3 checkerboard pattern, in which case FLOPS required would be that of a standard convnet with the same number of filters times the spatial kernel size (e.g. 9).  And even this depends on having a sparse implementation when there are more than 9 experts to select from.\n\nFig 3:  do you have any ideas on why LRLCN-ND performs worse than SMoE on the toy example? it's a very similar architecture for 1 output channel, the only difference is that it uses per-location softmax mixture of the 3 kernels, rather than top1 max.  will adding RC loss fix this?\n\nIf land and water are in different locations between images, is it possible to feed the land/water map in as a binary map and gate based on this, rather than location?\n", "limitations": "-", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The idea of incorporating location dependence for some applications, like weather forecasting, is sound, and well-suggested by the results in this paper.  Recovery of near-exact diffusion kernels in the toy data is good to see.  However, there are some additional baseline comparisons I think are missing, in particular, LCN, LRLCN and resnet with wider layers to match the number of parameters of SMoE on WeatherBench.  SMoE with a fixed set of random 128/256 experts at each location selected, as in the init, without learning the gating map, would also be interesting, as improvement over this would show the effect of learning the map in WeatherBench.  Measuring on an additional relevant application, such as face landmarks or another of the ones mentioned in the intro, would also strengthen the demonstrated reach of the method.\n\nThere are also some details around the order of the channels (experts) selected, and if that matters.  Are the top ones per location selected in index order, or sorted gate score order?  When a filter changes, and an expert i is replaced with expert j, are experts with indices between i and j will shifted in the output representation happens?  This would look like a sudden shift permutation to the next layer.  And if experts are selected in score order, then as scores change, input channels to the next layer will swap when scores change during training.\n\nAs an aside, I think it's debatable whether the system described is a mixture of experts.  It seems more to be a channel subselection, since the outputs are concatenated instead of mixed, and subselection is not dependent on the same input as the layers.  This is a very minor point, though, it seems close enough to me to be an accurate description.\n\nOverall, this is an interesting idea, and one that clearly seems effective at incorporating effects of spatial differences, such as land vs water in weather prediction.  A few comparisons appear to be missing in the experiments (described above), and while the toy heat diffusion experiment and existing ablations serve to profile the behavior of the system, I still would have liked to see more discussion on the effects of discontinuous subselection to downstream layers.\n", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "4hgYkyNguMy", "reviewer": "Reviewer_Y3eL", "summary": "The authors propose a convolutional neural architecture that learns parameters specialized towards different locations in the input image. This can be seen as a middle way between a ConvNet, where all parameters are shared across spatial locations and fully-connected networks, where all parameters are specific for their location. Specifically, the authors propose a learnable matrix that assigns each spatial location to a subset of convolutional filters. To learn this matrix, the authors propose a routing classification loss that aims to change assignments for locations with large gradients. The method is then evaluated on a small set of specialized datasets, where it performs well.", "questions": "Would the authors be able to highlight a series of key tasks where such a model could provide impact?", "limitations": "As far as I can see, the authors have not addressed potential limitations. However, I do not see any clear issues.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Making neural models larger has been an effective way to improve model performance. Making those large models sparse has been shown to be a successful way to combine large capacity models with fast (theoretical) inference and training speeds. As such the paper at hand is of interest to the research community. \n\nOverall the paper is well presented and the method is clearly explained. \n\nOne concern is that the authors make a very strong assumption in that SMOE relies on the spatial structure to be similar across all samples in a dataset. This leads to very limited applications and no generalization across datasets. Given this very large limitation, it would have been great if the authors could extend their motivation and highlight a series of key tasks where such a model could provide impact. ", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "2vnVeC7K6p7", "reviewer": "Reviewer_811C", "summary": "The paper presents a method based on spatial MoE for location-dependent learning. Unlike the translation-equivariant network (i.e., CNN) that shares the expert (convolution kernel) locally, the proposed method assigns the convolution kernel conditional on the location. Specifically, the work proposes a novel gating function to select an expert for each location and use the expert to extract information for the location. The paper validates the effectiveness of their method in multiple tasks and shows the SOTA results on medium-range weather prediction.", "questions": "Does the gating learn the translation-equivariance? It might be interesting to look at. The reason I ask is that, for location-dependent earth data, I still believe that the translation-equivariance could be helpful since patterns are repetitive. And as shown in Fig.3, I saw that there is a local pattern (e.g., repetitive local patches) in the expert selection.", "limitations": "The author addressed the limitations. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Originality:\nThe paper presents a novel viewpoint on utilizing deep networks. For some tasks, it's beneficial to violate the translation-equivariance. This is very interesting to the research community. Furthermore, the proposed method is also novel. The proposed gating function and classification loss obviate the need for a variety of regularizers.\n\nQuality:\nThe method supports the claims well with multiple experiments. The only concern I have is the missing analysis of mode collapse -- a typical pitfall of MoE. For completeness, it would be better to give more analysis -- for example, is there any mode collapse and any insight/discussion on why the proposed method can avoid it?\n\nClarity:\nThe paper is well-written in general. I only have minor comments.\n- The motivation behind the proposed tricks is missing. The paper didn't show the logic behind the proposed tricks. For example, why is the expert error damping essential, and the logic behind its design?\n\nSignificance:\nThe method is novel. The results are also strong and very interesting -- by violating translation-equivariance, the method shows improved performance.", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}]

Author rebuttal:

Decision: Accept