# How aligned are different alignment metrics?

Jannis Ahlert1
&Thomas Klein1,2∗
&Felix Wichmann1
&Robert Geirhos3 Equal contribution. 1University of Tübingen, Germany 2Max
Planck Institute for Intelligent Systems, Tübingen 3Google DeepMind, Toronto,
Canada. Correspondence to<EMAIL_ADDRESS>
###### Abstract

In recent years, various methods and benchmarks have been proposed to
empirically evaluate the alignment of artificial neural networks to human
neural and behavioral data. But how aligned are different alignment metrics?
To answer this question, we analyze visual data from Brain-Score (Schrimpf et
al., 2018), including metrics from the model-vs-human toolbox (Geirhos et al.,
2021), together with human feature alignment (Linsley et al., 2018; Fel et
al., 2022) and human similarity judgements (Muttenthaler et al., 2022). We
find that pairwise correlations between neural scores and behavioral scores
are quite low and sometimes even negative. For instance, the average
correlation between those $80$ models on Brain-Score that were fully evaluated
on all $69$ alignment metrics we considered is only $0.198$. Assuming that all
of the employed metrics are sound, this implies that alignment with human
perception may best be thought of as a multidimensional concept, with
different methods measuring fundamentally different aspects. Our results
underline the importance of integrative benchmarking, but also raise questions
about how to correctly combine and aggregate individual metrics. Aggregating
by taking the arithmetic average, as done in Brain-Score, leads to the overall
performance currently being dominated by behavior (95.25% explained variance)
while the neural predictivity plays a less important role (only 33.33%
explained variance). As a first step towards making sure that different
alignment metrics all contribute fairly towards an integrative benchmark
score, we therefore conclude by comparing three different aggregation options.

##  1 Introduction

A central question in the field of representational alignment is whether two
given perceptual systems apply the same transformation to their inputs,
thereby extracting equivalent representations from data. These systems could
be artificial deep neural networks (DNNs) and biological systems like primate
brains, or arbitrary other image-computable models. For a number of reasons,
the perceptual alignment of artificial to biological systems has seen growing
interest in recent years. Well-aligned DNNs could serve as models of
biological visual processing and ultimately be an important step towards
building neural prosthetics. Furthermore, measuring discrepancies between
machine and biological perception can help to identify shortcomings in DNNs
and improve machine perception (Wichmann & Geirhos, 2023). For example,
aligning DNNs with the human visual system promises to increase their
robustness and generalization abilities (Dapello et al., 2022; Sucholutsky &
Griffiths, 2023).

Since the question of representational alignment between brains and machines
cannot be answered theoretically, various empirical methods have been proposed
to measure the representational alignment of different systems. These efforts
include the comprehensive Brain-Score benchmark (Schrimpf et al., 2018; 2020),
which integrates $51$ different metrics to capture the alignment of different
models to biological systems. In addition to neural data extracted from
primates, other metrics also capture behavioral similarity, such as error
pattern analysis (Rajalingham et al., 2018; Geirhos et al., 2020; 2021) and
shape bias (Geirhos et al., 2018; Baker et al., 2018; Hermann et al., 2020).
Other commonly used datasets in the literature include behavioral similarity
judgements (e.g., Hebart et al., 2023; Muttenthaler et al., 2022).

In light of this wealth of different metrics and datasets, all of which intend
to quantify how “human-like” (or brain-like, or primate-like) machine learning
models are, a fundamental question arises: How aligned are different alignment
metrics? For instance, if two metrics are highly correlated, it may be
sufficient to evaluate models on one of them instead of both. On the other
hand, if two metrics lead to highly dissimilar results, they measure different
aspects and would lead to very different conclusions and model
rankings—assuming the metrics are sound. In contrast, if inconsistencies
between different metrics were for instance caused by methodological errors
that do not reflect different aspects of brain-likeness, those cases could
equally be surfaced by analyzing the agreement of metrics as we do in this
study. In order to determine how aligned different alignment metrics are, we
therefore analyze a broad set of up to $241$ models evaluated on $50$
different metrics, depending on data availability.

##  2 Methods

##### Nomenclature.

In the following, we will refer to any function that maps a model to a scalar
_score_ measuring some form of alignment as a _metric_ or _measure_. Any set
of metrics that together estimate alignment is referred to as a _benchmark_.
The point of this distinction is that for a benchmark, some method of
integrating its constituent metrics is necessary, like taking their average.

##### Data sources.

A central hub for collecting and integrating similarity measurements between
DNNs and biological neural data is the Brain-Score benchmark (Schrimpf et al.,
2018; 2020). We use data from Brain-Score as the foundation of our analyses,
but also integrate data from three other sources: Muttenthaler et al. (2022),
Fel et al. (2022) and Geirhos et al. (2021).

The Brain-Score benchmark (Schrimpf et al., 2018) provides comparative data
for two domains, vision and language. Throughout this paper, we focus
exclusively on vision and thus on the visual Brain-Score. Brain-Score
evaluation results can be thought of as a table, where the columns are
different metrics and the rows are models evaluated on these metrics. The
columns are semantically grouped into three sets: Neural, behavioral and
engineering benchmarks. The neural benchmark can be further subdivided into
groups that measure a model’s similarity to different brain areas: V1, V2, V4
and IT. The behavioral benchmark consists of image classification metrics like
error consistency, either for humans or monkey subjects. The engineering
benchmark, which does not enter the final scoring, captures technical
properties of models like ImageNet top-1 accuracy.

We are particularly interested in the agreement between neural and behavioral
measures of similarity between humans and machines, so we use data from
Geirhos et al. (2021), who record human classification performance on
(corrupted) ImageNet images and calculate multiple types of behavioral
similarity: _error consistency_ (Do humans and machines make mistakes on the
same images?) and _shape bias_ (When presented with shape-texture cue
conflicts, do models prefer shape, like humans?). Muttenthaler et al. (2022)
systematically analyze which design choices affect the human-machine alignment
for various DNNs, finding that neither architecture nor model scale are
relevant, but that the training data and objective function are driving
factors of alignment. They quantify human-machine alignment as a network’s
ability to predict human behavior on triplet tasks (Hebart et al., 2020), for
which human subjects are presented with three natural images and classify one
of them as the odd-one-out. Hence, we call this metric _odd-one-out
similarity_. Fel et al. (2022) define human-machine alignment as the
similarity between model attention maps and attention maps obtained from human
crowd-workers. Linsley et al. (2018) collected about 500k of these human
attention maps, which quantify which regions of an image humans attend. We
refer to this metric as _attention map similarity_. See Section A.2 for a
detailed breakdown of models used in our analyses.

##### Pairwise comparisons.

To perform pairwise comparisons between two benchmarks, we calculate
correlations between their respective scores for the same set of models. We
calculate Spearman’s rank correlation, which we believe to be appropriate
since we are interested in the ranking of models (rather than measuring linear
relationships as Pearson’s correlation does). We apply very conservative
Bonferroni-corrections to the thresholds of significance to account for
multiple comparisons. Specifically, we divide the threshold of $\alpha=0.05$
by the number of metric-pairs, resulting in $\alpha=0.0009$ for Figure 1 and
$\alpha=0.001$ for Figure 5.

##### Integrating metrics: arithmetic mean, z-transformed mean, mean rank.

Given a set of metrics, what is the best way of aggregating them to a single,
unified score? While a one-size-fits-all solution is unlikely to exist, we
observe that aggregation choices can influence conclusions and thus warrant
attention. Brain-Score uses a simple arithmetic average; in Section 4 we
explore alternatives. A typical technique for dealing with different scales
involves z-transforming values before aggregating them, which enforces that
all individual metrics have zero mean and unit standard deviation.
Alternatively, a ranking of models at some level of the Brain-Score hierarchy
could be achieved by first ranking all metrics individually, and then
calculating the average rank for every model.

##  3 How aligned are different metrics?

##### Correlations between Brain-Score benchmarks.

To investigate the internal consistency of metrics on Brain-Score, we
calculate Spearman’s rank correlation coefficient for every pair of metrics,
taking only those $95$ models into account for which all scores are available.
Removing duplicates and one model with an ImageNet accuracy below 1.9% leaves
88 models. We find that behavioral measures correlate strongly, especially
those that all use the same dataset (Geirhos et al., 2021). For many neural
metric groups such as V4 and IT, as well as sub-groups within the V1 region
relating to texture and receptive field size, internal consistency is higher
than consistency with other metrics. We relegate detailed results to the
Appendix, see Figure 4.

##### Moving beyond Brain-Score.

Next, we investigate how the two new metrics that are not yet included in the
Brain-Score project, odd-one-out similarity and attention map similarity,
correlate with metrics on Brain-Score in Figure 1. We again calculate pairwise
Spearman’s rank correlations between metrics for the set of models for which
scores are available on all metrics. The odd-one-out similarity metric by
Muttenthaler et al. (2022), which is of behavioral nature, is more correlated
with metrics of V1 than with other behavioral metrics or later neural areas
such as IT. The attention map similarity by Fel et al. (2022) is not
significantly correlated with any other metric.

Figure 1: How aligned are different alignment metrics? Pairwise Spearman’s
rank correlations of different Brain-Score metrics, as well as odd-one-out
similarity judgements and attention map similarity. Correlations that are
significant after Bonferroni-correction are bold. We include only those $42$
models that were evaluated on all metrics. Note that (a) variance of
correlation coefficients is quite high and (b) similar metrics tend to agree,
with the exception of the odd-one-out similarities. See also Figure 5 for the
$80$ Brain-Score models that have all scores on the Brain-Score metrics of
this heatmap.

##  4 Alternatives to the arithmetic average: How do aggregation choices
influence alignment results?

The overall score that a model achieves on Brain-Score is calculated as the
arithmetic mean of its neural and behavioral scores. In order for this to work
well, the two values should at least have roughly the same magnitude, and
ideally live in a common metric space. In Figure 2, we scatter behavioral
against neural scores, demonstrating the differences in their variance. Note
also Linsley et al. (2023), who find that models which perform better on
ImageNet (which is a behavioral metric) tend to have worse IT-alignment (a
neural metric). It is evident that the best behavioral scores of about $0.6$
are quite a bit higher than the best neural scores of $0.5$. Because of
different scaling of the metrics, models with mediocre neural scores can still
rank highly in the benchmark by virtue of their good behavioral scores. In the
overall model rank, we find examples of this: The two best models on Brain-
Score, a Convolutional Vision Transformer 111https://www.brain-
score.org/model/vision/1885 and a ResNeXt variant 222https://www.brain-
score.org/model/vision/646 by Mahajan et al. (2018) are only the 71st and 75th
best model according to neural scores, respectively.

Figure 2: Left: Relationship between neural and behavioral scores. Coloring
represents the model rank in Brain-Score, which is the average of the neural
and behavioral scores (brighter color indicates higher overall score). Right:
Same data after z-transforming both scores. The dominance of behavioral scores
prevails, because there are extreme outliers in the neural scores.

Assuming that we have obtained a matrix of model scores on different, possibly
uncorrelated metrics with different distributions, we should ask ourselves how
to best integrate these scores into a final model-score, or at least into an
overall ranking of models. We now consider three different possible ranking
schemes. The first strategy, which is the default in Brain-Score, is to
integrate metrics via their arithmetic mean, which is sensitive to scaling
issues and high-variance metrics. As an alternative, we consider _z-
transforming_ the score for every metric before averaging. We demonstrate how
this scheme would change scores in Figure 2. A drawback of this approach is
that such scores are no longer stable over time, but depend on the scores of
other models evaluated on that metric. Alternatively, if one does not care
much about absolute model scores but only about their relative order, a
possible integration scheme is obtaining the _rank order_ for every metric, to
then average over these ranks for the final score. Such a scheme would drop
quantitative information about differences, but avoids inconsistency issues
like the examples with high overall rank despite extremely low-rank
performance on some metrics. In practice, applying these ranking schemes to
Brain-Score does indeed lead to changes in model ordering (see Figure 3). Each
row corresponds to a ranking scheme, and each model is represented as one
point. A model’s position on the x-axis corresponds to its score: Good models
are further on the right, worse models further on the left, with scores
normalized to span the range $[0-1]$. To visualize the impact of the ranking
scheme, we color-code models by their position in the default ranking, with
higher scores colored brighter. Arguably, the differences in the resulting
rank orderings are not “dramatic”. However, if further studies were to e.g.
correlate brain-likeness against other quantities, the effects of such changes
could be quite large. We believe the issue of how to aggregate metrics
deserves more scrutiny in the future.

Figure 3: Comparison of rankings resulting from different integration schemes.
We compare three different ranking algorithms: (1) Arithmetic mean corresponds
to the current standard in Brain-Score, where scores are simply averaged. (2)
The z-transformed ranking is obtained by z-transforming scores first, then
averaging them. (3) Mean Rank is the result of averaging the ranks implied by
the metrics (ranks are inverted for consistency, so higher is still better).
All scores are normalized to the interval $[0-1]$ using Min-Max Normalization
and colored according to a model’s position under the original integration
scheme. Rank-order changes greater than $10$ positions (relative to the rank
implied by the arithmetic mean) are highlighted in red. Spearman’s rank
correlations between the original ranking and the new ones are 0.47 (mean
rank) and 0.92 (z-transformed mean).

##  5 Discussion

Recent years have seen a wealth of datasets, benchmarks and metrics to measure
the alignment between brains and machines. This is a promising research
avenue, but in light of an ever-expanding set of measures, this raises the
question of how aligned different alignment metrics truly are. We have
contributed an analysis of the relationships between different alignment
measures. The resulting correlations clearly indicate that different metrics
capture different aspects of perceptual alignment, sometimes leading to
contradictory results. For instance, models that score as highly human-like on
one metric may be among the worst models according to a different metric.
Consequently, alignment with human perception may best be thought of as a
multidimensional concept, and as a community we may need to spend more time
thinking about how to properly integrate different metrics. As a first step in
this direction, we have investigated how different aggregation choices
influence alignment results. Going beyond simple aggregation methods, in which
metrics are treated equally and independently, one could also consider aspects
like the semantic structure of metrics: Should we discount a model’s high IT-
score if it performs poorly on V1 and V2, indicating that it arrives at the
right solution via the wrong path? Ultimately, it could be valuable for the
community to derive a set of axiomatic requirements that an ideal integration
scheme should fulfill, and categorize existing aggregation choices according
to how well they conform to those. Overall, it seems clear that while there is
a need for unified and comprehensive benchmarks, it may be important to
consider alignment as an inherently multidimensional concept.

#### Acknowledgments

The authors would like to thank Martin Schrimpf for providing raw Brain-Score
data, Alexander Riedel for providing checkpoints of his models, Britta Dorn
for inspiring ranking algorithms based on computational social choice, Priyank
Jaini and Nina Wiedemann for feedback on the manuscript, as well as Simon
Kornblith, Lukas Muttenthaler and Martin Schrimpf for discussions.
Furthermore, we would like to thank the anonymous reviewers for their
constructive feedback.

Funded by the Deutsche Forschungsgemeinschaft (DFG, German Research
Foundation) under Germany’s Excellence Strategy—EXC number 2064/1—project
number 390727645 and SFB 1233—project number 276693517. This work was
supported by the German Federal Ministry of Education and Research (BMBF):
Tübingen AI Center, FKZ: 01IS18039A. TK acknowledges financial support via an
Emmy Noether Grant funded by the German Research Foundation (DFG) under grant
no. BR 6382/1-1, awarded to Wieland Brendel. The authors would like to thank
the International Max Planck Research School for Intelligent Systems (IMPRS-
IS) for supporting Thomas Klein.

## References

  * Baker et al. (2018) Nicholas Baker, Hongjing Lu, Gennady Erlikhman, and Philip J Kellman.  Deep convolutional networks do not classify based on global object shape.  _PLoS Computational Biology_ , 14(12):e1006613, 2018.
  * Dapello et al. (2022) Joel Dapello, Kohitij Kar, Martin Schrimpf, Robert Geary, Michael Ferguson, David D Cox, and James J DiCarlo.  Aligning model and macaque inferior temporal cortex representations improves model-to-human behavioral alignment and adversarial robustness.  _bioRxiv_ , pp. 2022–07, 2022.
  * Fel et al. (2022) Thomas Fel, Ivan Felipe Rodriguez, Drew Linsley, and Thomas Serre.  Harmonizing the object recognition strategies of deep neural networks with humans.  _Advances in Neural Information Processing Systems_ , 35:9432–9446, 2022.
  * Geirhos et al. (2018) Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A Wichmann, and Wieland Brendel.  ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness.  In _International Conference on Learning Representations_ , 2018.
  * Geirhos et al. (2020) Robert Geirhos, Kristof Meding, and Felix A Wichmann.  Beyond accuracy: quantifying trial-by-trial behaviour of cnns and humans by measuring error consistency.  _Advances in Neural Information Processing Systems_ , 33:13890–13902, 2020.
  * Geirhos et al. (2021) Robert Geirhos, Kantharaju Narayanappa, Benjamin Mitzkus, Tizian Thieringer, Matthias Bethge, Felix A Wichmann, and Wieland Brendel.  Partial success in closing the gap between human and machine vision.  In _Advances in Neural Information Processing Systems 34_ , 2021.
  * Hebart et al. (2020) Martin N Hebart, Charles Y Zheng, Francisco Pereira, and Chris I Baker.  Revealing the multidimensional mental representations of natural objects underlying human similarity judgements.  _Nature human behaviour_ , 4(11):1173–1185, 2020.
  * Hebart et al. (2023) Martin N Hebart, Oliver Contier, Lina Teichmann, Adam H Rockter, Charles Y Zheng, Alexis Kidder, Anna Corriveau, Maryam Vaziri-Pashkam, and Chris I Baker.  Things-data, a multimodal collection of large-scale datasets for investigating object representations in human brain and behavior.  _Elife_ , 12:e82580, 2023.
  * Hermann et al. (2020) Katherine Hermann, Ting Chen, and Simon Kornblith.  The origins and prevalence of texture bias in convolutional neural networks.  _Advances in Neural Information Processing Systems_ , 33:19000–19015, 2020.
  * Huang et al. (2017) Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger.  Densely connected convolutional networks.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 4700–4708, 2017.
  * Kolesnikov et al. (2020) Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigcerver, Jessica Yung, Sylvain Gelly, and Neil Houlsby.  Big transfer (bit): General visual representation learning.  In _Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part V 16_ , pp. 491–507. Springer, 2020.
  * Krizhevsky et al. (2012) Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton.  Imagenet classification with deep convolutional neural networks.  _Advances in neural information processing systems_ , 25, 2012.
  * Linsley et al. (2018) Drew Linsley, Dan Shiebler, Sven Eberhardt, and Thomas Serre.  Learning what and where to attend.  _arXiv preprint arXiv:1805.08819_ , 2018.
  * Linsley et al. (2023) Drew Linsley, Ivan F Rodriguez, Thomas Fel, Michael Arcaro, Saloni Sharma, Margaret Livingstone, and Thomas Serre.  Performance-optimized deep neural networks are evolving into worse models of inferotemporal visual cortex.  _arXiv preprint arXiv:2306.03779_ , 2023.
  * Ma et al. (2018) Ningning Ma, Xiangyu Zhang, Hai-Tao Zheng, and Jian Sun.  Shufflenet v2: Practical guidelines for efficient cnn architecture design.  In _Proceedings of the European conference on computer vision (ECCV)_ , pp. 116–131, 2018.
  * Mahajan et al. (2018) Dhruv Mahajan, Ross Girshick, Vignesh Ramanathan, Kaiming He, Manohar Paluri, Yixuan Li, Ashwin Bharambe, and Laurens van der Maaten.  Exploring the limits of weakly supervised pretraining.  In _Proceedings of the European Conference on Computer Vision (ECCV)_ , September 2018.
  * Muttenthaler et al. (2022) Lukas Muttenthaler, Jonas Dippel, Lorenz Linhardt, Robert A Vandermeulen, and Simon Kornblith.  Human alignment of neural network representations.  _arXiv preprint arXiv:2211.01201_ , 2022.
  * Rajalingham et al. (2018) Rishi Rajalingham, Elias B. Issa, Pouya Bashivan, Kohitij Kar, Kailyn Schmidt, and James J. DiCarlo.  Large-scale, high-resolution comparison of the core visual object recognition behavior of humans, monkeys, and state-of-the-art deep artificial neural networks.  _Journal of Neuroscience_ , 38(33):7255–7269, August 2018.  ISSN 15292401.  doi: 10.1523/JNEUROSCI.0388-18.2018.  Publisher: Society for Neuroscience.
  * Salman et al. (2020) Hadi Salman, Andrew Ilyas, Logan Engstrom, Ashish Kapoor, and Aleksander Madry.  Do adversarially robust imagenet models transfer better?  _Advances in Neural Information Processing Systems_ , 33:3533–3545, 2020.
  * Sandler et al. (2018) Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen.  Mobilenetv2: Inverted residuals and linear bottlenecks.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 4510–4520, 2018.
  * Schrimpf et al. (2018) Martin Schrimpf, Jonas Kubilius, Ha Hong, Najib J. Majaj, Rishi Rajalingham, Elias B. Issa, Kohitij Kar, Pouya Bashivan, Jonathan Prescott-Roy, Franziska Geiger, Kailyn Schmidt, Daniel L. K. Yamins, and James J. DiCarlo.  Brain-score: Which artificial neural network for object recognition is most brain-like?  _bioRxiv preprint_ , 2018.
  * Schrimpf et al. (2020) Martin Schrimpf, Jonas Kubilius, Michael J Lee, N Apurva Ratan Murty, Robert Ajemian, and James J DiCarlo.  Integrative benchmarking to advance neurally mechanistic models of human intelligence.  _Neuron_ , 2020.
  * Simonyan & Zisserman (2014) Karen Simonyan and Andrew Zisserman.  Very deep convolutional networks for large-scale image recognition.  _arXiv preprint arXiv:1409.1556_ , 2014.
  * Sucholutsky & Griffiths (2023) Ilia Sucholutsky and Thomas L Griffiths.  Alignment with human representations supports robust few-shot learning.  _arXiv preprint arXiv:2301.11990_ , 2023.
  * Szegedy et al. (2015) Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich.  Going deeper with convolutions.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 1–9, 2015.
  * Szegedy et al. (2016) Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna.  Rethinking the inception architecture for computer vision.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 2818–2826, 2016.
  * Szegedy et al. (2017) Christian Szegedy, Sergey Ioffe, Vincent Vanhoucke, and Alexander Alemi.  Inception-v4, inception-resnet and the impact of residual connections on learning.  In _Proceedings of the AAAI conference on artificial intelligence_ , volume 31, 2017.
  * Tan & Le (2019) Mingxing Tan and Quoc Le.  Efficientnet: Rethinking model scaling for convolutional neural networks.  In _International conference on machine learning_ , pp. 6105–6114. PMLR, 2019.
  * Wichmann & Geirhos (2023) Felix A Wichmann and Robert Geirhos.  Are deep neural networks adequate behavioural models of human visual perception?  _Annual Review of Vision Science_ , 9:501–524, 2023.
  * Xie et al. (2020) Cihang Xie, Mingxing Tan, Boqing Gong, Jiang Wang, Alan L Yuille, and Quoc V Le.  Adversarial examples improve image recognition.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pp. 819–828, 2020.

##  Appendix A Appendix

###  A.1 Limitations

Our analyses are limited by several technical issues. First, most of the
pairwise correlations calculated in Figure 1 are not statistically
significant, because so few models were present in all datasets. There also is
some uncertainty about whether model weights were matched perfectly.

###  A.2 Methodological Details (Model Selection)

At the time of writing, Brain-Score consists of $241$ models. However, not all
of those models were evaluated on all metrics: Only $88$ different models were
evaluated on all metrics that contribute to the final score (i.e. all neural
and behavioral metrics). Completing this evaluation would be valuable to allow
broader analyses, as would adding exact pointers to model checkpoints where
this is feasible.

For the analysis in Figure 1, we determine a subset of the models on Brain-
Score for which weights are publicly available and evaluate them on the
metrics of Fel et al. (2022) and Muttenthaler et al. (2022). This restriction
limits the scope of our analysis and we hope to increase this sample in the
future. Thus, the models currently evaluated on all considered metrics are:

  1. 1.

AlexNet (Krizhevsky et al., 2012) via torchvision

  2. 2.

DenseNet 121, 169 and 201 (Huang et al., 2017) via torchvision

  3. 3.

EfficientNet B0 and B7 (Tan & Le, 2019) via torchvision

  4. 4.

Inception V1 (Szegedy et al., 2015) and V3 (Szegedy et al., 2016) via
torchvision

  5. 5.

Inception V4 (Szegedy et al., 2017) via timm

  6. 6.

VGG 16 and 19 (Simonyan & Zisserman, 2014) via torchvision

  7. 7.

ShuffleNetV2-x1.0 (Ma et al., 2018) via torchvision

  8. 8.

Robust ResNet50-$\ell_{2}$ with $\epsilon=1$ and $\epsilon=3$ from (Salman et
al., 2020)

  9. 9.

13 MobileNetV2 variants (Sandler et al., 2018)

  10. 10.

6 BiT-S Variants by (Kolesnikov et al., 2020)

  11. 11.

6 AdvProp-EfficientNet variants by Xie et al. (2020)

  12. 12.

ResNet50 trained on SIN, SIN & IN, SIN & IN & finetuned on IN from (Geirhos et
al., 2018)

###  A.3 Further Results

Figure 4: How consistent are Brain-Score metrics? Pairwise Spearman’s rank
correlations of different metrics, from V1 to IT and finally to behavioral
measures. Note how the agreement between the behavioral metrics is much higher
than the agreement of the neural measures.

To calculate the correlations presented in Figure 1, we included all models
that were evaluated both on Brain-Score and the metrics proposed by
Muttenthaler et al. (2022) and Fel et al. (2022). However, this set of models
is relatively small ($42$ models total), so the correlations reported between
the different Brain-Score metrics are not necessarily representative. We
demonstrate how Figure 1 would have looked like if all Brain-Score models had
been included in Figure 5.

Figure 5: Pairwise Spearman’s rank correlations of Brain-Score metrics, from
V1 to IT and behavioral measures. In contrast to Figure 1, we include all
Brain-Score models that had a value available for each metric, not only the
subset that was also evaluated on the metrics by Muttenthaler et al. (2022)
and Fel et al. (2022). This amounts to a total of $80$ models. The average
correlation on this heatmap is 0.17, while the average correlation of the
smaller set of models in 1 is 0.21, exemplifying the need for further thorough
evaluations.

The neural score that a model obtains on Brain-Score is defined as the average
over four metrics measuring alignment to different brain regions (V1, V2, V4,
IT). This averaging strategy will work well if those scores are highly
correlated and similarly distributed. We therefore investigate the
distributions of neural scores in Figure 6, revealing similar distributions
after removal of models with incomplete evaluations.

Figure 6: Kernel Density Estimates of different neural metrics in Brain-Score.
This plot includes those 88 models for which all scores that contribute to the
overall Brain-Score ranking are available. Evidently, the distribution of
scores for areas V1, V4 and IT are very similar, but the scores for area V2
are much lower, either hinting at lack of progress towards predicting V2 or at
unfair calibration of this metric.

While we find that behavior already accounts for a disproportionate amount of
the overall score variance on the current Brain-Score leaderboard, we next
investigate which role missing values play in determining the aggregate scores
Figure 7. Coloring each model by the number of non-zero scores available
across metrics reveals a striking pattern: Most models fall into one of a few
homogeneous groups, having $50$, $33$, $30$ or $5$ non-zero metrics
respectively. Presumably, not all authors re-triggered scoring when new chunks
of metrics were added to Brain-Score. The plot reveals that behavioral scores
account for even larger portions of variance when only comparing models that
have the same number of non-zero metrics with each other.

Figure 7: How do missing scores shape Brain-Score? Scatterplot of neural and
behavioral averages for all models on Brain-Score. Models are colored
according to the number of benchmarks they have been evaluated on, as
determined by counting scores that are not exactly zero. Darker colors
indicate more non-zero scores, yellow models have only been evaluated on $5$
of $51$ possible metrics.


#  SEER-MoE: Sparse Expert Efficiency through Regularization for Mixture-of-
Experts

Alexandre Muzio† Alex Sun ‡ Churan He ‡
Stanford University †Google ‡NVIDIA

###### Abstract

The advancement of deep learning has led to the emergence of Mixture-of-
Experts (MoEs) models, known for their dynamic allocation of computational
resources based on input. Despite their promise, MoEs face challenges,
particularly in terms of memory requirements. To address this, our work
introduces SEER-MoE, a novel two-stage framework for reducing both the memory
footprint and compute requirements of pre-trained MoE models. The first stage
involves pruning the total number of experts using a heavy-hitters counting
guidance, while the second stage employs a regularization-based fine-tuning
strategy to recover accuracy loss and reduce the number of activated experts
during inference. Our empirical studies demonstrate the effectiveness of our
method, resulting in a sparse MoEs model optimized for inference efficiency
with minimal accuracy trade-offs.

##  1 Introduction

Recent advances in deep learning have propelled the field towards increasingly
large and complex models to achieve state-of-the-art performance across a
myriad of tasks. Among these, Mixture-of-Experts (MoEs) models have emerged as
a promising architecture, distinguished by their ability to dynamically
allocate computational resources based on the input (Fedus et al., 2021; Zoph
et al., 2022; Jiang et al., 2024; et al, 2024). This paradigm, characterized
by a sparse gating mechanism that routes inputs to a subset of specialized
expert networks, allows for the scalable expansion of model parameters while
maintaining a relatively constant computational footprint per input token.

However, the path forward for MoEs involves addressing significant challenges,
particularly regarding their substantial memory requirements. The advent of
very large MoE models such as Grok-1 xAI (2024), with 314B parameters
distributed across 8 experts, underscores the urgency of addressing this
issue. While the sparse nature of MoEs promises enhanced efficiency and
scalability, the sheer size of the models introduces a new set of
complexities, particularly in terms of memory footprint.

In light of these challenges, our work aims to investigate and develop pruning
and fine-tuning strategies that dynamically adjust the quantity and allocation
of experts in an MoE-based language model.

The main contributions of our paper is an in-depth study of Parameter Count /
FLOPs for MoE models and the SEER-MoE method, a novel 2-stage approach that
takes a step in the direction of reducing the memory footprint / compute
requirements for pretrained MoE models. Our first stage proposes to prune the
total number of experts in the model with a novel heavy-hitters counting
guidance to reduce the memory footprint for loading the entire MoE model. Our
second stage proposes an effecitve regularization-based finetuning strategy to
recover the accuracy loss from previous pruning while simultanously reducing
the number of activated experts during inference. The combination of both
stages yields a sparse MoEs model with cheaper memory requirements while being
optimized for inference efficiency, at the compensation of minimized accuracy
drops.

We perform extensive empirical studies with the popular Mixtral 8x7b Jiang et
al. (2024) MoEs model on both SST5 Socher et al. (2013) and MMLU Hendrycks et
al. (2020) to validate the effectiveness of our method, including an in-depth
ablation study to understand different design choices for each of our stage.

##  2 Related Work

Mixture-of-Experts (MoE) dates back at least since work from Jacobs et al.
(1991), which introduce a new model architecture composed of many separate
networks and each one handles a subset of the complete set of training cases.
Each expert specializes in a different region of the input space. Eigen et al.
(2014) extended the Mixture-of-Experts to a stacked model with multiple layers
of gating and experts, and exponentially increases the number of effective
experts through layers of combination. As the rapid advancement of LLM, MoE
(Jiang et al., 2024; et al, 2024; xAI, 2024) gained increased popularity for
it’s scalability, efficiency and STOA evaluation result from various
benchmarks. Most MoE architectures Fedus et al. (2021); Zoph et al. (2022);
Jiang et al. (2024) includes a specific gating network that learns the optimal
routing from input tokens to experts. However, the weights of the gating
network stay fixed, regardless of what task is being solved. To the best of
our knowledge, there is no such method that explores Top-K routing adaptation.

Sparsification aims to remove certain parts of the network. In Mixture-of-
Experts models, only a few experts (top-K) chosen by the router will be
activated in each layer to generate output, therefore removing unused experts
can linearly reduce the model size without loss in performance. (Lu et al.,
2024) introduced a heuristic search method to prune the number of experts in
post-training. The method is based on the enumeration of expert combinations
and choosing the target eliminating experts based on the lowest reconstruction
loss. They verified their approach’s effectiveness on Mixtral 8x7b (Jiang et
al., 2024). However, the method has high time complexity and isn’t applicable
to models with large expert counts nor cross-layer pruning (Zhang et al.,
2023) proposed an eviction algorithm targeting on KV cache based on "Heavy
Hitters" which holds a significant role for model performance. We hypothesize
that a similar strategy in expert activation counting can be applied to MoE
pruning.

##  3 Methodology

The computation and resources used by large MoEs models (FLOPs) mainly come
from two factors:

The total number of available experts at each MoE layer, denoted as $M_{l}$
for layer $l$, and determines how much VRAM we need to completely load the
model on the GPUs/TPUs.

The top-K number of experts activated, denoted as $k_{l}$ for layer $l$, and
controls how much computation is used per token.

In terms of notation, suppose we have a fine-tuning dataset $\mathcal{D}$
which consists of sample $(x_{i},y_{i})$ pairs, the gating network at layer
$l$ is $g_{l}$, and the $j$th expert at the $l$th layer is denoted as
$e_{l}^{j}$. Moreover, suppose the features collected at layer $l$ is
$f_{l}(x_{i})$.

Figure 1: SEER-MoE visualized in a two-stage process. (a) The initial model
with all experts and top-2 router. (b) Stage 1 involves expert pruning based
on heavy-hitters counting to identify and retain the most critical experts;
Stage 2 includes top-K adaptation through fine-tuning to optimize the number
of active experts, culminating in a model that balances efficiency and
performance. (c) SEER-MoE with expert pruning and top-K adaptation.

###  3.1 Parameter and Compute Scaling of MoE Transformers

Following the work of Kaplan et al. (2020), we extend the parameterization of
the Transformer architecture to that of the sparse MoE architecture (more
details in Appendix A). We utilize the Transformer hyperparameters together
with $n_{experts}$ (number of experts per layer) and $n_{topk}$ (number of
experts activated per token). Using $N$ to denote the number of non-embedding
parameters and considering $d_{attn}=d_{ff}/4=d_{model}$

$N\approx 4d^{2}_{model}n_{layer}(1+2n_{experts})$ (1)

And the FLOPs for the forward pass:

$C_{fwd}\approx
8d^{2}_{model}n_{layer}(1+2n_{topk})+2n_{layer}n_{ctx}d_{model}$ (2)

It is notable that for the MoE blocks of the network, the number of parameters
from the MoE Feedforward layers increase in proportion to $n_{experts}$ while
the FLOPs per Token to $n_{topk}$. This motivates us to think whether we can
target reducing both the compute FLOPs of the model and the memory separately.

Specifically for the Mixtral 8x7B Jiang et al. (2024) model, which utilizes
$2$ experts per token, Expert blocks computations account for about $55\%$ of
the total FLOPS. Additionally, for a model with the same architecture but with
only a single expert being activated, FLOPs reduces by $27\%$. This motivates
us to explore whether we are able to adapt existing models to use less
compute.

###  3.2 Expert Sparsification with Heavy-hitters Counting

Considering the large memory requirement of MoE models described and
considering that the bulk of the parameters belong to the Expert layers, we
propose reducing the total number of experts $M_{l}$ at each layer $l$. We do
so, by reducing the number of expert models. We propose to carry out this
investigation with a data-driven strategy.

For a MoEs model, an expert $e_{l}^{j}$ is activated for a token if its
corresponding router logit $g_{l}(f_{l}(x_{i}))^{j}$ is ranked in the top-K
after softmax. Specifically, we denote this selection process with a function
$\gamma(.)$:

$\displaystyle\gamma(e_{l}^{j},x_{i})$ $\displaystyle=1\text{ if
}j\in\text{ArgTopK}(\text{softmax}(g_{l}(f_{l}(x_{i}))),K),$ (3)
$\displaystyle\gamma(e_{l}^{j},x_{i})$ $\displaystyle=0\text{ else }$ (4)

For each expert $e_{l}^{j}$, we define the activation counts $a_{l}^{j}$ as
the total number of times it gets activated. Formally, this could be expressed
as follows:

$\displaystyle
a_{l}^{j}=\sum_{(x_{i},y_{i})\in\mathcal{D}}\mathbbm{1}[\gamma(e_{l}^{j},x_{i})=1]$
(5)

Equivalently, Eqn.6 is performing a Monte-Carlo estimate of the marginal
probability $P(e_{l}^{j}\text{ gets activated})$ using the dataset
$\mathcal{D}$, which provides theoretical motivation for our adopted technique
here.

Soft Counting We also propose another variant with a softer and more relaxed
version of heavy-hitters counting. Intuitively, for certain tokens, if expert
$e_{l}^{j}$ is activated, we don’t know whether it wins over other experts by
a slight margin or gets activated with high confidence. Since the binary
activation count can not capture this exact magnitude of confidence in
activating certain experts, we propose to directly leverage the softmax
probabilities as soft counts. Formally, this can be defined as:

$\displaystyle
a_{l}^{j}=\sum_{(x_{i},y_{i})\in\mathcal{D}}\text{softmax}(g_{l}(f_{l}(x_{i})))^{j}$
(6)

Layer Expert Pruning Now with the statistics of the heavy-hitters counts, we
could leverage them as powerful guidance to remove experts that are unlikely
to be activated for data from $P_{\mathcal{D}}$. Suppose we want to keep a
total of $\hat{M}_{l}$ experts per layer, the kept experts at layer $l$ are
denoted as:

$\displaystyle\text{ArgTopK}(\bigcup_{j\in[1,M_{l}]}\\{a_{l}^{j}\\},\hat{M_{l}})$
(7)

We repeat this for every layer in the MoEs model to get an entire mask.

Global Expert Pruning Since the counts have uniform magnitude range across all
layers, we can also carry out a global sorting and pruning to remove the
experts with the least probability of getting activated. Suppose we only want
to keep a total of $\hat{M}$ experts in the network. The kept ones are denoted
as:

$\displaystyle\text{ArgTopK}(\bigcup_{l\in[1,L]}\bigcup_{j\in[1,M_{l}]}\\{a_{l}^{j}\\},\hat{M})$
(8)

Compared with the above Layer Expert Pruning option, this will provide a
nonuniform expert sparsity pattern across layers but could potentially be of
higher-quality.

To recap, for pruning experts from the MoEs model to reduce storage burden and
memory footprints, we propose to perform pruning based on heavy-hitters
counting. This counting could be either actual activation counts or a soft
counting with softmax probabilities. The actual removal of experts could be
conducted either layer-wise or globally. This gives us a total of four
configurations with combinations of pairs of counting and removal strategies,
and we are going to provide detailed ablation results in the Experiment
section.

###  3.3 Enhancing Expert Efficiency: Advanced Finetuning Strategies

With the goal of reducing the number of experts activated for each token
during inference, while still maintaining competitive performance, we propose
different fine-tuning procedures for MoE model.

####  3.3.1 Top-K adaptation

Starting from a pretrained model trained with top-k > $1$, we posit that by
fine-tuning the model on a downstream task while reducing $k$ during training
is a feasible and simple strategy to adapt the model to utilize. We focus on
fine-tuning since we are interested in utilizing existing pretrained more
efficiently.

Given that we are trying to target the best open-source available MoE model,
which is the Mixtral 8x7b standard fine-tuning was not feasible given the
amount of memory required (even using 8xA100 80GB), therefore we opted for
explore QLoRA Dettmers et al. (2023) fine-tuning on the self-attention blocks
to reduce the memory footprint of the optimization.

In this work, we propose Top-K reduction procedures with simplicity in mind:
Static top-k with $k<K$ and Annealing top-k from $K\rightarrow k$ with $k<K$.
We also explore additional methods such as QLoRA fine-tuning targeting only
the gating network.

####  3.3.2 Entropy-based gating regularization

Entropy, in the context of information theory, is a measure of the
unpredictability or randomness of a distribution. For a categorical
probability distribution (which is the case of the expert gating network),
entropy is defined as $H(X)=-\sum_{i=1}^{n}p(x_{i})\log p(x_{i})$. We posit
that a gating network with a more peaky distribution, meaning lower entropy,
relies more heavily on a single expert. Therefore, by minimizing the entropy
of the gating network’s distribution, we encourage the model to make more
decisive selection of experts while reducing the computational overhead
associated with activating multiple experts.

The final loss we propose is $loss=\mathcal{L}_{\text{cross
entropy}}+\lambda\mathcal{L}_{\text{entropy}}$, and explicitly:

$\text{loss}=-\sum_{t=1}^{T}\sum_{i}y_{t,i}\log(p_{t,i})+\lambda\left(-\sum_{t=1}^{T}\sum_{j}p_{t,j}\log(p_{t,j})\right)$
(9)

Here $\lambda$ is a hyperparameter, and the first term denotes the standard
cross-entropy loss, and the second term denotes the entropy loss. The top-2
gating mechanism inherently provides a form of redundancy, which can be
beneficial for robustness and handling uncertainty. Moving to a more peaky,
top-1 distribution could reduce this redundancy, potentially making the model
more sensitive to errors in the expert selection.

###  3.4 SEER-MoE: Sparse Expert Efficiency through Regularization for
Mixture-of-Experts

We introduce SEER-MoE (Sparse Expert Efficiency through Regularization for
Mixture-of-Experts), a two-stage approach designed to enhance the
computational efficiency of pslightlyed MoE models, specifically targeting the
Mixtral 8x7b (but is also applicable to other MoE models) such as Grok-1:

Stage 1: MoE-specific sparsification decreases the total number of available
experts ($M_{l}$) in each MoE layer via effectively pruning the less
significant ones.

Stage 2: Top-K adaptation i regularization techniques during fine-tuning,
encouraging the model to rely on fewer experts without compromising the
quality of the learned representations while better utilizing the information
from the experts that have been pruned in Stage 1.

Finally, considering the total Model Parameter/FLOPs analysis from Section
3.1, by reducing the number of experts by 25% and using a single activated
expert, we reduce the model parameters by $\approx 25\%$ and FLOPs by $\approx
27\%$ while only slightly degrading model performance.

##  4 Experiments

This section outlines the experimental setup, including model configurations,
training details, and evaluation metrics used to assess the performance of our
proposed SEER-MoE method.

###  4.1 Experimental details

We utilized the Mixtral 8x7b model for all our experiments, chosen for its
superior performance over OpenMoE in our initial tests (for more quality
results, check Appendix C). Our fine-tuning employed QLoRA on 8 NVIDIA A100
GPUs (80GB), running for 1000 steps (unless specified otherwise) using Adam
optimizer with a weight decay of 0.01.

Distributed training (Technical challenges with sharding). Despite facing
challenges with distributed training and sharding, we managed to optimize our
setup for efficient training. We utilized QLoRA fine-tuning to minimize memory
usage, crucial for our hardware constraints. We also explored PyTorch’s Fully
Sharded Data Parallelism (FSDP) Zhao et al. (2023) as a viable option but
given that QLoRa was showing good results with focused on QLoRA. We
implemented the full distributed training loop using the Gugger et al. (2022)
library.

###  4.2 Data

We chose MMLU (Hendrycks et al., 2020) for multitask language understanding
with 57 subjects and Stanford Sentiment Treebank-5 (SST5) dataset (Socher et
al., 2013) for sentiment classification.

MMLU dataset contains questions, choices, and correct labels. The question and
choice pairs are formulated to prompt the following format in Appendix B. The
answer will be extracted from model output and compared against the correct
label to determine output accuracy, which is used to evaluate various expert
sparsification strategies (counting, pruning, subject-specific masking)
quantitatively.

SST5 dataset contains texts and sentiment labels (very negative, negative,
neutral, positive, very positive). The text is applied to the instruction
template from Appendix B to form the prompt to the model. Accuracy will be
computed by comparing the extracted label from the model’s text output with
the ground-truth label. We use the the SST5 training set ($\approx$ 8.5k
examples) for fine-tuning model on different configurations (expert sparsity,
top-K, fine-tuning precedure) and evaluate on the validation set ($\approx$
1.1k examples).

###  4.3 Evaluation Method

The evaluation metrics we used are result accuracy and reduced computation
(FLOPs) and memory footprint.

Accuracy is measured by the percentage of answers extracted from model output
matches with the validation dataset’s labels. For MMLU task, we calculate the
accuracy from number of correct answer among all questions and total number of
1531 questions in validation set. For SST task, we calculate the accuracy from
number of correctly assigned sentiment with the total number of 1101 texts in
validation set.

FLOPs and Memory Reduction is calculated based on model sparsification and
ablation configuration. The FLOPs per token of the pruned model are calcuated
based on Table 4. With expert removal, the model would contain less experts
thus less parameters. The memory reduction can be calculated by number of
expert reduction at each layer.

##  5 Results and Analysis

###  5.1 Expert Sparsification with Heavy-hitters Counting

We start by evaluating the effectiveness of our proposed pruning strategy to
reduce the storage and memory footprint of MoE models with Heavy-hitters
Counting.

Baselines We have 3 baselines to compare against: Dense baseline without
expert, randomly pruning the experts and the state-of-the-art expert pruning
stategy proposed by Lu et al. (2024).

For clarity, we report the results with the dense model numbers serving as the
reference to indicate how much pruning affects the model. In terms of the
baseline statistics of the dense model, from our evaluation on MMLU, we get a
Mean Accuracy of $60.55\%$ over all subjects and a Memory Usage of $86$GB with
$bf16$. Notice that our dense accuracy of Mixtral on MMLU does not exactly
match the numbers reported in Lu et al. (2024) due to potential hardsware
mismatch and prompt difference. However, we are comparing with the accuracy
drop from the respective dense baselines which ensures a fair comparison on
the same scale.

All reported results are zero-shot without fine-tuning.

Comparison We demonstrate the comparison of our top-confirming configuration
with the baselines in Table 1. As observed, at either $25\%$ or $50\%$ expert
sparsity, we significantly surpass all the baselines by a clear margin,
achieving much smaller accuracy drop. For example, compared with the latest
SOTA approach Lu et al. (2024), at the same $25\%$ expert sparsity, we only
lose $\mathbf{3.85}$ accuracy points, almost halve the accuracy drop of $6.40$
of Lu et al. (2024).

Moreover, we could observe that with the expert sparsity introduced into the
model, we successfully reduce the memory footprint for loading Mixtral by
$24\%$ with $25\%$ expert sparsity and $45\%$ with $50\%$ expert sparsity.
While it is not possible to fully load the entire dense Mixtral $8x7b$ model
on a single $80GB$ $A100$ GPU(Dense: $86GB$), our proposed apporach provides a
viable solution at the compensation of minimized accuracy loss by pruning
unimportant experts from the model guided by Heavy Hitters Counting. We are
also going to show later that this accuracy drop could be further minimized
with Task-Specific Finetuning.

Method Total Expert Sparsity $\uparrow$ Accuracy Drop from Dense $\downarrow$
Memory Usage$\downarrow$ Speedup$\uparrow$ Dense $0$ $0$ $\times 1$ $\times 1$
Random $25\%$ $6.17$ $\mathbf{\times 0.76}$ $\mathbf{\times 1.20}$ Lu et al.
(2024) $25\%$ $6.40$ $\mathbf{\times 0.76}$ $\mathbf{\times 1.20}$ Ours $25\%$
$\mathbf{3.85}$ $\mathbf{\times 0.76}$ $\mathbf{\times 1.20}$ Random $50\%$
$15.19$ $\mathbf{\times 0.55}$ $\mathbf{\times 1.27}$ Lu et al. (2024) $50\%$
$16.12$ $\mathbf{\times 0.55}$ $\mathbf{\times 1.27}$ Ours $50\%$
$\mathbf{13.78}$ $\mathbf{\times 0.55}$ $\mathbf{\times 1.27}$

Table 1: Comparison with baseline approaches. Ours achieves the minimized
accuracy drop from the dense baseline at all expert sparsity levels. Notably,
we beat the state-of-the-art Lu et al. (2024) with a clear margin.

Method Counting Strategy Pruning Strategy Subject-Specific Pruning Accuracy
Drop from Dense $\downarrow$ Dense n/a n/a n/a $0$ Ours Activation Layer Yes
$16.59$ Ours Activation Global Yes $14.63$ Ours Soft Layer Yes $15.54$ Ours
Soft Global Yes $7.90$ Ours Soft Layer No $12.80$ Ours Soft Global No
$\mathbf{3.85}$

Table 2: Ablation results of ours. Using soft counting with global pruning and
no subject-specific mask yields the best result.

Ablation Here, we will compare different heavy hitters counting option and
expert removal strategies. Recall that we could use either actual activation
counts or soft counting to gather count data and conduct either layer pruning
or global pruning to actually remove the experts. Moreover, for MMLU
containing $57$ total subjects, we also study whether using a Subject-Specific
Pruning strategy benefits the performance. Concretely, with Subject-Specific
Pruning, we are going to gather counting statistics and perform pruning for
each subject independently, which creates a unique expert mask for solving
each subject. Without this option, we gather counting numbers from samples of
all subjects and adopt the same expert mask to solve all subjects. We
demonstrate comprehensive ablation results in the following Table 2.

From the table, we could make the following three observations:

Observation #1: Global pruning works better than layer pruning. Given our
count statistics are uniform in magnitude across all layers, global pruning
offers more flexible solutions without the layer constraint. Expectedly, we
see the results are better with it.

Observation #2: Using soft counting works better than actual activation
counts. Recall that soft counting, by directly accumulating the softmax
probability for each expert, gives us a sense of the confidence in selecting
each expert to better cope with the scenarios when certain expert barely wins
over others. This is validated by the superiorty in results shown in the
table.

Figure 2: Heavy Hitters Counting Heatmap with Mixtral 8x7b on MMLU.

Observation #3: Subject-specific pruning does not help to gain better
performance on MMLU. Although we expect the subject-specific pruning could
offer improvements by varying the expert mask adaptively based on the subject,
the results suggest otherwise. There could be a theoretical support for this
phenomenon. Recall that as discussed in Sec. 3.2, heavy-hitters counting is
equivalently performing a maginal probability Monte-Carlo estimation. Without
subject-specific pruning, more samples are used to build this estimation,
which could make it more solid.

Visualization and Analysis Here, we provide a heatmap visualization of the
collected heavy-hitter statistics of Mixtral on MMLU in Fig. 2. As observed in
the figure, we could observe that some experts are heavily activated and
leveraged during inference for example Expert $\\#2$ from Layer $26$ and $30$;
whereas some experts are barely activated for example Expert $\\#7$ from Layer
$22$ and $23$. From this discrepancy of acitvation patterns, we could see why
heavy-hitters counting could serve as an effective guidance for pruning
experts.

###  5.2 Finetuning

We assessed the efficacy of various strategies for reducing the number of
activated experts within the Mixtral 8x7b model by fine-tuning on a sentiment
classification task, SST5. The model was finetuned using the training set and
evaluated on the evaluation set via text generation employing greedy decoding.
For details regarding the prompts we used, refer to Appendix B.

We established a baseline by fine-tuning with top-2 routing to gauge the
performance delta attributed to the reduced number of activated experts. This
was further compared with a baseline in-context learning approach. We also
evaluate the different procedures proposed in Section 3.3. Table 3 summarizes
our findings.

Our results reveal a performance gap when using a single expert as opposed to
two, with zero-shot accuracy dropping by 8.2 percentage points. However,
through QLoRA fine-tuning (FT), this gap is mostly bridged, yielding a similar
accuracy of 50.8%. Notably, employing two experts still holds a marginal
advantage, as indicated by the 53.6% accuracy rate post fine-tuning.
Nevertheless, when allowed two experts, the QLoRA approach demonstrated a
slight improvement over the zero-shot baseline.

Method Top-K SST5 acc $\uparrow$ Zero-shot 1 42.6% Zero-shot 2 50.8% QLoRA FT
1 50.8% QLoRA FT 2 53.6% QLoRA on router 1 diverges QLoRA on router 2 51.4% FT
+ Entropy loss, $\lambda=1$ 1 45.5 FT + Entropy loss, $\lambda=0.1$ 1 50.7% \+
Annealing Top-K 1 48.6% \+ Annealing Top-K (more steps) 1 51.4% \+ Annealing
Top-K \+ Entropy Loss 1 51.8%

Table 3: Comparison of fine-tuning strategies and their impact on SST5
accuracy.

These findings suggest that fine-tuning can recover some of the losses
attributed to reducing the number of activated experts. Finally, the approach
that works the best combines Annealing with Entropy minimization approach,
performing better than naive Top-1 finetuning and only 1.8% less than Top-2
finetuning.

###  5.3 SEER-MoE: putting everything together

We now explore whether combining both approaches yield additional gains. We
first sparsify the experts via the heavy-hitters counting and then finetune.
Our goal with this is to understand how we can reduce FLOPs utilization
without losing performance.

Counting / Pruning Strategy Top-K FT Method Top-K SST5 acc $\uparrow$
Activation / Global (25%) QLoRA FT 2 49.0% Activation / Global (25%) QLoRA FT
1 47.5% Soft / Global (25%) QLoRA FT 1 46.7% Activation / Global (25%) QLoRA
FT \+ Annealing Top-K \+ Entropy Loss 1 48.0%

Figure 3: Full stage approach results on SST5.

Notably, our SEER-MoE approach, achieved a competitive accuracy of 48.0% while
only activating a single expert. This highlights the potential of SEER-MoE to
maintain high model performance even under significantly reduced computational
overhead. Remarkably, the accuracy attained mirrors that of the two-expert
configuration following only QLoRA fine-tuning, underscoring the effectiveness
of our comprehensive strategy in reducing FLOPs without detriment to accuracy.

Furthermore, it is evident from the results that the activation-based pruning
combined with a single-expert QLoRA fine-tuning confers a substantial accuracy
gain over the soft pruning approach. This suggests that the more targeted
activation-based pruning method combines effectively with our fine-tuning
paradigm.

##  6 Conclusion

Our SEER-MoE framework effectively mitigates some of the computational
inefficiencies of Mixture-of-Experts (MoE) models with minor compromise to
performance. Through a two-stage process that includes expert sparsification
and Top-K adaptation via fine-tuning, we have significantly cut down both
FLOPs and memory usage for Mixtral 8x7b. Testing on SST5 and MMLU benchmarks
shows that SEER-MoE achieves strong performance while reducing the number of
active experts and parameters, making MoE models more viable across various
applications and hardware constraints.

## References

  * et al (2024) Machel Reid et al. 2024.  Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context.
  * Dettmers et al. (2023) Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. 2023.  Qlora: Efficient finetuning of quantized llms.
  * Eigen et al. (2014) David Eigen, Marc’Aurelio Ranzato, and Ilya Sutskever. 2014.  Learning factored representations in a deep mixture of experts.
  * Fedus et al. (2021) William Fedus, Barret Zoph, and Noam Shazeer. 2021.  Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.  _CoRR_ , abs/2101.03961.
  * Gugger et al. (2022) Sylvain Gugger, Lysandre Debut, Thomas Wolf, Philipp Schmid, Zachary Mueller, Sourab Mangrulkar, Marc Sun, and Benjamin Bossan. 2022.  Accelerate: Training and inference at scale made simple, efficient and adaptable.  https://github.com/huggingface/accelerate.
  * Hendrycks et al. (2020) Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. 2020.  Measuring massive multitask language understanding.  _arXiv preprint arXiv:2009.03300_.
  * Jacobs et al. (1991) Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. 1991.  Adaptive mixtures of local experts.  _IEEE Neural Computation_.
  * Jiang et al. (2024) Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel, Guillaume Bour, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, Marie-Anne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. 2024.  Mixtral of experts.
  * Kaplan et al. (2020) Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. 2020.  Scaling laws for neural language models.  _CoRR_ , abs/2001.08361.
  * Lu et al. (2024) Xudong Lu, Qi Liu, Yuhui Xu, Aojun Zhou, Siyuan Huang, Bo Zhang, Junchi Yan, and Hongsheng Li. 2024.  Not all experts are equal: Efficient expert pruning and skipping for mixture-of-experts large language models.
  * Socher et al. (2013) Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts. 2013.  Recursive deep models for semantic compositionality over a sentiment treebank.  In _Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing_ , pages 1631–1642, Seattle, Washington, USA. Association for Computational Linguistics.
  * xAI (2024) xAI. 2024.  Open release of grok-1.  https://x.ai/blog/grok-os.  Accessed: March 17th, 2024.
  * Xue et al. (2024) Fuzhao Xue, Zian Zheng, Yao Fu, Jinjie Ni, Zangwei Zheng, Wangchunshu Zhou, and Yang You. 2024.  Openmoe: An early effort on open mixture-of-experts language models.  _arXiv preprint arXiv:2402.01739_.
  * Zhang et al. (2023) Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, Zhangyang Wang, and Beidi Chen. 2023.  H2o: Heavy-hitter oracle for efficient generative inference of large language models.
  * Zhao et al. (2023) Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, and Shen Li. 2023.  Pytorch fsdp: Experiences on scaling fully sharded data parallel.
  * Zoph et al. (2022) Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. 2022.  St-moe: Designing stable and transferable sparse expert models.

##  Appendix A Parameters/FLOPs per Token for MoE Transformers

Table 4 describes the amount of parameters and FLOPs per token for a
Transformer MoE model.

Operation | Parameters | FLOPs per Token
---|---|---
Embed | $(n_{vocab}+n_{ctx})d_{model}$ | $4d_{model}$
Attention: QKV | $n_{layer}d_{model}3d_{attn}$ | $2n_{layer}d_{model}3d_{attn}$
Attention: Mask | — | $2n_{layer}n_{ctx}d_{attn}$
Attention: Project | $n_{layer}d_{attn}d_{model}$ | $2n_{layer}d_{attn}d_{embd}$
MoE Feedforward | $n_{experts}n_{layer}2d_{model}d_{ff}$ | $2n_{topk}n_{layer}2d_{model}d_{ff}$
MoE Gating | $n_{experts}n_{layer}d_{model}$ | $2n_{experts}n_{layer}d_{model}$
De-embed | — | $2d_{model}n_{vocab}$
Table 4: Parameter counts and compute (forward pass) estimates for a MoE
Transformer model. Nonlinearities, biases, and layer normalization are
omitted.

##  Appendix B Prompts

In this section we describe the different prompts that we use for different
tasks to query the desired answer. We have experimented with different prompts
and the ones presented below are chosen due to better performance or that
match previously reported benchmarks.

SST5 For Mixtral:

    [INST] You are a helpful assistant. Your task is of sentiment classification.
    Categorize the following text as either "very negative", "negative",
     "neutral", "positive" or "very positive":
    \{TEXT\}
    Only generate the label, without explanations:[/INST]

SST5 For OpenMoE:

    <<SYS>> You are a helpful assistant. Your task is of sentiment classification. <</SYS>>
    <s>[INST] Categorize the following text in one of the following sentiments
    ’very negative’, ’negative’, ’neutral’, ’positive’ or ’very positive’:
    \{TEXT\} [/INST]

MMLU For Mixtral 8x7b:

    [INST] The following are multiple choice questions (with answers) about \{SUBJECT\}.

    \{QUESTION\}

    \{CHOICES\}

    Only respond with the letter of the correct answer and no explanation.
    Answer:
    [/INST]

MMLU For OpenMoE:

    <<SYS>> You are a helpful assistant. Your task is of multiple choice question answering
    based on your knowledge.
    For subject \{SUBJECT\}
    Choose the best answer (A), (B), (C), or (D)
     to the following question without explanation:<</SYS>>

    <s>[INST] \{QUESTION\} from the following choices:

    \{CHOICES\}

    [/INST]

##  Appendix C OpenMoE results

We also experimented with OpenMoE Xue et al. (2024) models but we were not
able to get reasonable results. We evaluate the OpenMoE-8B-Chat model on MMLU
and SST5 and the results can be seen in Table 5.

Model | Setup | SST5 Acc. | MMLU Acc.
---|---|---|---
OpenMoE-8B-Chat (1.1T+SFT) | expert-topk=1 | 35.0% | 25.7%
OpenMoE-8B-Chat (1.1T+SFT) | expert-topk=2 | 35.6% | 26.5%
Table 5: Evaluation on MMLU and SST5 for different for OpenMoE models with
top-1 and top-2.

Considering the random-guessing baseline for MMLU to be 25% and for SST5 to be
20%, we did not further pursue utilizing these models for additional
experiments.