Spatial Mixture-of-Experts
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Many data have an underlying dependence on spatial location; it may be weatheron the Earth, a simulation on a mesh, or a registered image.
Yet this featureis rarely taken advantage of, and violates common assumptions made by manyneural network layers, such as translation equivariance.
Further, many works thatdo incorporate locality fail to capture ﬁne-grained structure.
To address this, weintroduce the Spatial Mixture-of-Experts (SM O E) layer, a sparsely-gated layer thatlearns spatial structure in the input domain and routes experts at a ﬁne-grainedlevel to utilize it.
We also develop new techniques to train SM O Es, including aself-supervised routing loss and damping expert errors.
Finally, we show strongresults for SM O Es on numerous tasks, and set new state-of-the-art results formedium-range weather prediction
Introduction
Many datasets exhibit an underlying, location-based structure, where the value at a point dependson where that point is.
For example, weather predictions [71] depend on their location on Earth;many scientiﬁc simulations are parameterized on an underlying mesh [8]; data (e.g., faces) may bespeciﬁcally aligned [84]; 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 inthe future.
Numerous architectures have been successfully applied to such tasks.
Convolutional neuralnetworks (CNNs) [73] and transformers [67] show promise for medium-range weather prediction.
Locally-connected networks (LCNs), which use independent, rather than shared, ﬁlters at eachpoint, have been applied to weather post-processing [33], face recognition [40, 70, 84], and othertasks [15, 64], speciﬁcally to learn local features.
Low-rank local connectivity (LRLCN) [27] relaxesthe translation equivariance of CNNs while requiring fewer parameters than LCNs.
Other approaches,such as CoordConv [58], add an additional inductive bias by providing explicit input coordinates.
However, for tasks on data with location-based structure, prior approaches suffer from variouslimitations.
Convolution assumes that data is translation equivariant, which does not hold for suchtasks [40, 84].
Approaches like LCNs require many parameters.
Many architectures have beendesigned for classiﬁcation tasks and fail to perform well on regression because they opererate at toocoarse granularity and are unable to capture key details.
This limits their applicability to importanttasks, such as medium-range weather prediction or climate simulations.
Indeed, on a simple heatdiffusion task with location-dependent diffusivities (see §3.1), many approaches do not learn thelocation dependence at all and instead converge to an “average” diffusivity.
To address this, we introduce a novel neural network layer, the Spatial Mixture-of-Experts (SM O E)layer (§2).
An SM O E uses a learned gating function to sparsely select and route from a shared set ofexperts to each spatial location (e.g., pixel) in an input sample.
This enables experts to specialize tothe unique characteristics of different “regions” within the data and easy interpretability by analyzingthe gate.
SM O Es require the assumption that all samples in the dataset have a similar underlyingspatial structure, but this often holds (at least approximately) for many datasets, such as weather,where each example is on the same latitude/longitude grid. For the SM O E gating function, we
I ndu c
t i v e
b
i a s
Parameters Fully-connected
Convolution
CoordConv
Locally connected
Spatial Mixture-of-Experts (this work)
T il e d L C N
Low-rank locally connected
SMoEwith same experts at every location
SMoEwith different experts at every location (Not to scale)
Convolution:
Locally-connected:
Low-rank locally-connected:
Spatial Mixture-of-Experts:
Filters
Filters ! ! ! " !
# ! $ !
" " " # " $ " % " &
Basis filters
Learned combiners ! !
and " !
Experts Learned gate
SMoEslearn fine-grained routing to locations
Mixture-of-Experts:
Spatial Mixture-of-Experts:
Experts
G a t e
G a t e
Experts • Coarse-grain routing • >Capacity for constant compute • Classification • Fine-grain routing • Learn location- dependent functions • Regression
Inputs
Input introduce tensor routing , a simple and cheap routing function that effectively learns this spatialstructure (§2.1).
We also ﬁnd that end-to-end training approaches, as in standard mixture-of-experts(MoEs) [75, 79], do not capture location-dependence well.
Instead, we introduce a separate routingclassiﬁcation loss (§2.2) to train the gate, a self-supervised loss which provides a better learningsignal to the gate about incorrect routings.
We also introduce expert error damping (§2.3), whichlimits expert updates from misroutings, and further improves performance.
Both methods rely onextracting information from the error signal (gradient of the layer output) which would otherwise beuninformative.
Figure 1 provides an overview of SM O Es and a qualitative comparison to importantrelated work.
With these methods, SM O Es are able to learn ﬁne-grain, location-dependent whichother architectures miss, and consequentially deliver much better performance.
SM O Es are related to and inspired by prior work in MoEs (e.g., [75, 79, 93]), but there are keydifferences.
Existing MoEs use coarse-grained routing at the sample [91], token [79], or patch [75]level, and experts produce coarse output (e.g., entire samples or channels) while SM O Es route ata ﬁne-grained, per-pixel level.
Fine-grained routing is important for enabling experts to specializeto ﬁne-scale information, which may be crucial [67, 77].
Such MoEs also typically aim to increasemodel capacity while keeping compute costs constant, whereas the goal of SM O Es is to capturespatial dependencies.
Other work has aimed to incorporate spatial dependence, such as LRLCNs [27],which learn content- and location-dependent weights to combine a set of basis ﬁlters.
However,we found that these prior methods failed to capture the ﬁne-grained features necessary for goodperformance on regression tasks (§3.1).
We conduct experiments on several benchmark datasets with SM O Es and conduct extensive ablationstudies of SM O E design decisions (§3).
Using a simple heat diffusion dataset, we showcase thelimitations of other models and the power of SM O Es in a controlled situation.
We then apply
SM O Es to medium-range weather prediction and outperform the state-of-the-art [73] on WeatherBench [71] when not using additional data.
Finally, we show that SM O Es can also be applied toimage classiﬁcation tasks, where we match or outperform LRLCNs while using fewer parameters.
1.1 Related Work
Mixture-of-Experts and conditional computation.
While MoEs have long been used [12, 43,46], since the advent of deep learning, there has been much work in applying MoEs, conditionalcomputation, and dynamic routing to DNNs [1, 4, 5, 7, 9–11, 16–18, 20, 22, 24, 26, 29, 30, 34, 45,52, 53, 63, 68, 69, 75, 76, 79, 91, 93, 98].
Often, the goal is to increase model capacity without acorresponding increase in compute costs and use coarse-grained routing.
Notably, MoEs route atthe sample (e.g., [22, 26, 63, 76]), token (e.g., [29, 52, 79]), or patch (e.g., [75]) level; are typicallytrained with extensive auxiliary losses; and often target classiﬁcation tasks.
This coarse routing, inparticular, means there is limited opportunity for experts to specialize.
MoEs speciﬁcally for visiontasks also typically operate at a sample-level (e.g., [1, 4, 34, 91, 93]) and use experts to specializeﬁlter or channel choices.
In contrast, our SM O Es induce an inductive bias via ﬁne-grained, locationdependent routing for problems with such underlying structure; typically do not need auxiliary lossesbeyond the routing classiﬁcation loss; and work well for dense regression tasks. Local spatial structure. Locally connected networks (i.e., convolution ﬁlters without weight sharing)were historically successful in vision tasks such as facial recognition [15, 40, 70, 84]. Further
Input
[Figure] ℰ " #
Gate $ , tensor routing % ∼ ℰ ×"×# ( §2.1 )
[Figure]
top ! sparse routing MSE loss + = 1 . / 0 − 2 " Error signal 3+ 30 = .
(0 − 2)
Routing Classification Loss ( §2.2 ) Identify incorrect routings with error signal magnitude ✓ ✘ ✓ ✓ Construct labels 10 1 ! 01 0 !
Compute Loss
Binary cross- entropy Expert Error Damping ( §2.3 )
Incorrect routings ✓ ✘ ✓ ✓ Damp error signals for incorrect routings ! Forward Backward
[Figure]
channels ❶ ❷ ❸ work incorporated periodic or tiled weight sharing [32, 64, 96].
However, more recent work oftenexclusively uses convolution, which has been found to perform better [65] while requiring signiﬁcantlyfewer parameters.
CoordConv [58] explicitly provides pixel coordinates as input.
Because of theirstructure, a single convolution or CoordConv layer cannot learn location-dependent sparse activationsas SM O Es do, and multi-layer networks give limited improvements in practice.
Low-rank locallyconnected networks [27] learn to combine a small number of basis ﬁlters in a location- and inputdependent manner.
These combiners lack sparsity and apply all basis ﬁlters at each point, and aresoftmax-normalized, which can result in a few ﬁlters dominating, limiting diversity.
Separate work hasused attention to provide position-dependent scaling or integrate context [31, 44, 57, 85, 87, 92].
Thishas culminated in vision transformers [23] and variants [25, 55, 60, 88, 89], which use specializedattention mechanisms to integrate spatial information.
Other work uses Markov [54, 59] or conditionalrandom ﬁelds [13, 37, 97], or graph-based methods [14, 90, 95], to learn long-range dependencies.
Sparsity.
Our use of sparse routing resembles work on sparsity in DNNs in general [39].
Manyapproaches learn a ﬁne-grained mask to identify which weights to keep or prune (e.g., [56, 83]),while other approaches sparsify activations (e.g., [3, 41, 49, 61, 62]).
These works use sparsity toimprove runtime performance, typically at inference time.
Weather prediction.
Medium-range weather prediction, or forecasting roughly three to seven days inadvance [6], is of broad societal import [50].
There has been much interest in applying deep learningto this task [78], and WeatherBench [71] serves as a community benchmark.
Deep learning has alsobeen successfully applied to the related tasks of ensemble post-processing [33, 72] and now-casting(forecasting a few hours in advance) [2, 28, 74, 80–82].
Standard CNNs are currently state-of-theart on WeatherBench [73].
However, weather data has extensive ﬁne-grained, location-dependentstructure which take advantage of with SM O Es for improved performance.
Spatial Mixture of Experts
We now introduce the Spatial Mixture-of-Experts (SM O E) layer.
An SM O E uses a learned gatingfunction (§2.1) to select speciﬁc experts from a shared set to apply at each point (e.g., pixel) in aninput.
The gate learns the underlying spatial structure of the data, and routes speciﬁc experts todifferent “regions”, allowing them to specialize.
An SM O E is predicated on the assumption that thespatial structure is similar across all samples in a dataset.
We also assume data are on a Cartesianmesh (i.e., grid), although this could be generalized.
To train SM O Es, we introduce a self-supervisedrouting classiﬁcation loss (§2.2) and expert error damping (§2.3), which we found essential forachieving the best performance.
Figure 2 provides an overview of SM O Es and their training.
We ﬁrst deﬁne the SM O E layer in full generality, then discuss the particular implementation we use.(See §3.1 for ablations.)
Let x 2 R C ⇥ H ⇥ W be an input sample (we assume 2D data for simplicity).
The SM O E layer consists of a set E of experts, of which E  |E| will be selected at each point, and agating function G : R C ⇥ H ⇥ W !
R |E| ⇥ H ⇥ W .
The gating function has exactly E nonzeros at eachspatial point ( H ⇥ W ), which correspond to the selected experts at that point.
Each expert e 2 E isapplied at the points in x where it has been selected, and may include additional context from x (e.g.,surrounding points if e is convolution), and produces an output of size F .
The outputs of each expertat a point are then concatenated to form the output channel dimension (of size E · F ) of the SM O E.
More precisely, let gather( · ) select only the non-zero entries of its input.
Then an SM O E is:
[Equation]
where · is element-wise multiplication, [ e ; . .
stacks tensors, and y 2 R EF ⇥ H ⇥ W .
When E ⌧ |E| , 122 the gating function induces signiﬁcant sparsity; this allows an SM O E to compute experts only at thepoints where they are to be applied, avoiding most computation.
This formulation yields a weighted SM O E, where the expert outputs are scaled by the gating function.
An alternative, which is slightly more efﬁcient and can be more readily interpretable, is an unweighted
SM O E, in which case experts are not scaled, and the gather only uses G ( x ) to select experts.
In this work, we focus on a simple, yet powerful, SM O E layer using convolution ﬁlters as experts:
[Equation]
c , where ? is cross-correlation, i 2 [ |E| ] , and we have elided the gather for simplicity.
2.1 Gating Functions
The gating function G in an SM O E is critical for learning the underlying location-dependenceof the data.
While G can be any function, a good gating should be relatively cheap.
We followprior work on MoEs (e.g., [75, 79]) and use top- E routing to select experts and sparsify the gate:
G ( x ) =
top E ( g ( x )) , where g ( x ) 2 R |E| ⇥ H ⇥ W is a learnable gating layer and top E selects the Elargest entries along the ﬁrst (expert) dimension and sets the remaining entries to 0.
However, wefound 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 MLPs [75, 79], but wefound this did not perform well.
Convolution or CoordConv [58] also did not perform well.
Instead,we ﬁnd that a simple tensor routing gating layer worked best.
With this, g ( x ) =
D 2 R |E| ⇥ H ⇥ W ,where D is a learnable tensor that directly encodes the underlying location dependence and routingstructure without depending on the input.
Further, D uses one parameter per expert per location andrequires no computation beyond optimizer updates, making it an efﬁcient choice.
We initialize D using a uniform distribution over [  3 |E| / EF , 3 |E| / EF ] , which corresponds to a Kaiming 143 uniform initialization with fan-in EF / |E| [35].
However, in many practical cases, some data about 144 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 gatingfunction to beneﬁt from prior knowledge while still being able to adjust the routing.
Additionally, anetwork may contain many SM O E layers, in which case we can share D between layers that have thesame 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 routingpolicies and avoid mode collapse (e.g., [9, 75, 79]).
When training with the routing classiﬁcation losswe propose, we did not ﬁnd additional auxiliary losses to be necessary.
2.2 Training and the Routing Classiﬁcation 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 SM O Es and tensor routing, particularly on regression tasks, as the gate did not learn well: it rarelychanged its routing decisions from initialization.
We hypothesize that this is due to a “mismatch” inthe 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, whereasselecting an expert requires a threshold (see §B).
To address this, we train the gate with a separate,self-supervised loss function, the routing classiﬁcation (RC) loss.
The key idea is to consider routingas a dense, multi-label classiﬁcation task: selecting the “correct” experts at each point (cf. semanticsegmentation).
The RC loss does exactly this, and trains the gate by constructing appropriate labels.
This also helps avoid mode collapse, where only a 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 ofexperts at each point.
However, such information is not directly available.
Instead, we use the errorsignal into the SM O E layer (i.e., the gradient of the layer output w.r.t.
the loss) as a proxy, and saythat 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 errorloss L (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:
[Equation]
Hence, the error signal is simply the (scaled, signed) error of the predictions.
While this exact relationceases to hold as backpropagation proceeds, the intuition behind the error signal magnitude remains.
We now deﬁne the routing classiﬁcation loss.
Given the error signal " into an SM O E layer, and anerror quantile q (a hyperparameter), we say that selecting an expert at a point was incorrect if " atthat expert and point is greater than the q th quantile of " , and correct otherwise.
We use quantiles asthey are independent of " ’s scale, which may be hard to determine and change during training.
Wethen construct the labels for each point as follows:
Unselected experts start with label 0 .
A correctlyselected expert has label 1 .
Finally, if an expert was incorrectly selected, its label is 0 and we add 179/ ( |E|  E ) to the label value of each unselected expert (note |E|  E is the number of unselected 180 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 SM O E gate to be trained directly based on whether it routed correctly,experts that were incorrectly routed to still perform gradient updates based on this routing.
Thisresults in experts updating to improve their performance for locations they may not be applied at infuture 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 incorrectupdates.
We ﬁnd that this can improve SM O E performance and reduce convergence time.
Expert error damping is similar to the RC loss, and we classify incorrect routings in the same way.
We then scale the error signal into the experts by a constant factor at each point where the routingwas 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 SM O E layer, primarily focusing on the simple SM O Ewith convolution ﬁlters we use.
In an ideal implementation, the ﬂop cost of an SM O E is thecost of the gating function plus the cost of applying the selected experts.
When using a tensorrouting gating layer, there are no ﬂops in the gating layer.
The ﬂops from applying the selectedconvolutional experts is equivalent to a standard convolutional layer with a number of ﬁlters equalto the number of selected experts.
Hence, SM O Es are quite efﬁcient ﬂop-wise.
However, recentwork has shown that data movement costs are also critical to performance [42].
Because of this,we do not expect even well-optimized implementations to match the runtime of a convolution layerdue to the sparse routing, limited locality in accessing expert ﬁlters, and other operations, althoughwork on hardware-accelerated sparsity [19] should offer beneﬁts.
Despite this, during inference someadditional optimizations may be available because the tensor routing gate does not depend on theinput and computations could be reordered to maximize locality. Unfortunately, efﬁciently implementing the irregular access patterns in an SM O E is challengingin standard Python deep learning frameworks, and likely requires custom compute kernels to bewritten. Instead we opt for a naïve implementation in PyTorch [66] where we apply all experts at [Figure]
[Table]
[Figure]
[Table] [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure] all points and then use gather and scatter operations to efﬁciently implement sparse routing. Experts are concatenated in sorted routing score order (i.e., the order given by G ( x ) ); while this canresult in channel orders changing during training, we ﬁnd it has little impact (see §3.2).
We ﬁnd thisimplementation sufﬁcient even for large-scale tasks.
Experiments
We now present experimental results using SM O Es and a variety of baseline models, as well asablations of different SM O E components.
First, we describe a simple location-dependent heatdiffusion dataset, which we use to study the ability of different architectures to learn locationdependence in a controlled environment (§3.1).
We then present results on the WeatherBenchmedium-range weather forecasting challenge, where SM O Es achieve state-of-the-art accuracy (§3.2).
Lastly, we show results on several image classiﬁcation tasks, illustrating the breadth of SM O Eapplicability even in situations without strict location-dependence (§3.3).
All results were run using PyTorch [66] version 1.11 on a large cluster with 16 GB V100 GPUs.
Wesummarize training details throughout this section, and provide full details in §A. We use existing PyTorch implementations when available, and implemented methods ourselves otherwise.
Unlessnoted, all SM O E results used 3 ⇥ 3 convolution kernel experts, unweighted tensor routing gates, RC loss, expert error damping, error quantile q = 0 .
7 , and damping factor 0 .
We report the meanand standard deviation (mean ± std ) over ten runs for all experiments, except for WeatherBench and ImageNet, which used three runs.
In total, we used about 24k GPU hours for experiments.
3.1 Location-Dependent Heat Diffusion
To study location-dependence in a controlled manner, we generated a heat diffusion dataset withlocation-dependent diffusivities.
The task is, given the current heat at each point, to predict what theheat will be at the next timestep.
Because we know the exact diffusivities and their distribution in thedata, it is easy to identify how well a model has learned this task.
We ﬁrst describe the dataset and itsgeneration in detail, then discuss results and ablation studies on SM O E components.
Dataset.
The dataset consists of a region map, where each location is assigned a type, and each typecorresponds to a different heat diffusivity.
The region map and diffusivity assignment are then ﬁxedfor the dataset.
We generate the region map randomly using a simple preferential attachment method.
This deﬁnes the location-dependence that we wish to learn.
To generate samples, we ﬁrst randomlydistribute drops of heat across the domain, then apply ﬁve-point diffusion stencils at each point, usingthe 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 thenconsists of the generated timesteps for many different starting states.
If a model is able to learn the diffusion stencils and the location-diffusivity correspondence, itcan exactly predict the next timestep.
Further, the diffusivity stencils are also simple, and exactlycorrespond to a 3 ⇥ 3 convolution kernel with no nonlinearity. [Table]
[Table] The particular dataset we use consists of 100 , 000 64 ⇥ 64 samples, with 1 , 000 initial states evolvedfor 100 timesteps each.
Adding more samples did not signiﬁcantly change results.
There are threeregion types, with diffusivities 0 .
25 , 0 . 025 , and 0 .
Figure 3 (center) shows the region map anddiffusion stencils.
We report results using “% within 1%”, the percentage of locations in a sample thatare within 1% relative error of the true value, as this is more interpretable than a mean-square error.
Results.
Figure 3 (left) shows results on the heat diffusion dataset for SM O Es and a number of baselines: CNNs, LCNs, fully-connected (FC) layers, CoordConv [58], CondConv [93], LRLCNs [27],vision transformers (ViT) [23], and vision MoEs (V-MoE) [75].
For the LCN and FC layers, weuse only a single layer as additional ones showed no beneﬁt.
Convolution, CoordConv, CondConv,and LRLCN are the best network found among a set with up to three layers, 3 ⇥ 3 kernels, 12ﬁlters per layer, batchnorm, and ReLU activations.
ViT and V-MoE use one transformer block witha patch size of 4 ⇥ 4 , an embedding dimension of 128 , and four heads.
LRLCNs use three basisﬁlters and an input-dependent combiner.
We also tried unshared combining weights with no inputdependence (LRLCN-ND).
V-MoEs select one expert from a set of three.
Our SM O Es use a singlelayer with three experts and select E = 1 expert per point.
All models were trained with batch size32 , Adam [47] with a learning rate of 0 .
001 (decayed by 10 ⇥ after no validation improvement for 15epochs), and early stopping after no improvement for 30 epochs.
Additional hyperparameter tuningdid not signiﬁcantly improve results.
We report SM O E parameters as expert + gate parameters.
SM O Es achieve perfect performance on this dataset.
Further, by examining the learned routingand experts (Fig. 3 (right)), we can see that it has indeed correctly learned the diffusion stencilsand location-dependence.
LCNs also achieve this, but require 3 ⇥ more parameters, and require110 epochs to converge (versus 8 for SM O Es).
Fully-connected layers failed to learn the datawell, likely due to the challenge of optimizing so many parameters.
Other methods all converge tobetween 91 and 94% within 1%.
Examining their predictions and weights, we observe that theydo not appear to have learned the location-dependence of the diffusivities, and instead convergedto predicting with an “average” diffusivity across the domain.
We also tried larger (deeper and/orwider) convolutional networks, but performance did not improve.
MoE methods (CondConv and V-MoE) also fail in this manner, as their coarse-grained experts are unable to specialize. Further, the
LRLCN-ND fails in this manner, despite its architecture being similar to an SM O E when there is oneoutput channel (a location-dependent, softmax-weighted combination of three basis kernels).
Webelieve the LRLCN-ND exhibits a similar gradient “mismatch” as discussed earlier (§2.2).
We now discuss a number of different ablations of the SM O E architecture and design.
What if the “right” expert conﬁguration is not known?
While in the above experiments, wewere able to select the SM O E expert conﬁguration (number of experts, number of selected experts,expert ﬁlter size) so that it is both necessary and sufﬁcient to learn the task, in many situations thisinformation may not be available.
We considered alternative SM O E conﬁgurations varying each ofthese parameters: ∂ using six experts; ∑ experts with 5 ⇥ 5 kernels; ∏ and selecting two experts perlocation from six total.
For case ∏ , we summed the two SM O E output channels together.
In all three cases, the SM O E achieved 100.0 ± 0 % within 1% on the heat diffusion task.
In ∂ , wefound that they learned duplicate diffusion stencils and still routed them appropriately.
∑ learned theﬁve-point stencil plus a boundary of near-zero values, thus being nearly identical to the 3 ⇥ 3 kernel.
Finally, ∏ learned diffusion stencils that summed together to produce the correct diffusivity.
Thus,we can see that SM O Es are robust and adapt well to these sorts of architecture choices.
RC loss and expert error damping.
Table 1 shows results for training SM O Es with and withoutour routing classiﬁcation loss (§2.2) and expert error damping (§2.3).
Without the RC loss, SM O Eperformance is at par with other baselines in Fig.
3, but once it is added, performance improves sig- [Table] [Table] niﬁcantly as the gating function now learns the location-dependency in the data.
Adding expert errordamping further improves performance and convergence by limiting the impact of gate misroutingson expert learning.
However, damping on its own offers little beneﬁt, as it does not improve gatelearning.
These results show that these reﬁnements are critical for good performance.
Gate function.
Table 2 shows the performance of different gating functions (§2.1) on the SM O E.
We consider six options: A single fully-connected layer (as is commonly used in MoEs [75, 79]); asingle 3 ⇥ 3 convolution, CoordConv [58], or LCN layer; a gate with three CoordConv layers withbatchnorm and ReLU; and our tensor routing gate.
When training, we also considered auxiliary lossesand other methods for improving performance (see §C) and report the best result.
Our tensor routingoffers the best performance.
An LCN performs second-best, likely because it also uses separateparameters per location, but uses 9 ⇥ as many parameters and requires signiﬁcant computation.
Othermethods do not appear able to effectively capture location-dependence.
Other ablations.
We conduct a number of additional ablation studies in §C, including using auxiliarylosses and routing noise during training, routing normalizations, and expert functions.
3.2 Medium-Range Weather Prediction
We now discuss results on the WeatherBench [71] medium-range weather forecasting benchmark.
This benchmark uses the ERA5 reanalysis dataset [38], with hourly global weather data for 1979–2018.
We use the data subset suggested by Rasp et al.
[71] at 5.625° resolution ( 32 ⇥ 64 grid points)and train on data from 1979–2015, validate on 2016, and report test results for 2017–2018; we donot use additional pretraining data.
We otherwise follow the training methodology of Rasp and
Thuerey [73].
The target quantities to predict are geopotential at 500 hPa (Z500) and temperatureat 850 hPa (T850) with a three- and ﬁve-day lead time.
We report results using latitude-weightedroot-mean-square error (RMSE).
As a baseline, we use the ResNet architecture [36] introduced by Rasp and Thuerey [73], whichcurrently reports the best results on WeatherBench.
This architecture consists of 19 residual blockseach with two [ 3 ⇥ 3 convolution !
LeakyReLU ! batchnorm !
dropout] layers, plus an initial 316⇥ 7 convolution layer.
All convolutions but the last have 128 ﬁlters.
We consider three additionalbaselines.
The ﬁrst is identical to the above, but with twice as many ﬁlters (256) in each convolution.
Second, we replace 3 ⇥ 3 convolutions with LRLCN [27] layers.
Finally, we use a four-layer ViT [23]with patch size 2 ⇥ 2 , hidden dimension 1024, and eight heads (the best performing conﬁguration).
We adapt the Rasp and Thuerey ResNet to use SM O Es by considering three different conﬁgurations:adding an SM O E layer after the ﬁrst convolution; adding an SM O E layer after the ﬁnal convolutionallayer; and replacing all 3 ⇥ 3 convolutions with SM O E layers.
Each SM O E selects the same numberof experts as the original layer had ﬁlters, and has twice as many experts (i.e., |E| = 256 and 324 E = 128 ).
We also share the tensor routing gate across all SM O E layers with the same number ofexperts, so the overhead of the gate is minimal.
Because the weather data is on a ﬁxed grid with underlying location-dependence (the Earth), weexpect SM O Es to convey some beneﬁt by specializing to the characteristics of different regions.
In
Table 3, we observe that this is indeed the case, and adding SM O Es improves results in all situations,with the most signiﬁcant improvement coming through replacing all 3 ⇥ 3 convolution layers with
SM O Es.
This showcases the advantage of incorporating appropriate location-dependent biases.
Wider ResNets offer limited improvement; this is in line with scaling results reported by Rasp and
Thuerey [73].
The location-dependent ﬁlters of LRLCNs improve over ResNets, but fail to match
SM O Es.
We were unable to achieve good performance with ViTs, but did observe that performanceis highly sensitive to patch size.
Incorporating prior knowledge into gates.
While the exact nature of the location-dependence ofthis data is unknown, we do have a broad prior on some aspects of it, such as whether a point is landor sea.
This information can be incorporated into an SM O E by initializing the tensor routing gateusing it to bias routing to different experts.
To this end, we use the land-sea mask from ERA5 toinitialze the gate to route land locations to half the experts and sea locations to the other half.
Notethis does not ﬁx the routing, as the gate is able to adjust as it learns.
Further, the land-sea mask isalready 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 conﬁguration sets anew state-of-the-art for WeatherBench when not using additional data.
Indeed, it nearly matches theperformance of a ResNet with 150 years of additional pretraining data [73].
We also tried a conﬁguration where the gate was initialized randomly and ﬁxed rather than learned(“rand ﬁxed gate init”).
This performs worse than our baseline, as the network cannot adapt its routingchoices, and each expert sees even fewer points in each sample than a standard network, resulting inless learning.
Thus, learning the routing function is critically important to good performance.
Mode collapse. Many MoEs suffer from expert or mode collapse (e.g., [9, 75, 79]), where only asmall number of experts are selected. This is typically avoided with a combination of routing noiseand auxiliary “load balance” losses. On the heat diffusion dataset, we found these losses to offer nobeneﬁt (§C).
We also examined SM O Es for mode collapse on WeatherBench, and did not observe it.
By using the RC loss, we directly train the gate, including updating routing weights toward otherexperts when a mistake is made, and so avoid such issues.
Expert selection order.
During training, the order experts are concatenated may change (either dueto changes in relative routing scores, or different experts being selected), which will impact the orderof channels seen by subsequent layers.
When training on WeatherBench, we found this not to havea signiﬁcant impact: expert order stabilizes early in training, allowing layers to operate on stablerepresentations.
Further, most “swapping” that does occur happens among experts a gate has lowconﬁdence in, so it is limited to a subset of channels.
Lastly, we present results on several image classiﬁcation tasks; we focus here on ImageNet-1k [21]and discuss results on additional datasets in §D. While these datasets do not have a strict locationdependent structure, relaxing the strict translation equivariance of convolutions can bring beneﬁts,and enables a direct comparison with Elsayed et al.
We follow their experimental methodologyand train using the recipe of Vryniotis [86].
We either insert an SM O E layer after the ﬁrst or lastconvolutional layer of ResNet-50 [36] or replace all 3 ⇥ 3 convolutions with SM O E layers.
Our SM O Es have twice as many experts as the original convolution layer and select half of them, tokeep output dimensions constant.
Gating layers are shared among all equally-sized blocks.
Forcomparison, we also train ResNet-50 with all 3 ⇥ 3 convolutions replaced by LRLCN [27] layers;and a Wide ResNet-50-2 [94], which has comparable parameters to SM O Es.
Table 4 shows that SM O Es outperform LRLCNs when we replace all 3 ⇥ 3 convolutions, while using56% of the parameters.
However, a wide ResNet performs best overall.
Nevertheless, this shows that
ImageNet classiﬁcation does indeed beneﬁt from relaxing translation equivariance.
Discussion
We presented the Spatial Mixture-of-Experts layer, a novel layer that learns underlying locationdependencies in data and then uses ﬁne-grained routing to specialize experts to different areas.
Wealso introduce a routing classiﬁcation loss and expert error damping, which enable SM O Es to performwell on dense regression tasks.
Prior work shows limited effectiveness on these tasks: Either itdoes not capture location-dependence (e.g., convolutions) or it operates at a coarse-grained level (e.g., standard MoEs).
By overcoming these challenges, we demonstrate a new capability for neuralnetworks, and set a new state-of-the-art for medium-range weather prediction.
Many other problems of broad societal import have a similar spatial structure, particularly in scientiﬁcdomains [8] and we expect SM O Es to be applicable to them.
However, tasks such as facialrecognition and surveillance have also historically shown beneﬁt from such improvements [84] and
SM O Es should therefore be used with care.
SM O Es show that learning location-dependence 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 developimproved implementations for ﬁne-grained, sparse routing; and to generalize SM O Es from operatingon grids to general graphs, which would enable them to be applied to many additional tasks.