* Jaeger Jaeger, M. (2007).  Parameter learning for relational bayesian networks.  In Proceedings of the 24th international conference on Machine learning, pp. 369–376.
  * Kazemi, et al. Kazemi, S. M., Kimmig, A., Van den Broeck, G., & Poole, D. (2016).  New liftable classes for first-order probabilistic inference.  In Proceedings of the 30th International Conference on Neural Information Processing Systems, NIPS’16, p. 3125–3133.
  * Kersting & De Raedt Kersting, K., & De Raedt, L. (2001).  Adaptive bayesian logic programs.  In Proceedings of the 11th International Conference on Inductive Logic Programming, ILP ’01, p. 104–117, Berlin, Heidelberg. Springer-Verlag.
  * Kersting & De Raedt Kersting, K., & De Raedt, L. (2007).  Bayesian logic programming: Theory and tool.  In Introduction to Statistical Relational Learning. MIT Press.
  * Kisynski & Poole Kisynski, J., & Poole, D. (2009).  Lifted aggregation in directed first-order probabilistic models..  In Proceedings of the 21st International Joint Conference on Artificial Intelligence, IJCAI’09, pp. 1922–1929.
  * Koller & Friedman Koller, D., & Friedman, N. (2009).  Probabilistic graphical models: principles and techniques.  MIT press.
  * Koller, et al. Koller, D., Friedman, N., Džeroski, S., Sutton, C., McCallum, A., Pfeffer, A., Abbeel, P., Wong, M.-F., Heckerman, D., Meek, C., et al. (2007).  Introduction to statistical relational learning.  MIT press.
  * Kowalski Kowalski, R. (1974).  Predicate logic as programming language.  In IFIP congress, Vol. 74, pp. 569–544.
  * Kumar & Kuželka Kumar, N., & Kuželka, O. (2021).  Context-specific likelihood weighting.  In International Conference on Artificial Intelligence and Statistics, pp. 2125–2133. PMLR.
  * Kumar, Kuželka, & De Raedt Kumar, N., Kuželka, O., & De Raedt, L. (2021).  Learning distributional programs for relational autocompletion.  Theory and Practice of Logic Programming, 21, 1–34.
  * Li & Russell Li, L., & Russell, S. J. (2013).  The blog language reference.  Tech. rep., Technical Report UCB/EECS-2013–51, EECS Department, University of California.
  * Mateescu & Dechter Mateescu, R., & Dechter, R. (2009).  Mixed deterministic and probabilistic networks.  Annals of Mathematics and Artificial Intelligence, 54, 3–51.
  * Michels, Hommersom, & Lucas Michels, S., Hommersom, A., & Lucas, P. J. F. (2016).  Approximate probabilistic inference with bounded error for hybrid probabilistic logic programming.  In Proceedings of the Twenty-Fifth International Joint Conference on Artificial Intelligence, IJCAI’16, p. 3616–3622. AAAI Press.
  * Milch, et al. Milch, B., Marthi, B., Sontag, D., Russell, S., Ong, D. L., & Kolobov, A. (2005).  Approximate inference for infinite contingent bayesian networks.  In International Workshop on Artificial Intelligence and Statistics, AISTATS’05, pp. 238–245. PMLR.
  * Milch & Russell Milch, B., & Russell, S. J. (2009).  Extending bayesian networks to the open-universe case.  In R. Dechter, H. Geffner, and J. Y. Halpern, eds. Heuristics, Probability and Causality: A Tribute to Judea Pearl.
  * Minka, et al. Minka, T., Winn, J., Guiver, J., & Knowles, D. (2010).  Infer .net 2.4, 2010. microsoft research cambridge..
  * Natarajan, et al. Natarajan, S., Tadepalli, P., Dietterich, T. G., & Fern, A. (2008).  Learning first-order probabilistic models with combining rules.  Annals of Mathematics and Artificial Intelligence, 54(1), 223–256.
  * Ngo & Haddawy Ngo, L., & Haddawy, P. (1995).  Probabilistic logic programming and bayesian networks.  In Asian computing science conference, pp. 286–300. Springer.
  * Ngo & Haddawy Ngo, L., & Haddawy, P. (1997).  Answering queries from context-sensitive probabilistic knowledge bases.  Theoretical Computer Science, 171(1-2), 147–177.
  * Niepert Niepert, M. (2012).  Lifted Probabilistic Inference: An MCMC Perspective.  Universitätsbibliothek Mannheim.
  * Nilsson & Małuszyński Nilsson, U., & Małuszyński, J. (1995).  Logic, programming and Prolog.  Wiley Chichester.
  * Nitti, De Laet, & De Raedt Nitti, D., De Laet, T., & De Raedt, L. (2016).  Probabilistic logic programming for hybrid relational domains.  Machine Learning, 103(3), 407–449.
  * Poole Poole, D. (1993).  Probabilistic horn abduction and bayesian networks.  Artificial intelligence, 64(1), 81–129.
  * Poole Poole, D. (1997).  Probabilistic partial evaluation: Exploiting rule structure in probabilistic inference.  In Proceedings of the Fifteenth International Joint Conference on Artifical Intelligence - Volume 2, IJCAI’97, p. 1284–1291.
  * Poole Poole, D. (2003).  First-order probabilistic inference.  In Proceedings of the 18th International Joint Conference on Artificial Intelligence, IJCAI’03, p. 985–991.
  * Poole Poole, D. (2008).  The independent choice logic and beyond.  In Probabilistic inductive logic programming, pp. 222–243. Springer.
  * Poole & Zhang Poole, D., & Zhang, N. L. (2003).  Exploiting contextual independence in probabilistic inference.  Journal of Artificial Intelligence Research, 18, 263–313.
  * Poole & Mackworth Poole, D. L., & Mackworth, A. K. (2010).  Artificial Intelligence: foundations of computational agents.  Cambridge University Press.
  * Sato & Kameya Sato, T., & Kameya, Y. (1997).  Prism: A language for symbolic-statistical modeling.  In Proceedings of the Fifteenth International Joint Conference on Artifical Intelligence - Volume 2, IJCAI’97, p. 1330–1335.
  * Shachter Shachter, R. (1998).  Bayes-ball: The rational pastime (for determining irrelevance and requisite information in belief networks and influence diagrams).  In Proceedings of the Thirteenth Conference on Uncertainty in Artificial Intelligence, UAI’97, p. 480–487.
  * Shen, Choi, & Darwiche Shen, Y., Choi, A., & Darwiche, A. (2016).  Tractable operations for arithmetic circuits of probabilistic models.  In Proceedings of the 30th International Conference on Neural Information Processing Systems, NIPS’16, p. 3943–3951.
  * Taghipour, et al. Taghipour, N., Fierens, D., Davis, J., & Blockeel, H. (2013).  Lifted variable elimination: Decoupling the operators from the constraint language.  Journal Artificial Intelligence Research, 47(1), 393–439.
  * Team Team, S. D. (2015).  Stan: A c++ library for probability and sampling, version 2.5. 0..
  * Van Den Broeck Van Den Broeck, G. (2011).  On the completeness of first-order knowledge compilation for lifted probabilistic inference.  In Proceedings of the 24th International Conference on Neural Information Processing Systems, NIPS’11, p. 1386–1394.
  * Van Den Broeck & Darwiche Van Den Broeck, G., & Darwiche, A. (2013).  On the complexity and approximation of binary evidence in lifted inference.  In Proceedings of the 26th International Conference on Neural Information Processing Systems, NIPS’13, pp. 81–86.
  * Van Den Broeck, et al. Van Den Broeck, G., Taghipour, N., Meert, W., Davis, J., & De Raedt, L. (2011).  Lifted probabilistic inference by first-order knowledge compilation.  In Proceedings of the Twenty-Second International Joint Conference on Artificial Intelligence, IJCAI’11, p. 2178–2185. AAAI Press.
  * Vennekens, Denecker, & Bruynooghe Vennekens, J., Denecker, M., & Bruynooghe, M. (2009).  Cp-logic: A language of causal probabilistic events and its relation to logic programming.  Theory and practice of logic programming, 9(3), 245–308.
  * Vennekens, Verbaeten, & Bruynooghe Vennekens, J., Verbaeten, S., & Bruynooghe, M. (2004).  Logic programs with annotated disjunctions.  In 20th International Conference on Logic Programming, Vol. 3132 of Lecture Notes in Computer Science, pp. 431–445. Springer.
  * Warren Warren, D. S. (1992).  Memoing for logic programs.  Communications of the ACM, 35, 93–111.
  * Wielemaker, et al. Wielemaker, J., Schrijvers, T., Triska, M., & Lager, T. (2011).  Swi-prolog.  Theory and Practice of Logic Programming, 12, 67 – 96.
  * Wu, et al. Wu, Y., Srivastava, S., Hay, N., Du, S., & Russell, S. (2018).  Discrete-continuous mixtures in probabilistic programming: Generalized semantics and inference algorithms.  In International Conference on Machine Learning, pp. 5343–5352. PMLR.
  * Zeng & Van Den Broeck Zeng, Z., & Van Den Broeck, G. (2020).  Efficient search-based weighted model integration.  In Proceedings of The 35th Uncertainty in Artificial Intelligence Conference, pp. 175–185. PMLR.
  * Zhang & Poole Zhang, N. L., & Poole, D. L. (1996).  Exploiting causal independence in bayesian network inference.  Journal of Artificial Intelligence Research, 5, 301–328.

##  A Missing Proofs

###  A.1 Proof of Proposition 2

In this section, we present the detailed proof of Proposition 2.

###### Proof.

First, we show that if $\mathtt{rv(A)}$ is in the least Herbrand model of
$rv(\mathbb{P})$ (denoted by $M_{rv(\mathbb{P})}$) then $\mathbb{P}$ defines
$\mathtt{A}$ as an RV: Since $\mathbb{P}$ is well-defined, there will be at
least one clause for $\mathtt{A}$ in $ground(\mathbb{P})_{\omega}$ whose body
is true with respect to $\omega$ due to the exhaustiveness condition. So,
$\mathbb{P}$ defines $\mathtt{A}$ as an RV.

Now, we show the converse: Let $\omega$ be an assignment where each
$\mathtt{B}$ is assigned a value if $\mathtt{rv(B)}$ is in
$M_{rv(\mathbb{P})}$. Given $\omega$, we construct a ground program
$ground(\mathbb{P})_{\omega}$ such that there is a clause $\mathtt{A\sim
D\leftarrow Body}$ in $ground(\mathbb{P})_{\omega}$ with atoms
$\\{\mathtt{B_{1}\cong V_{1},\dots,B_{m}\cong V_{m}}\\}$ in $\mathtt{Body}$
and these atoms are true with respect to $\omega$. Then, due to the way DC#
programs have been defined in Definition 5, there must be a definite clause
$\mathtt{rv(A)\leftarrow rv(B_{1}),\dots,rv(B_{m})}$ in
$ground(rv(\mathbb{P}))$. Now, $\\{\mathtt{rv(B_{1}),\dots,rv(B_{m})}\\}$ are
in $M_{rv(\mathbb{P})}$ since those RVs are assigned in $\omega$. So,
$\mathtt{rv(A)}$ will be in $M_{rv(\mathbb{P})}$ since the body of the
definite clause will be true. This completes the proof. ∎

###  A.2 Proof of Proposition 3

In this section, we present the detailed proof of Proposition 3.

###### Proof.

First, we show that if $\mathtt{A}$ directly influences $\mathtt{B}$ then
$\mathtt{pa(B,A)}$ is in the least Herbrand model of $dep(\mathbb{P})$: Since
$\mathtt{A}$ directly influences $\mathtt{B}$ there must be a clause
$\mathcal{C}\equiv\mathtt{B\sim D\leftarrow B_{1},\dots,B_{n}}\in
ground(\mathbb{P})_{\omega}$ for some possible world $\omega$ such that
$\mathtt{n>0}$, $\mathtt{A}$ is a RV term in the body of $\mathcal{C}$, and
the body is true with respect to $\omega$. This implies that there must be a
clause $\mathtt{rv(B)\leftarrow rv(A_{1}),\dots,rv(A_{m})}$ in
$ground(rv(\mathbb{P}))$ such that $\\{\mathtt{A_{1},\dots,A_{m}}\\}$ is the
set of RV terms in the body of $\mathcal{C}$ and
$\mathtt{A\in\\{\mathtt{A_{1},\dots,A_{m}}\\}}$. Consequently, there must be a
definite clause $\mathcal{C}^{\prime}\equiv\mathtt{pa(B,A)\leftarrow
rv(B),rv(A_{1}),\dots,rv(A_{m})}$ in $ground(dep(\mathbb{P}))$ by
construction. $\mathtt{\\{A_{1},\dots,A_{m}\\}}$ is assigned in $\omega$ so
due to Proposition 2:
$rv(\mathbb{P})\models\mathtt{(rv(A_{1}),\dots,rv(A_{m}))}$, consequently,
$rv(\mathbb{P})\models\mathtt{rv(B)}$. This is also true in $dep(\mathbb{P})$
since all clauses in $rv(\mathbb{P})$ are in $dep(\mathbb{P})$. So, the body
of clause $\mathcal{C}^{\prime}$ will be true and
$\mathtt{dep(\mathbb{P})\models pa(B,A)}$.

Now, we show the converse: Since $dep(\mathbb{P})\models\mathtt{pa(B,A)}$,
there must be a definite clause $\mathtt{pa(B,A)\leftarrow
rv(B),rv(A_{1}),\dots,rv(A_{m})}$ in $ground(dep(\mathbb{P}))$ such that
$\mathtt{A\in\\{A_{1},\dots,A_{m}\\}}$ and the clause’s body is true in
$ground(dep(\mathbb{P}))$. This implies that there must be a definite clause
$\mathtt{rv(B)\leftarrow rv(A_{1}),\dots,rv(A_{m})}\in ground(rv(\mathbb{P}))$
whose body is true in $ground(rv(\mathbb{P}))$. Due to the way DC# programs
are defined (Definition 5), there must be a distributional clause
$\mathtt{B\sim D\leftarrow Body}\in ground(\mathbb{P})_{\omega}$ for some
$\omega$ such that $\mathtt{\\{A_{1}\cong V_{1},\dots,A_{m}\cong V_{m}\\}}$
belongs to $\mathtt{Body}$ that is true in $\omega$. So, $\mathtt{A}$ directly
influences $\mathtt{B}$, which completes the proof.

∎

###  A.3 Proof of Lemma 1

In this section, we present the detailed proof of Lemma 1.

###### Proof.

Let us denote the variables in $\mathbf{Z}$ that are marked on the top
(requisite) by $\mathbf{Z}_{\star}$ and that are not marked on the top (not
requisite) by $\mathbf{Z}_{\bar{\star}}$. The required probability $\mu$ is
then given by,

$\displaystyle\mu=P(\mathbf{x}_{q}\mid\mathbf{e})=\frac{\sum_{\mathbf{x},\mathbf{z}_{\star},\mathbf{z}_{\bar{\star}}}P(\mathbf{x},\mathbf{z}_{\star},\mathbf{z}_{\bar{\star}},\mathbf{e})f(\mathbf{x})}{\sum_{\mathbf{x},\mathbf{z}_{\star},\mathbf{z}_{\bar{\star}}}P(\mathbf{x},\mathbf{z}_{\star},\mathbf{z}_{\bar{\star}},\mathbf{e})}=\frac{\sum_{\mathbf{x},\mathbf{z}_{\star}}P(\mathbf{x},\mathbf{z}_{\star},\mathbf{e})f(\mathbf{x})\sum_{\mathbf{z}_{\bar{\star}}}P(\mathbf{z}_{\bar{\star}}\mid\mathbf{x},\mathbf{z}_{\star},\mathbf{e})}{\sum_{\mathbf{x},\mathbf{z}_{\star}}P(\mathbf{x},\mathbf{z}_{\star},\mathbf{e})\sum_{\mathbf{z}_{\bar{\star}}}P(\mathbf{z}_{\bar{\star}}\mid\mathbf{x},\mathbf{z}_{\star},\mathbf{e})}$

Since
$\sum_{\mathbf{z}_{\bar{\star}}}P(\mathbf{z}_{\bar{\star}}\mid\mathbf{x},\mathbf{z}_{\star},\mathbf{e})=1$,
we can write,

$\displaystyle\mu=\frac{\sum_{\mathbf{x},\mathbf{z}_{\star}}P(\mathbf{x},\mathbf{z}_{\star},\mathbf{e})f(\mathbf{x})}{\sum_{\mathbf{x},\mathbf{z}_{\star}}P(\mathbf{x},\mathbf{z}_{\star},\mathbf{e})}$

Now let us denote the observed variables in $\mathbf{E}$ that are visited
(requisite) by $\mathbf{E}_{r}$ and those that are not visited (not requisite)
by $\mathbf{E}_{n}$. We can write,

$\displaystyle\mu=\frac{\sum_{\mathbf{x},\mathbf{z}_{\star}}P(\mathbf{x},\mathbf{z}_{\star},\mathbf{e}_{r})P(\mathbf{e}_{n}\mid\mathbf{x},\mathbf{z}_{\star},\mathbf{e}_{r})f(\mathbf{x})}{\sum_{\mathbf{x},\mathbf{z}_{\star}}P(\mathbf{x},\mathbf{z}_{\star},\mathbf{e}_{r})P(\mathbf{e}_{n}\mid\mathbf{x},\mathbf{z}_{\star},\mathbf{e}_{r})}$

The variables in $\mathbf{X}\cup\mathbf{Z}_{\star}$ pass the Bayes-balls to
all their parents and all their children, but $\mathbf{E}_{n}$ is not visited
by these balls. The correctness of the Bayes-ball algorithm ensures that there
is no active path from $\mathbf{X}\cup\mathbf{Z}_{\star}$ to any $E_{n}$ in
$\mathbf{E}_{n}$ given $\mathbf{E}_{r}$. Thus
$\mathbf{X},\mathbf{Z}_{\star}\perp\mathbf{E}_{n}\mid\mathbf{E}_{r}$ and
$P(\mathbf{e}_{n}\mid\mathbf{x},\mathbf{z}_{\star},\mathbf{e}_{r})=P(\mathbf{e}_{n}\mid\mathbf{e}_{r})$.
After cancelling out the common term $P(\mathbf{e}_{n}\mid\mathbf{e}_{r})$, we
get,

$\displaystyle\mu=\frac{\sum_{\mathbf{x},\mathbf{z}_{\star}}P(\mathbf{x},\mathbf{z}_{\star},\mathbf{e}_{r})f(\mathbf{x})}{\sum_{\mathbf{x},\mathbf{z}_{\star}}P(\mathbf{x},\mathbf{z}_{\star},\mathbf{e}_{r})}$

Now let us denote observed variables in $\mathbf{E}_{r}$ that are only visited
by $\mathbf{E}_{\smwhitestar}$ and that are visited as well as marked on top
by $\mathbf{E}_{\star}$. After cancelling out the common term, we get the
desired result,

$\displaystyle\mu=\frac{\sum_{\mathbf{x},\mathbf{z}_{\star}}P(\mathbf{x},\mathbf{z}_{\star},\mathbf{e}_{\star}\mid\mathbf{e}_{\smwhitestar})P(\mathbf{e}_{\smwhitestar})f(\mathbf{x})}{\sum_{\mathbf{x},\mathbf{z}_{\star}}P(\mathbf{x},\mathbf{z}_{\star},\mathbf{e}_{\star}\mid\mathbf{e}_{\smwhitestar})P(\mathbf{e}_{\smwhitestar})}=\frac{\sum_{\mathbf{x},\mathbf{z}_{\star}}P(\mathbf{x},\mathbf{z}_{\star},\mathbf{e}_{\star}\mid\mathbf{e}_{\smwhitestar})f(\mathbf{x})}{\sum_{\mathbf{x},\mathbf{z}_{\star}}P(\mathbf{x},\mathbf{z}_{\star},\mathbf{e}_{\star}\mid\mathbf{e}_{\smwhitestar})}$

∎

###  A.4 Proof of Theorem 5

In this section, we present the detailed proof of Theorem 5.

###### Proof.

The expectation $\mathbb{E}_{Q_{\star}}[W_{\mathbf{\dot{e}}_{\star}}]$ is
given by

$\displaystyle\sum_{\mathbf{x},\mathbf{z}_{\star}}\prod_{u_{i}\in\mathbf{x}\cup\mathbf{z}_{\star}}P(u_{i}\mid\mathbf{pa}(U_{i}))\prod_{v_{i}\in\mathbf{\dot{e}}_{\star}}P(v_{i}\mid\mathbf{pa}(V_{i})).$

The basis $\mathbf{\dot{S}}_{\star}$ is a subset of
$\mathbf{X}\cup\mathbf{Z}_{\star}$ by Definition 16 . Let us denote
$(\mathbf{X}\cup\mathbf{Z}_{\star})\setminus\mathbf{\dot{S}}_{\star}$ by
$\mathbf{Z}_{\diamond}$. We can now rewrite the expectation as follows,

$\displaystyle\sum_{\mathbf{\dot{s}}_{\star},\mathbf{z}_{\diamond}}\prod_{u_{i}\in\mathbf{\dot{e}}_{\star}\cup\mathbf{\dot{s}}_{\star}}P(u_{i}\mid\mathbf{pa}(U_{i}))\prod_{v_{i}\in\mathbf{z}_{\diamond}}P(v_{i}\mid\mathbf{pa}(V_{i})).$

We will show that $Pa\notin\mathbf{Z}_{\diamond}$ for any
$Pa\in\mathbf{Pa}(U_{i})$, which will then allow us to push the summation over
$\mathbf{z}_{\diamond}$ inside. Let us consider two cases:

  * •

For $U_{i}\in\mathbf{\dot{E}}_{\star}$, let $Pa\in\mathbf{Pa}(U_{i})$ be an
unobserved parent of $U_{i}$, then there will be a direct causal trail from
$Pa$ to $U_{i}$, consequently $Pa$ will be in the set
$\mathbf{\dot{S}}_{\star}$.

  * •

For $U_{i}\in\mathbf{\dot{S}}_{\star}$, there will be a causal trail
$U_{i}\rightarrow\cdots\ B_{j}\ \cdots\rightarrow E$ such that
$E\in\mathbf{\dot{E}}_{\star}$ and such that either no $B_{i}$ is observed or
there is no $B_{i}$. Let $Pa\in\mathbf{Pa}(U_{i})$ be an unobserved parent of
$U_{i}$ then there will be a direct causal trail from $Pa$ to $U_{i}$,
consequently, there will be such causal trail from $Pa$ to $E$ and $Pa$ will
be in the set $\mathbf{\dot{S}}_{\star}$.

Hence, we push the summation over $\mathbf{z}_{\diamond}$ inside and use the
fact that
$\sum_{\mathbf{z}_{\diamond}}\prod_{v_{i}\in\mathbf{z}_{\diamond}}P(v_{i}\mid\mathbf{pa}(V_{i}))=1$,
to get the desired result. ∎

###  A.5 Proof of Theorem 6

In this section, we present the detailed proof of Theorem 6.

###### Proof.

Since
$\mathbf{X},\mathbf{Z}_{\star},\mathbf{E}_{\star},\mathbf{E}_{\smwhitestar}$
are variables of the Bayesian network $\mathcal{B}$ and they form a sub-
network $\mathcal{B}_{\star}$ such that $\mathbf{E}_{\smwhitestar}$ do not
have any parent, we can always write,

$\displaystyle
P(\mathbf{x},\mathbf{z}_{\star},\mathbf{e}_{\star}\mid\mathbf{e}_{\smwhitestar})=\prod_{u_{i}\in\mathbf{x}\cup\mathbf{z}_{\dagger}\cup\mathbf{e}_{\dagger}}P(u_{i}\mid\mathbf{pa}(U_{i}))\prod_{v_{i}\in\mathbf{z}_{\ddagger}\cup\mathbf{e}_{\ddagger}}P(v_{i}\mid\mathbf{pa}(V_{i}))$

such that
$p\in\mathbf{x}\cup\mathbf{z}_{\star}\cup\mathbf{e}_{\star}\cup\mathbf{e}_{\smwhitestar}$
for all $p\in\mathbf{pa}(U_{i})$ or $p\in\mathbf{pa}(V_{i})$. Now consider the
summation over all possible assignments of variables in
$\mathbf{X},\mathbf{Z}_{\star}$, that is:
$\sum_{\mathbf{x},\mathbf{z}_{\star}}P(\mathbf{x},\mathbf{z}_{\star},\mathbf{e}_{\star}\mid\mathbf{e}_{\smwhitestar})$.
We can always write,

$\displaystyle\sum_{\mathbf{x},\mathbf{z}_{\star}}P(\mathbf{x},\mathbf{z}_{\star},\mathbf{e}_{\star}\mid\mathbf{e}_{\smwhitestar})=\sum_{\psi\in\Psi}\sum_{\mathbf{z}_{\ddagger}[\psi]}P(\mathbf{x}[\psi],\mathbf{z}_{\dagger}[\psi],\mathbf{z}_{\ddagger}[\psi],\mathbf{e}_{\dagger}[\psi],\mathbf{e}_{\ddagger}[\psi]\mid\mathbf{e}_{\smwhitestar})$
(8)

To simplify notation, from now we denote $\\{\mathbf{x}[\psi]$,
$\mathbf{z}_{\dagger}[\psi]$, $\mathbf{Z}_{\ddagger}[\psi]$,
$\mathbf{e}_{\dagger}[\psi]$, $\mathbf{e}_{\ddagger}[\psi]\\}$ by
$\\{\mathbf{x}$, $\mathbf{z}_{\dagger}$, $\mathbf{Z}_{\ddagger}$,
$\mathbf{e}_{\dagger}$, $\mathbf{e}_{\ddagger}\\}$. After using the definition
of contextual assignments, we have that,

$\displaystyle
P(\mathbf{x},\mathbf{z}_{\dagger},\mathbf{z}_{\ddagger},\mathbf{e}_{\dagger},\mathbf{e}_{\ddagger}\mid\mathbf{e}_{\smwhitestar})=\prod_{u_{i}\in\mathbf{x}\cup\mathbf{z}_{\dagger}\cup\mathbf{e}_{\dagger}}P(u_{i}\mid\mathbf{ppa}(U_{i}))\prod_{v_{i}\in\mathbf{z}_{\ddagger}\cup\mathbf{e}_{\ddagger}}P(v_{i}\mid\mathbf{pa}(V_{i}))$

Since $p\notin\mathbf{z}_{\ddagger}$ for any $p\in\mathbf{ppa}(U_{i})$, we can
push the summation over $\mathbf{z}_{\ddagger}$ inside to get,

$\displaystyle\sum_{\psi\in\Psi}\sum_{\mathbf{z}_{\ddagger}}P(\mathbf{x},\mathbf{z}_{\dagger},\mathbf{z}_{\ddagger},\mathbf{e}_{\dagger},\mathbf{e}_{\ddagger}\mid\mathbf{e}_{\smwhitestar})=\sum_{\psi\in\Psi}\prod_{u_{i}\in\mathbf{x}\cup\mathbf{z}_{\dagger}\cup\mathbf{e}_{\dagger}}P(u_{i}\mid\mathbf{ppa}(U_{i}))\sum_{\mathbf{z}_{\ddagger}}\prod_{v_{i}\in\mathbf{z}_{\ddagger}\cup\mathbf{e}_{\ddagger}}P(v_{i}\mid\mathbf{pa}(V_{i})).$
(9)

However, we get a strange term
$\sum_{\mathbf{z}_{\ddagger}}\prod_{v_{i}\in\mathbf{z}_{\ddagger}\cup\mathbf{e}_{\ddagger}}P(v_{i}\mid\mathbf{pa}(V_{i}))$.
Let $\mathbf{S}_{\ddagger}$ denote the basis of residual
$\mathbf{e}_{\ddagger}$. We have that
$\mathbf{S}_{\ddagger}\subseteq\mathbf{Z}_{\ddagger}$ by Definition 17. Let us
denote $\mathbf{Z}_{\ddagger}\setminus\mathbf{S}_{\ddagger}$ with
$\mathbf{Z}_{\diamond}$. Now the strange term can be rewritten as,

$\displaystyle\sum_{\mathbf{s}_{\ddagger},\mathbf{z}_{\diamond}}\prod_{u_{i}\in\mathbf{e}_{\ddagger}\cup\mathbf{s}_{\ddagger}}P(u_{i}\mid\mathbf{pa}(U_{i}))\prod_{v_{i}\in\mathbf{z}_{\diamond}}P(v_{i}\mid\mathbf{pa}(V_{i})).$

In the proof of Theorem 5, we showed that the summation over variables not in
$\mathbf{S}_{\ddagger}$ can be pushed inside; hence, $\mathbf{Z}_{\diamond}$
can be pushed inside. After using the fact that
$\sum_{\mathbf{z}_{\diamond}}\prod_{v_{i}\in\mathbf{z}_{\diamond}}P(v_{i}\mid\mathbf{pa}(V_{i}))=1$,
we conclude that the strange term is actually the expectation
$\mathbb{E}_{Q_{\star}}[W_{\mathbf{e}_{\ddagger}}]$. Using Equation (3), (8),
(9) and rearranging terms, the result follows. ∎

###  A.6 Proof of Lemma 7

In this section, we present the detailed proof of Lemma 7.

###### Proof.

It is clear that a subset of unobserved variables is assigned. Let
$\mathbf{Z}_{\ddagger}$ be a set of unobserved variables left unassigned. Let
$E\in\mathbf{E}_{\star}$ be an observed variable. Consider two cases:

  * •

All ancestors of $E$ are in
$\mathbf{Z}_{\ddagger}\cup\mathbf{E}_{\smwhitestar}\cup\mathbf{E}_{\star}$.

  * •

Some ancestors of $E$ are in
$\mathbf{Z}_{\ddagger}\cup\mathbf{E}_{\smwhitestar}\cup\mathbf{E}_{\star}$ and
some are in $\mathbf{X}\cup\mathbf{Z}_{\dagger}$. Let
$A\in\mathbf{X}\cup\mathbf{Z}_{\dagger}$ and let $A\rightarrow\cdots\ B_{i}\
\cdots\rightarrow E$ be a causal trail. Some $B_{i}$ are observed in all such
trails.

Clearly, $E$ will not be visited from any parent in the first case, and in the
second case, the visit will be blocked by observed variables. Consequently,
$E$ will not be weighted, which completes the proof. ∎

###  A.7 Proof of Theorem 8

In this section, we present the detailed proof of Theorem 8.

###### Proof.

Variables in $\mathbf{Z}_{\ddagger}$ are not assigned in the simulation;
hence, it follows immediately from Lemma 9 that the assignment is contextual.
Assume by contradiction that $A\in\mathbf{X}\cup\mathbf{Z}_{\dagger}$,
$E\in\mathbf{E}_{\ddagger}$ and there is a causal trail $A\rightarrow\cdots\
B_{i}\ \cdots\rightarrow E$ such that no $B_{i}$ is observed or there is no
$B_{i}$. Since $A$ is assigned, all children of $A$ will be visited, and
following the trail, the variable $E$ will also be visited from its parent
since there is no observed variable in the trail to block the visit.
Consequently, $E$ will be weighted, which contradicts our assumption that $E$
is not weighted. Hence, the assignment is also safe. ∎

###  A.8 Proof of Lemma 9

In this section, we present the detailed proof of Lemma 9.

###### Proof.

Since $A$ is assigned/weighted and rules in $\mathbb{P}$ are exhaustive, a
rule $\mathcal{R}\in\mathbb{P}$ with $A$ in its head must have fired. Let
$\mathbf{d}$ be a body and $\mathcal{D}$ be a distribution in the head of
$\mathcal{R}$. Since each $d_{i}\in\mathbf{d}$ must be true for $\mathcal{R}$
to fire, $\mathbf{d}\subseteq\mathbf{c}$. We assume that rules in $\mathbb{P}$
are mutually exclusive. Thus, among all rules for $A$, only $\mathcal{R}$ will
fire even when an assignment of some variables in $\mathbf{B}$ is also given.
Hence, by definition of the rule $\mathcal{R}$, we have that,

$\mathcal{D}=P(A\mid\mathbf{d})=P(A\mid\mathbf{c})=P(A\mid\mathbf{c},\mathbf{B})$

∎

###  A.9 Proof of Theorem 10

In this section, we present the detailed proof of Theorem 10.

###### Proof.

Since random variable $A$ is sampled/weighted, $k$ clauses out of $n$ ground
distributional clauses for $A$ must have fired. Those $n-k$ clauses that do
not fire must have at least one atom in their body that is false. Let
$\mathbf{d}$ be an assignment of random variables in the body of the $k$
clauses and $\mathcal{D}$ be a distribution obtained after applying a
combining rule on the multiset of distributions in the head of the $k$
clauses. Since each $d\in\mathbf{d}$ must be true for $k$ clauses to fire,
$\mathbf{d}\subseteq\mathbf{c}$. Even if an assignment of some variables in
$\mathbf{B}$ is also given, the $n-k$ clauses cannot fire since one atom in
the body of these clauses is already false. Hence, by definition, we have that

$\mathcal{D}=P(A\mid\mathbf{d})=P(A\mid\mathbf{c})=P(A\mid\mathbf{c},\mathbf{B})$

∎

##  B Advanced Constructs in DC# Framework

This section describes advanced constructs such as negation, aggregates, and
statistical models in the DC# framework, which are useful while writing
complex relational models (?).

######  Example 18.

The following program illustrates advanced constructs such as negation,
aggregates, and statistical models.

$\displaystyle\mathtt{client(ann)\sim val(true).}$
$\displaystyle\mathtt{loan(l\\_1)\sim bernoulli(0.9).}$
$\displaystyle\mathtt{loan(l\\_2)\sim bernoulli(0.9).}$
$\displaystyle\mathtt{age(C)\sim gaussian(40,10.5)\leftarrow client(C)\cong
true.}$ $\displaystyle\mathtt{has\\_loan(C,L)\sim bernoulli(0.2)\leftarrow
client(C)\cong true,loan(L)\cong true.}$ $\displaystyle\mathtt{status(L)\sim
discrete([0.3:appr,0.7:decl])\leftarrow loan(L)\cong true.}$
$\displaystyle\mathtt{credit\\_score(C)\sim gaussian(700,10.9)\leftarrow
has\\_loan(C,L)\cong true,status(L)\cong appr.}$
$\displaystyle\mathtt{credit\\_score(C)\sim gaussian(600,20.5)\leftarrow
has\\_loan(C,L)\cong true,status(L)\cong decl.}$
$\displaystyle\mathtt{credit\\_score(C)\sim gaussian(M,30.2)\leftarrow
has\\_loan(C,L)\cong true,\backslash\\!+\\!status(L)\cong\\_.}$
$\displaystyle\mathtt{credit\\_score(C)\sim gaussian(750,15.9)\leftarrow
age(C)\cong Y,}$ $\displaystyle\qquad\mathtt{mode(X,(has\\_loan(C,L)\cong
true,status(L)\cong X),appr),linear([Y],[20.1,30.9],M).}$

Next, we explain the advanced constructs of the above program one by one in
detail.

##### Negation

The exhaustiveness condition for well-defined programs is somewhat contrived.
In practice, it is difficult to specify conditional probability distributions
of a RV given all possible assignments of its parents because real-world data
is often inadequate and contains missing values. In such cases, it is
convenient to describe models using negative literals in the bodies of DCs.

The second last clause in the program of Example 18 has a negation in the
body. The negation is interpreted as negation as failure as usual in logic
programming. We understand from the above program that the loan’s status can
take three values: $\mathtt{appr}$ (“approved”), $\mathtt{decl}$ (“declined”),
and “undefined”. That is, its value can be unknown (or undefined) in some
possible worlds. The loan’s status takes the undefined value with a
probability of $1$ when the loan’s identity is false, and the program
specifies a probability distribution over such worlds.

To avoid floundering (?), we allow writing only those clauses
$\mathtt{A_{0}\sim D\leftarrow L_{1},\dots,L_{n}}$ in DC# programs, which
satisfy this condition: if a logical variable $\mathtt{J}$ occurs in the RV
term of a negative literal $\mathtt{L_{i}}$ then $\mathtt{J}$ occurs in the RV
term of a positive literal in $\\{\mathtt{A_{0},L_{1},\dots,L_{i-1}}\\}$. This
means writing unsafe negations are not allowed. For example, writing the
clause

$\mathtt{credit\\_score(C)\sim
gaussian(500,30.2)\leftarrow\backslash\\!+\\!status(L)\cong\\_.}$

is not allowed because the logical variable $\mathtt{L}$ occurs in the RV term
of the negative literal but it does not occur in any positive literal.

##### Aggregates

An alternative approach to combining rules, extensively used for describing
models for relational data, is the aggregation function or aggregate. It is a
function that maps a multiset of values to a single value. Standard examples
are average (if values are numerical), mode (most frequently occurring value),
maximum, minimum, cardinality, etc. For example, the mode maps a multiset
$\mathtt{\\{appr,decl,appr\\}}$ to a value $\mathtt{appr}$. To express these
functions, built-in aggregation predicates are used in conjunction with the
second-order predicate $\mathtt{findall/3}$ as follows:
$\mathtt{findall(T,G,L),aggr(L,R)}$, where $\mathtt{T}$ is a target variable
such that it occurs in a goal $\mathtt{G}$, and $\mathtt{L}$ unifies with the
multiset of the instantiations that $\mathtt{T}$ gets successively on
backtracking over $\mathtt{G}$. This is just like Prolog’s
$\mathtt{findall/3}$ predicate (?). The atom $\mathtt{findall(T,G,L)}$
succeeds with an empty multiset if goal $\mathtt{G}$ has no solutions. A
built-in predicate $\mathtt{aggr/2}$ implements an aggregation function (e.g.,
mode) that maps the multiset $\mathtt{L}$ to a single value $\mathtt{R}$. When
the multiset $\mathtt{L}$ is empty, the aggregation atom $\mathtt{aggr(L,R)}$
fails. We will use a shorthand notation $\mathtt{aggr(T,G,R)}$ for
$\mathtt{findall(T,G,L),aggr(L,R)}$ and call it an aggregation atom.

The last clause in the program of Example 18 has an aggregation atom in its
body. Here, the idea is to consider the joint influence of statuses of all
loans that a client has on the distribution of the client’s credit score
(instead of letting each loan have its own independent influence as in Example
7). There are three situations in which the aggregation atom will fail, and
the body of the last clause will be true: (i) the client has no loan, (ii) it
is undefined that the client has loans or not, (iii) the client has loans, but
statuses of those loans are undefined.

The support for aggregates, which in turn requires support for the second-
order predicate $\mathtt{findall/3}$, takes DC# outside the domain of first-
order probabilistic logic. However, the semantics of programs with aggregates
can also be understood in terms of the ground programs. Just like Prolog, atom
$\mathtt{aggr(T,G,R)}$ gets its truth value after checking for truths of all
instantiations of goal $\mathtt{G}$ in the ground program.

##### Statistical Model Atom

It maps outcomes of RVs in the body of a DC to parameters of the distribution
in the head. Formally, a DC with a statistical model is a clause of the form
$\mathtt{A_{0}\sim D_{\phi}\leftarrow A_{1},\dots,A_{n},M_{\psi}}$, where
$\mathtt{M_{\psi}}$ is an atom implementing a mathematical function that
relates the values of RVs in $\mathtt{\\{A_{1},\dots,A_{n}\\}}$ with
parameters $\phi$ in distribution $\mathtt{D_{\phi}}$ via parameters $\psi$.

An example is shown in the last clause of the program of Example 18. Here,
$\mathtt{linear}$ is a built-in predicate that implements a linear model
relating the value of client’s age and mean $\mathtt{M}$ of the distribution
in the head using parameters $\mathtt{\psi=[20.1,30.9]}$ as follows:
$\mathtt{M\ is\ 20.1\times X+30.9}$.

Many well-known statistical models, such as linear regression, logistic
regression, softmax regression, etc., can be used to describe very complex
relational models (?). Importantly, we have integrated the full expressiveness
of logic programming with the strengths of statistical models in learning
intricate patterns.

###  B.1 Inference in Programs with Advanced Constructs

For inference, first, we should identify RVs from DC# programs with advanced
constructs.

######  Definition 19.

Let $\mathbb{P}$ be a well-defined DC# program with advanced constructs. A RV
set $rv(\mathbb{P})$ of the program is the set of definite clauses obtained by
transforming each clause $\mathtt{A\sim D\leftarrow
L_{1},\dots,L_{m}\in\mathbb{P}}$ like this:

  1. 1.

Let $\mathtt{Body}$ be the empty set. For each positive or negative literal
$\mathtt{L_{i}}$ of the form $\mathtt{T\cong V}$ $\mathtt{rv(T)}$ is added to
$\mathtt{Body}$.

  2. 2.

A clause $\mathtt{rv(A)\leftarrow Body}$ is added to $rv(\mathbb{P})$.

######  Example 19.

The RV set for the program in Example 18 is:

$\displaystyle\mathtt{rv(client(ann)).}$
$\displaystyle\mathtt{rv(loan(l\\_1)).}$
$\displaystyle\mathtt{rv(loan(l\\_2)).}$
$\displaystyle\mathtt{rv(age(C))\leftarrow rv(client(C)).}$
$\displaystyle\mathtt{rv(has\\_loan(C,L))\leftarrow
rv(client(C)),rv(loan(L)).}$ $\displaystyle\mathtt{rv(status(L))\leftarrow
rv(loan(L)).}$ $\displaystyle\mathtt{rv(credit\\_score(C))\leftarrow
rv(has\\_loan(C,L)),rv(status(L)).}$
$\displaystyle\mathtt{rv(credit\\_score(C))\leftarrow rv(age(C)).}$

The above definition differs from Definition 6 in only one way. It
additionally specifies the way negative literals of the form
$\mathtt{\backslash\\!+\\!T\cong V}$ should be handled. However, nothing
special is done for them because just like positive literals that compare
outcomes of ground RV terms with values, negative literals compare the
outcomes with values or with the “undefined” value. Additionally, comparison
and statistical model atoms are ignored while constructing the RV set. This is
because they do not contain RV terms. The aggregation atoms do contain RV
terms, but they do not constrain instantiations of logical variables appearing
in other atoms. So, they are also ignored.

However, we must ensure that Proposition 2 is valid even after using negations
that relax the exhaustiveness condition. Now, it may happen that
$rv(\mathbb{P})\models\mathtt{rv(A)}$ but no distribution is specified for
$\mathtt{A}$ in some ground programs. In this case, however, the “undefined”
value is assigned to $\mathtt{A}$, so the exhaustiveness condition is still
implicitly present and Proposition 2 is valid.

Next, we should be able to identify direct influence relationships among RVs.
However, the second-order aggregation atoms complicate the construction of
dependency sets. We add an additional rule in Definition 8 to deal with it.

######  Definition 20.

Let $rv(\mathbb{P})$ be a RV set of a well-defined DC# program $\mathbb{P}$
with advanced constructs. The dependency set $dep(\mathbb{P})$ is the union of
$rv(\mathbb{P})$ and the set of definite clauses $pa(\mathbb{P})$ obtained by
transforming each clause $\mathcal{C}\equiv\mathtt{A\sim D\leftarrow
L_{1},\dots,L_{m}\in\mathbb{P}}$ with non empty body as follows:

  1. 1.

Let $\mathtt{rv(A)\leftarrow rv(T_{1}),\dots,rv(T_{n})}$ be a clause in
$rv(\mathbb{P})$ corresponding to $\mathcal{C}$.

  2. 2.

For each $\mathtt{rv(T_{i})}$ a clause $\mathtt{pa(A,T_{i})\leftarrow
rv(A),rv(T_{1}),\dots,rv(T_{n})}$ is added to $pa(\mathbb{P})$.

  3. 3.

For each literal $\mathtt{L_{i}}$ of the form $\mathtt{aggr(T,Q,R)}$,

    1. (a)

Let $\mathtt{Q}$ is of the form $\mathtt{(Q_{1},\dots,Q_{n})}$ and let
$\mathtt{Body_{1}}$ be the empty set. For each literal $\mathtt{Q_{i}}$ of the
form $\mathtt{T\cong V}$ an atom $\mathtt{rv(T)}$ is added to
$\mathtt{Body_{1}}$.

    2. (b)

For each literal $\mathtt{Q_{i}}$ of the form $\mathtt{T\cong V}$:
$\mathtt{pa(A,T)\leftarrow rv(A),rv(T_{1}),\dots,rv(T_{n}),Body_{1}}$ is added
to $pa(\mathbb{P})$.

However, note that in Prolog, all free variables appearing in
$\mathtt{findall/3}$ are bound with the existential operator. So, all free
variables in $\mathtt{aggr/3}$ should be renamed before constructing the
dependency set.

######  Example 20.

In addition to the RV set of Example 19, the dependency set of Example 18
contains the following clauses. Notice the renamings in the last two clauses.

$\displaystyle\mathtt{pa(age(C),client(C))\leftarrow
rv(age(C)),rv(client(C)).}$
$\displaystyle\mathtt{pa(has\\_loan(C,L),client(C))\leftarrow
rv(has\\_loan(C,L)),rv(client(C)),rv(loan(L)).}$
$\displaystyle\mathtt{pa(has\\_loan(C,L),loan(L))\leftarrow
rv(has\\_loan(C,L)),rv(client(C)),rv(loan(L)).}$
$\displaystyle\mathtt{pa(status(L),loan(L))\leftarrow
rv(status(L)),rv(loan(L)).}$
$\displaystyle\mathtt{pa(credit\\_score(C),has\\_loan(C,L))\leftarrow
rv(credit\\_score(C)),rv(has\\_loan(C,L)),rv(status(L)).}$
$\displaystyle\mathtt{pa(credit\\_score(C),status(L))\leftarrow
rv(credit\\_score(C)),rv(has\\_loan(C,L)),rv(status(L)).}$
$\displaystyle\mathtt{pa(credit\\_score(C),age(C))\leftarrow
rv(credit\\_score(C)),rv(age(C)).}$
$\displaystyle\mathtt{pa(credit\\_score(C),has\\_loan(C,L_{1}))\leftarrow
rv(credit\\_score(C)),rv(age(C)),}$
$\displaystyle\qquad\qquad\mathtt{rv(has\\_loan(C,L_{1})),rv(status(L_{1})).}$
$\displaystyle\mathtt{pa(credit\\_score(C),status(L_{2}))\leftarrow
rv(credit\\_score(C)),rv(age(C)),}$
$\displaystyle\qquad\qquad\mathtt{rv(has\\_loan(C,L_{2})),rv(status(L_{2})).}$

After identifying RVs and direct influences among them, the simulation of
programs defined in Algorithm 5 is straightforward. One can easily add rules
to the proof procedure (Algorithm 6) for resolving these advanced constructs.
Predicates $\mathtt{rv/1,pa/2}$ should be tabled (?) for efficiency.


<EMAIL_ADDRESS>
# Buffer Overflow in Mixture of Experts

Jamie Hayes  Alphabetical order  Google DeepMind  Ilia Shumailov  Alphabetical
order  Google DeepMind  Itay Yona  Alphabetical order  Google DeepMind

###### Abstract

Mixture of Experts (MoE) has become a key ingredient for scaling large
foundation models while keeping inference costs steady. We show that expert
routing strategies that have cross-batch dependencies are vulnerable to
attacks. Malicious queries can be sent to a model and can affect a model’s
output on other benign queries if they are grouped in the same batch. We
demonstrate this via a _proof-of-concept_ attack in a _toy experimental
setting_.

##  1 Overview

In the early days of computing, time-sharing gave a major boost to user
productivity and hardware utilisation. Yet, it also brought significant
security challenges; computers were now shared, meaning unauthorised access
had to be mediated, while individual private data had to be protected through
explicit isolation. Machine Learning (ML) is now experiencing its time-sharing
moment – hardware costs and high utilisation requirements are pushing
accelerators that host ML models to be shared across different users,
naturally leading to a question: are there any new security concerns?

In this work we demonstrate that when Mixture of Experts (MoE) is used in, for
example, a transformer architecture (Vaswani et al., 2017), an adversary can
change the model prediction on data of other users who happen to be placed
into the same batch. This is because many routing strategies that assign
inputs to experts are not inter-batch data (order) invariant. This
vulnerability affects current MoE implementations that set a buffer capacity
limit on the number of inputs each expert can process.

We describe two ways to generate adversarial data and show that both provide
an adversary with an ability to run integrity attacks  _i.e._ cause
deterministic fault injection, as well as, availability attacks _i.e._ denial
of expert attacks against other users. Finally, we investigate several
mitigation strategies that nullify the attack or significantly reduce its
efficiency.

##  2 Technical details

We first describe MoE (Jacobs et al., 1991; Jordan and Jacobs, 1994) and its
later sparse variants (Shazeer et al., 2017; Fedus et al., 2022; Riquelme et
al., 2021; Du et al., 2022), and following this, introduce attack details.

###  2.1 Mixture of Experts

We follow the set-up described in Riquelme et al. (2021) for the deep learning
setting. Given an input token $z\in\mathbb{R}^{d}$, the output of $z$ through
an MoE layer is given by:

$\displaystyle MoE(z)=\sum_{i=1}^{n}g_{i}(z)e_{i}(z),$ (1)

where $e_{i}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ are the $n$ experts, and
$g_{i}:\mathbb{R}^{d}\rightarrow\mathbb{R}$ is a gating function that assigns
an input conditioned weight to an expert. In the settings we consider, $g$ and
$e$ are parameterized by fully-connected neural networks. MoE is referred to
as _sparse_ if $g$ assigns non-zero weight to fewer than $n$ experts. This is
can save on inference costs as the outputs of all $n$ experts do not need to
be evaluated. In most MoE implementations, the number of zero weights is
preset by only routing the top-$k$ values of $g$ to their assigned experts,
where $k<n$. After a subset of top-$k$ values of $g$ have been selected, it is
common to adjust the weights in eq. 1 by re-normalizing this subset through a
softmax function.

One of the primary motivations for using MoE is improved model performance for
a fixed number of activated parameters (Shazeer et al., 2017). With MoE, the
whole network is not needed to evaluate a given input. For example, the
popular open-source MoE transformer model, Mixtral-8$\times$7B (Jiang et al.,
2024), has 46.7B parameters, but effectively only uses a smaller subset of
12.9B parameters per token.

###  2.2 Routing strategies

Our investigation focuses on transformer models that use MoE, where inputs to
the model are sequences of tokens from a vocabulary. Let
$Z\in\mathbb{R}^{B\cdot T\times d}$ be an input to the MoE layer, where $B$ is
the batch size, $T$ is the number of tokens per sequence in the batch, and $d$
is a $d$-dimensional representation of a token at this layer. For each $z\in
Z$, we first compute $g(z)=\\{g_{1}(z),g_{2}(z),\ldots,g_{n}(z)\\}$, where $n$
is the number of experts. This gives the matrix

$\displaystyle G=\begin{bmatrix}g_{11}&g_{12}&\dots&g_{1n}\\\
g_{21}&g_{22}&\dots&g_{2n}\\\ \vdots&\vdots&\vdots&\vdots\\\ g_{(B\cdot
T)1}&g_{(B\cdot T)2}&\dots&g_{(B\cdot T)n}\\\ \end{bmatrix},$ (2)

where $g_{ij}$ denotes the routing weight of token $z_{i}$ to expert $e_{j}$.
There are various methods to rank and prioritise how tokens are sent to
experts. Riquelme et al. (2021); Fedus et al. (2022); Shazeer et al. (2017);
Lepikhin et al. (2020) use variations of the following strategy, which is
referred to as the _vanilla_ routing strategy: Sequentially, starting from the
first row of $G$, find the column with top-1 value for that row and send that
token to the expert associated with that column. Once top-1 assignments have
been made for all $B\cdot T$ tokens, repeat for top-2 assignments, and keep
repeating this until all top-$k$ assignments for each row have been made. This
routing strategy is given in pseudo-code in Algorithm 1.

Algorithm 1 Vanilla expert routing strategy

Args: number of experts $n$, number of experts per token $k$, batch size $B$,
number of tokens per example $T$, dimensionality of token representation $d$,
expert functions $e_{i}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$
($i\in\\{1,\ldots,n\\}$), gating function
$g:\mathbb{R}^{d}\rightarrow\mathbb{R}^{n}$, input to MoE layer
$Z\in\mathbb{R}^{B\cdot T\times d}$.

$G\leftarrow$ Initialize a $(B\cdot T)\times n$ matrix by passing each $z\in
Z$ into the gating function $g$.

$h=1$

while $h\leq k$ do

$i=1$

while $i\leq B\cdot T$ do

$g_{ij}\leftarrow\text{Get $jth$ column for row $i$ of $G$ (the column with
the top-$h$ value).}$

Assign $z_{i}$ to expert $e_{j}$

$i=i+1$

$h=h+1$

In theory, each expert in Algorithm 1 can process $B\cdot T$ tokens. This
could happen if, for each $k$, the top-$k$ value in each row of $G$ had the
same expert assignment. An uneven assignment of experts is generally an
undesirable property as this results in under utilization of some experts. To
partially mitigate this issue, a buffer capacity limit is usually set for each
expert (Riquelme et al., 2021; Gale et al., 2023). This is a static number,
$B_{e}$, that dictates how many tokens each expert can process, which in
practice is set by a capacity slack variable $C$, where:

$\displaystyle B_{e}=\text{round}\bigg{(}\frac{kCBT}{n}\bigg{)}.$ (3)

If $C<1$, some routing assignments will be ignored, while imbalances in
routing assignments between experts can be tolerated if $C>1$. We refer the
interested reader to Fedus et al. (2022), who discuss the trade-offs in
setting the buffer capacity limit.

###  2.3 Threat model and attack method

Figure 1: _Overall attack flow. The adversary pushes their data into the
shared batch, that already contains user data. As tokens get distributed
across different experts, adversarial data fills the expert buffers that would
be preferred by the user, dropping or routing their data to experts that
produce suboptimal outputs._

Our attack relies on MoE routing that uses finite sized buffer queues for each
individual expert. When these queues are filled, tokens could be dropped
111Note, although dropping tokens is generally considered to be wasteful, they
still influence downstream layers as they are directly passed to the next
layer through a residual connection. Our experiments will drop tokens when an
expert buffer is full rather than re-route to different experts. or passed to
another expert whose queue is not full (Fedus et al., 2022). This opens up a
new attack vector; if adversarial data is mixed with benign user data in a
batch, the adversarial data can influence the expert choice of the benign data
by filling the buffer of certain experts. We depict the overall attack process
in Figure 1. For the vanilla routing strategy in Algorithm 1, tokens are
sequentially assigned to experts according to the order they appear in the
batch. This means inputs at the beginning of the batch can affect the routing
assignments of tokens further down in the order. Imagine there are two data-
points in the batch $x_{1}$ and $x_{2}$, each point has two tokens,
$\\{z^{1}_{i},z^{2}_{i}\\}$, and there are two experts $e_{1}$ and $e_{2}$
with a buffer capacity of two. The input to the MoE layer is therefore
$\\{z^{1}_{1},z^{2}_{1},z^{1}_{2},z^{2}_{2}\\}$, and let’s assume that all
tokens should be sent to $e_{1}$. According to the vanilla routing strategy,
the first experts buffer will be filled first with $z^{1}_{1}$ and then
$z^{2}_{1}$, and so the assignments for $z^{1}_{2}$ and $z^{2}_{2}$ are
ignored. Crucially, because each expert has a buffer capacity limit, the
assignment of tokens to experts is batch dependent — one data-point can change
the expert assignment of another data-point. Our attack exploits this inter-
batch routing dependency to find adversarial data-points that can negatively
affect the outputs of other data-points in the batch.

Our attack set up is as follows: Let
$f_{\theta}:\mathcal{V}^{T}\rightarrow\mathbb{R}^{s}$ be a transformer model
that uses MoE layers and is parameterized by a set of weights $\theta$, where
$\mathcal{V}^{T}$ denotes a set of $T$ tokens from a vocabulary $\mathcal{V}$
of size $s$. The transformer model takes as input a sequence of $T$ tokens and
outputs a probability distribution over $\mathcal{V}$. It is typical to batch
inputs together to maximize throughput, we assume that given a batch of inputs
$X=\\{x_{1},x_{2},\ldots,x_{B-1},x^{*}\\}$, where each $x\in X$ is a set of
$T$ tokens, that the adversary controls all elements in the batch other than
$x^{*}$. We refer to this subset of adversarially controlled inputs as
$X^{\mathcal{A}}$. Note that $f_{\theta}(X)\in\mathbb{R}^{B\times s}$, and we
denote the output of $x^{*}$ as
$f_{\theta,x^{*}}(X=X^{\mathcal{A}}\cup\\{x^{*}\\})\in\mathbb{R}^{s}$. We also
assume the adversary can control the order of batch elements, and always
ensures that $x^{*}$ is last in the batch. Let’s assume that the next
predicted token in the sequence after $x^{*}$ should be $y\in\mathcal{V}$
222We abuse notation and let $y$ represent both the vocabulary element _and_
index within $\mathcal{V}$. and before the attack,
$\operatorname*{arg\,max}f_{\theta,x^{*}}(X)=y$. The attack goal is to find
adversarial inputs $X^{\mathcal{A}}$ such that
$\operatorname*{arg\,max}f_{\theta,x^{*}}(X)\neq y$. We find this by
optimizing $X^{\mathcal{A}}$ to minimize the loss
$l_{X^{\mathcal{A}}}(x^{*})=f_{\theta,x^{*}}(X^{\mathcal{A}}\cup\\{x^{*}\\})_{y}-\operatorname*{arg\,max}_{\hat{y}\neq
y}f_{\theta,x^{*}}(X^{\mathcal{A}}\cup\\{x^{*}\\})$, where
$f_{\theta,x^{*}}(X^{\mathcal{A}}\cup\\{x^{*}\\})_{y}$ denotes the output at
token $y$ in $\mathcal{V}$. We refer to the set of adversarial inputs after
this optimization process as $\tilde{X}^{\mathcal{A}}$.

Throughout our experiments we use simple random search to construct
$\tilde{X}^{\mathcal{A}}$. We randomly sample a subset of tokens from the
vocabulary and check if this decreases $l_{X^{\mathcal{A}}}(x^{*})$. The step-
by-step process of this search is given in Algorithm 2.

Algorithm 2 Random search attack

Args: Loss function $l$, batch size $B$, number of tokens per sample in the
batch $T$, number of iterations $M$, number of tokens to change per iteration
$r$, input $x^{*}$ with target $y$, vocabulary $\mathcal{V}$.

Output: Optimized adversarial data $\tilde{X}^{\mathcal{A}}$

$X^{\mathcal{A}}\leftarrow\\{z^{1}_{1},z^{2}_{1},\ldots,z^{T}_{1},z^{1}_{2},\ldots,z^{T}_{B-1}\\}$
$\triangleright$ Initialize $(B-1)\cdot T$ tokens at random from
$\mathcal{V}$.

smallest loss $\leftarrow 999$ $\triangleright$ We track the adversarial
inputs that achieve the smallest loss on $l(\cdot)$.

$\tilde{X}^{\mathcal{A}}\leftarrow X^{\mathcal{A}}$ $\triangleright$
Initialize the current best choice of adversarial inputs to
$\tilde{X}^{\mathcal{A}}$.

$i=1$

while $i\leq M$ do

loss $\leftarrow l_{X^{\mathcal{A}}}(x^{*})$ $\triangleright$ Compute the loss
on $x^{*}$.

if loss $<$ smallest loss then $\triangleright$ Check if current loss is
smaller than any previously computed loss.

smallest loss $\leftarrow$ loss $\triangleright$ Assign the tracked smallest
loss to current loss.

$\tilde{X}^{\mathcal{A}}\leftarrow X^{\mathcal{A}}$ $\triangleright$ Assign
the best choice of adversarial inputs to the current ones.

else

$X^{\mathcal{A}}\leftarrow\tilde{X}^{\mathcal{A}}$ $\triangleright$ Re-
initialize the current adversarial inputs to the tracked adversarial inputs
that achieved the smallest loss so far.

$X^{\mathcal{A}}\leftarrow$ Replace $r$ tokens at random.

$i=i+1$

return $\tilde{X}^{\mathcal{A}}$ $\triangleright$ Return the optimized
adversarial inputs.

This attack assumes that an adversary is interacting with a black-box model
that is utilising MoE, and observes logit outputs. It also assumes the
adversary can ensure their data is always grouped in the same batch as the
target point $x^{*}$. We make no assumptions about the MoE configuration,
except that it has per-expert queues that are finite and uses the batch order
dependent vanilla routing strategy. This final assumption is not unreasonable
and is how many MoE based transformer models are trained and deployed.

##  3 Attack demonstration

We demonstrate our attack on the popular open-source MoE transformer model,
Mixtral-8$\times$7B, which has a vocabulary size of 32,000, uses 8 experts
(per layer), and each token can be assigned to a maximum of 2 experts. We set
the target $x^{*}$ to the prompt “Solve the following equation: 1+1=” which
has a sequence length of 12 tokens. We confirmed that the most likely next
token output by the model is “2”, with a softmax probability of 25.9% 333One
may wonder why this probability is so low. This is likely because we are using
the base model that has not been instruction tuned and we do not use few-shot
learning.. We set the batch size to 8, and so there are a total of
$(B=8)\times(T=12)$ tokens in the batch. Mixtral-8$\times$7B does not set an
expert buffer capacity limit in its default MoE implementation and so is _not_
vulnerable to our attack. We augment the MoE to use the vanilla routing
strategy with an expert buffer capacity of 38 tokens (out of a possible 96).
We confirmed that this change did not affect the model’s output on $x^{*}$.
However, such a small buffer capacity will almost certainly affect the quality
of outputs on longer sequences.

The number of attack iterations is set to 1,000, and we modify $5$ tokens per
data-point in the adversarial controlled batch $X^{\mathcal{A}}$ per
iteration. The number of tokens per data-point is 12 and there are 7 data-
points in $X^{\mathcal{A}}$, so the adversary can modify a total of 35 tokens
out of a possible 84 per iteration.

Results of the attack are shown in Figure 2. Initially, “2” has the largest
probability, and the next most likely token “1” has a probability
$\approx$20%. At the beginning of the attack, the number of tokens assigned to
the second expert from the adversarial batch $X^{\mathcal{A}}$ is 37 and there
is one token assigned to this expert from $x^{*}$. At iteration 60, the buffer
of the second expert is filled entirely with tokens from $X^{\mathcal{A}}$,
which means the token in $x^{*}$ that was assigned to this expert is dropped.
This change in expert assignments (and expert assignments in downstream
layers) is enough to cause “1” to become the most likely next token. We also
plot the difference in softmax probabilities between tokens “2” and “1”
throughout the attack.

Figure 2: _Probability of correct token, “ 2”, and (most likely) incorrect
token, “1”, throughout the random search attack. By the end of the attack the
output token with largest probability is the incorrect token, “1”._

##  4 Anecdotal evidence of transferability to different prompts

It is unrealistic to assume an adversary ensure that $\tilde{X}^{\mathcal{A}}$
will always grouped with a target point. The attack becomes much more
practical if one can show that errors can be induced on data-points other than
$x^{*}$ when included in a batch with $\tilde{X}^{\mathcal{A}}$. This opens up
the possibility of untargeted attacks, where an adversary constructs
$\tilde{X}^{\mathcal{A}}$ that block prior chosen experts that are preferred
for certain tasks (e.g. arithmetic or translation).

We check how well the optimized adversarial data-points,
$\tilde{X}^{\mathcal{A}}$, transfer to arithmetic prompts other than $x^{*}$.
We conjecture that if these other data-points are sampled from the same
distribution as $x^{*}$, they are more likely to have similar expert routing
assignments, and so are more likely to be affected by
$\tilde{X}^{\mathcal{A}}$, which was optimized specifically for $x^{*}$. We
measure the change in most likely next token probabilities on prompts that are
similar to $x^{*}$ in the Table 1.

Table 1: _Transferability of $\tilde{X}^{\mathcal{A}}$ constructed against
$x^{*}$ (“Solve the following equation: 1+1=”)._

Prompt & completion before | Prompt & completion after | Probability change | Did the attack
---|---|---|---
batching with $\tilde{X}^{\mathcal{A}}$ | batching with $\tilde{X}^{\mathcal{A}}$ | of the correct token | induce an error?
“Solve the following equation: 2+2=4” | “Solve the following equation: 2+2=4” | $-1.5\%$ | No
“Solve the following equation: 4+4=8” | “Solve the following equation: 4+4=4” | $-1.8\%$ | Yes

The attack successfully transferred to one other prompt; the inclusion of the
adversarial data-points in the batch induced an error in the model’s output on
this prompt. Even for the prompt where the attack failed to induce an error,
the probability of the correct token, “4”, was reduced. This is anecdotal
evidence that adversarial data can transfer and induce model output errors on
prompts unseen during the construction of $\tilde{X}^{\mathcal{A}}$.

##  5 Does the position in batch matter?

Any expert routing strategy that uses an expert buffer capacity limit that is
smaller than the total number of tokens in the batch will be vulnerable to an
attack that attempts to fill up certain expert buffers to block other token’s
assigned experts. The most vulnerable data-point in the batch depends on the
MoE routing strategy. The vanilla routing strategy sequentially assigns
experts by position in the batch, and so a data-point that is last in the
batch is more vulnerable than data-points that are near the front, as there
are more opportunities for other data-points to change the expert assignment
of the tokens from the targeted data-point. We verified this by running the
random search attack while always placing $x^{*}$ to be first in the batch;
the attack did not induce an error in the model’s output on $x^{*}$. However,
we note that the expert assignments of a data-point that is first in the batch
can still be affected by other points if the vanilla routing strategy assigns
tokens to multiple experts, i.e., uses top-$k$ gating values where $k>1$.
Consider the previous illustrative example where four tokens
$\\{z^{1}_{1},z^{2}_{1},z^{1}_{2},z^{2}_{2}\\}$ can be sent to experts $e_{1}$
or $e_{2}$, each with a buffer capacity of two. If we set the targeted point
to $x^{*}=x_{1}=\\{z^{1}_{1},z^{2}_{1}\\}$, and let the adversary control
$\tilde{X}^{\mathcal{A}}=x_{2}=\\{z^{1}_{2},z^{2}_{2}\\}$, then
$\tilde{X}^{\mathcal{A}}$ cannot affect the top-1 expert assignments of
$x^{*}$, but can change the top-2 assignment. For example, if the top-2
assignment for $z^{1}_{1}$ is $e_{2}$ but the top-1 assignments for
$\tilde{X}^{\mathcal{A}}$ are also $e_{2}$, then $z^{1}_{1}$ would not be
assigned to $e_{2}$ as its buffer has been filled with the adversarial tokens.

##  6 Attack sensitivity to buffer capacity to limit

The smaller the expert buffer capacity limit, the easier it will be to find a
set of adversarial tokens, $\tilde{X}^{\mathcal{A}}$, that block the
assignment of tokens from $x^{*}$ to its preferred expert(s). We repeat our
attack under different expert buffer capacity limits and present results in
the Table 2.

Table 2: _Attack success for different expert buffer capacity limits._ $\nicefrac{{\text{Buffer capacity, }B_{e}}}{{\text{Number of tokens in batch, }B\cdot T}}$ | Capacity variable, $C$ | Did the attack induce an error?
---|---|---
0.2 | 0.8 | Yes
0.4 | 1.6 | Yes
0.5 | 2 | No
0.6 | 2.4 | No
0.8 | 3.2 | No
1 | 4 | No

As expected, the attack is successful for smaller buffer capacity limits, and
as the size of the buffer grows it becomes harder to affect expert assignments
of the target $x^{*}$, causing the attack to fail.

##  7 Example of a denial-of-expert attack

Until now the attack goal has been to induce model output errors for a target
input $x^{*}$. Overflowing an expert’s buffer or changing the routing
assignment (for deeper MoE layers in model) to cause such an error is a
byproduct of this goal. We now show we can explicitly optimize to fill a
chosen expert’s buffer. We set the buffer capacity limit to 20 tokens per
expert and optimize the adversarial data-points $X^{\mathcal{A}}$ to fill the
buffer of the third expert (out of eight) in the first MoE layer by changing
the objective function in Algorithm 2 to count how many tokens from $x^{*}$
are processed by this expert. Note, this changes the attack assumptions, as
the adversary now needs white-box access to the model during the construction
of $\tilde{X}^{\mathcal{A}}$ because they need to monitor an expert’s buffer.
From Figure 3, we see that initially 6 out of the 12 tokens from $x^{*}$ are
processed by this expert; within 28 iterations we can find a batch of data-
points $\tilde{X}^{\mathcal{A}}$ that cause all of these 6 tokens from $x^{*}$
to be dropped.

Figure 3: _Attack that constructs adversarial inputs that block the preferred
expert of the majority of tokens from the target $x^{*}$._

##  8 Mitigations

The attack relies on two optimizations made by MoE: (1) the usage of expert
buffer capacity limits, and (2) batch dependent expert routing assignments. If
the buffer capacity limit is removed, or set to the total number of tokens in
the batch, then the attack will fail. We confirmed this by checking that our
attack fails on the default routing strategy in the Mixtral-8$\times$7B that
does not set a buffer capacity limit. However, the use of expert buffer
capacity limits is useful from a hardware utilization perspective as it avoids
wasted computation and memory. We therefore list other mitigation methods,
that will not offer guarantees of protection, but do decrease the efficiency
of the attack. For each suggestion that follows, we confirmed that our random
search attack fails when this mitigation is deployed.

Randomize batch order. The attack relies on the adversary controlling which
position the target prompt $x^{*}$ appears in the batch. Model owners should
always randomize the order of inputs in a batch. For particular sensitive
queries where an increase in inference cost is tolerable, one could also run
inference twice with two different orderings in the batch. If the two outputs
on $x^{*}$ are significantly different, this should be a signal to the model
owner that an inspection of the underlying causes is necessary.

Use a large capacity slack, $C$. As we have already seen from Table 2, using a
larger capacity slack $C$ nullifies the attack as it becomes more difficult to
fill an experts buffer with only tokens from adversarial inputs.

Sampling from gate weights instead of selecting the top-$k$. The expert
assignments for a token $z$ are deterministic; the top-$k$ values of
$g(z)=\\{g_{1}(z),g_{2}(z),\ldots,g_{n}(z)\\}$ are routed to their chosen
experts. If the attack can infer which of the $k$ experts are selected for
$z$, they can construct adversarial inputs to block those experts. Given that
$g(z)$ outputs a probability distribution over $n$ experts, we could sample
from this distribution to select $k$ experts rather than selecting the
top-$k$. Sampling from $g(z)$ results in randomness in the expert assignments
for $z$, making it harder for an attack to decide which expert buffers to
fill.

##  9 Discussion

Combining queries from untrusted sources improves hardware utilization, but
opens the door for exploitation. The machine learning community is well aware
that dropping tokens due to buffer capacity limits, or routing tokens to non-
preferred experts, can reduce performance (Fedus et al., 2022). We encourage
the machine learning community to consider risks that come with inter-batch
dependent outputs beyond subpar model performance.

It is unclear if an attack on MoE could represent a practical risk to deployed
models. Although the practicality of our attack is limited in its current
form, there are obvious directions for future work. Firstly, we expect that
random search is extremely inefficient, and developing stronger gradient-based
attacks could help model owners understand worst-case behavior. Secondly,
another route to measuring worst-case behaviour is to record the variance in
model outputs when tokens are dropped at each MoE layer, or tokens are
assigned to random experts. If model outputs are robust to changes in expert
routing then the risk of an attack diminishes. Thirdly, there are various
routing strategies beyond the vanilla strategy we experiment with; it is
important to understands the trade-offs between performance and security that
come with these different routing choices. Fourthly, models that use MoE are
commonly trained with load balancing auxiliary losses. This may help increase
the robustness to an attack, as each expert may perform well at processing
tokens across a broad range of topics. Future work could investigate how load
balancing during training can help mitigate the attack. Finally, we have not
experimented with instruction-tuned models, which we expect to be more likely
to be deployed in security sensitive settings (e.g. coding assistants).

## References

  * Du et al. (2022) N. Du, Y. Huang, A. M. Dai, S. Tong, D. Lepikhin, Y. Xu, M. Krikun, Y. Zhou, A. W. Yu, O. Firat, B. Zoph, L. Fedus, M. Bosma, Z. Zhou, T. Wang, Y. E. Wang, K. Webster, M. Pellat, K. Robinson, K. Meier-Hellstern, T. Duke, L. Dixon, K. Zhang, Q. V. Le, Y. Wu, Z. Chen, and C. Cui.  Glam: Efficient scaling of language models with mixture-of-experts, 2022\.
  * Fedus et al. (2022) W. Fedus, B. Zoph, and N. Shazeer.  Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022.
  * Gale et al. (2023) T. Gale, D. Narayanan, C. Young, and M. Zaharia.  Megablocks: Efficient sparse training with mixture-of-experts.  _Proceedings of Machine Learning and Systems_ , 5, 2023.
  * Jacobs et al. (1991) R. A. Jacobs, M. I. Jordan, S. J. Nowlan, and G. E. Hinton.  Adaptive mixtures of local experts.  _Neural Computation_ , 3(1):79–87, 1991.  10.1162/neco.1991.3.1.79.
  * Jiang et al. (2024) A. Q. Jiang, A. Sablayrolles, A. Roux, A. Mensch, B. Savary, C. Bamford, D. S. Chaplot, D. d. l. Casas, E. B. Hanna, F. Bressand, et al.  Mixtral of experts.  _arXiv preprint arXiv:2401.04088_ , 2024.
  * Jordan and Jacobs (1994) M. I. Jordan and R. A. Jacobs.  Hierarchical mixtures of experts and the em algorithm.  _Neural Computation_ , 6(2):181–214, 1994.  10.1162/neco.1994.6.2.181.
  * Lepikhin et al. (2020) D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen.  Gshard: Scaling giant models with conditional computation and automatic sharding.  _arXiv preprint arXiv:2006.16668_ , 2020.
  * Riquelme et al. (2021) C. Riquelme, J. Puigcerver, B. Mustafa, M. Neumann, R. Jenatton, A. S. Pinto, D. Keysers, and N. Houlsby.  Scaling vision with sparse mixture of experts, 2021.
  * Shazeer et al. (2017) N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean.  Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017.
  * Vaswani et al. (2017) A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.