Quadratic Gating MoE

• Quadratic gating MoE is a technique that replaces linear mappings with quadratic functions to enhance expert routing by capturing input curvatures.
• The method leverages both quadratic-polynomial and quadratic-monomial gates to achieve statistical efficiency and faster convergence rates in model predictions.
• Integrating quadratic gating with nonlinear experts and active-attention strategies leads to superior empirical performance across vision, language, and time-series tasks.

Quadratic gating in Mixture-of-Experts (MoE) models refers to the use of quadratic functions, rather than conventional linear mappings, for the gating network responsible for routing inputs to experts. This approach enhances the expressivity and efficiency of routing decisions, drawing theoretical and architectural parallels with attention mechanisms. Two main quadratic gating variants are recognized—quadratic-polynomial and quadratic-monomial gates—both yielding provable statistical efficiency and empirical gains over standard linear gating.

1. Quadratic Gating Formulation

In a dense MoE setting, the model output has the form

$$\text{MoE}(x; \Theta) = \sum_{i=1}^N g_i(x; \Theta_g) \cdot E_i(x; \Theta_i),$$

where $g_i(x; \Theta_g)$ is the gate value for expert $i$, often computed using softmax over scores $s_i(x)$:

$$g(x; \Theta_g) = \operatorname{Softmax}(s_1(x; \theta_g), \dots, s_N(x; \theta_g))$$

with $s_i(x)$ typically linear, $s_i(x) = w_i^\top x + b_i$. Quadratic gating generalizes $s_i$ to second-degree forms. The two canonical types are:

• Quadratic-polynomial gate:

$$s_i(x) = x^\top A_i x + b_i^\top x + c_i, \qquad g_i(x) = \frac{\exp s_i(x)}{\sum_{j=1}^N \exp s_j(x)}$$

• Quadratic-monomial gate (no linear term):

$$s_i(x) = x^\top A_i x + c_i, \qquad g_i(x) = \frac{\exp s_i(x)}{\sum_{j=1}^N \exp s_j(x)}$$

Here, each $A_i \in \mathbb{R}^{d \times d}$ is a learned quadratic parameter matrix, allowing the gating function to capture input curvatures and more complex expert assignment boundaries (Akbarian et al., 2024).

2. Relationship to Attention Mechanisms

Quadratic gating arises naturally when analyzing the affinity computation in self-attention. Standard self-attention computes

$$\alpha_i(q, k) = \frac{\exp(q^\top k_i)}{\sum_j \exp(q^\top k_j)}$$

where

$q = W^q x + b^q, \quad k_i = W^k x + b^k_i$

The inner product

$$(W^q x + b^q)^\top (W^k x + b^k_i) = x^\top (W^{q\top} W^k) x + (b^q)^\top W^k x + (W^{q\top} x)^\top b^k_i + (b^q)^\top b^k_i$$

exhibits a quadratic-polynomial structure in $x$. When bias terms are omitted or shared, the model implicitly adopts the monomial quadratic form. Thus, quadratic-gated MoE layers can be understood as a class of attention mechanisms where expert routing is governed by a second-order kernel over the input, with each expert defined by a distinct quadratic form (Akbarian et al., 2024).

3. Statistical and Theoretical Properties

Theoretical analysis is conducted in a regression framework where observations $Y_i$ follow

$Y_i = f_{G_*}(X_i) + \varepsilon_i,$

with $f_{G_*}$ an MoE function using a quadratic gate. Estimation proceeds via least-squares minimization,

$$\widehat G_n = \arg\min_{G\in \mathcal{G}_N} \sum_{i=1}^n (Y_i - f_G(X_i))^2$$

• Convergence Rates: Both quadratic-polynomial and monomial gates achieve prediction error rates

$$\| f_{\widehat G_n} - f_{G_*} \|_{L^2} = \widetilde{O}_P(n^{-1/2})$$

• Identifiability:
  • Poly-polynomial gates: If each expert $h(x; \theta)$ is nonlinear (e.g., ReLU/tanh) and satisfies “poly-strong identifiability,” then the estimated parameters for each gate and expert converge at parametric $n^{-1/2}$ rate (locally), except in over-specified regimes, where lower rates such as $n^{-1/6}$ or $n^{-1/12}$ may appear.
  • Monomial gates: Under “mono-strong identifiability,” parameter recovery in over-specified settings is faster (at least $n^{-1/4}$), as the absence of a linear term in the gate simplifies the system’s derivative structure.
• Failure with Linear Experts: Linear expert functions ($h(x; \beta) = \beta_1^\top x + \beta_0$) break identifiability, leading to degraded rates $O(1/\log n)$, implying quadratic gates must be paired with nonlinear experts for statistical efficiency (Akbarian et al., 2024).

4. Gate and Model Design Recommendations

Empirical and theoretical analysis yield several design guidelines for practitioners:

• Prefer monomial over polynomial gates: The monomial variant (no linear term) avoids slow convergence in over-specified cases and aligns with standard Transformer query/key configurations.
• Use nonlinear experts: Experts implemented as neural networks with ReLU or tanh activations satisfy identifiability and promote efficient learning.
• Low-rank quadratic parameterization: To limit model size, $A_i$ can be factored as $U_i V_i^\top$ with $U_i, V_i \in \mathbb{R}^{d \times r}$ and $r \ll d$, reducing parameters per expert from $O(d^2)$ to $O(dr)$, which is favorable for large-scale models.
• Active-Attention mechanism: Inspired by the MoE-attention connection, enhancing the value head with a nonlinearity (e.g., GELU, ReLU, tanh) in

$$\mathrm{Act\text{-}Att}(Q, K, V) = \operatorname{Softmax}(QK^\top/\sqrt{d}) \; \bar\sigma(V)$$

matches the benefits of nonlinear experts and increases expressivity of aggregation (Akbarian et al., 2024).

5. Empirical Benchmarks and Performance

Comprehensive evaluation demonstrates superior sample efficiency and performance of quadratic gating structures, particularly when paired with nonlinear experts and active-attention strategies.

Synthetic Regression

• On one-dimensional regression data from a 3-expert quadratic-polynomial MoE, monomial gates with ReLU experts achieved convergence at the $-1/2$ slope in log–log plots, confirming theoretical predictions.
• Polynomial gates or linear experts displayed slower convergence, in line with theory.

Vision, Language, and Time-Series Tasks

Model Variant	Dataset	Evaluation Metric	Performance Impact
ViT/CaiT (Active-Att)	CIFAR-10, ImageNet	Top-1 accuracy	+1–2 pp over baseline
Transformer/Performer	WikiText-103	Perplexity	Lower with ReLU/GELU value heads
PatchTST, Transformer	8 Time-series datasets	MSE	5–10% improvement with nonlinearities

• Nonlinear value maps (e.g., tanh, sigmoid) were particularly effective in smaller-scale/high-overfitting regimes.

Quadratic vs Linear Gating in GPT-2 MoE

A direct comparison was conducted on GPT-2 (124M, 10B tokens, FineWeb-Edu), using a dense model, an MoE with linear Top-2 router, and an MoE with rank-32 quadratic Top-2 router, with identical active-parameter budgets.

• Dense GPT-2: validation loss 3.0211, accuracy 30.74%
• MoE (linear): validation loss 2.9928, accuracy 29.95%
• MoE (quadratic): validation loss 2.9712, accuracy 32.11%

Quadratic gating led to both lower loss and higher zero-shot accuracy on the HellaSwag benchmark (Akbarian et al., 2024).

6. Interpretive Insights and Broader Implications

The adoption of quadratic gating—especially the monomial type combined with nonlinear experts—improves MoE architectures by increasing the expressivity of routing functions to match the pairwise similarity structure exploited by attention mechanisms. This approach preserves fast parametric convergence rates and yields consistent gains in large-scale vision, language, and temporal modeling. A notable connection to self-attention suggests expanding the toolkit for efficient expert routing within and beyond Transformer architectures. Theoretical identifiability considerations strongly disfavor linear experts for quadratic gating, guiding future architectural design (Akbarian et al., 2024).