Softmax–Laplace Model in HMoE

• The Softmax–Laplace Model is a gating mechanism in HMoE architectures that distinguishes between Softmax and Laplace functions to route expert subnetworks.
• It eliminates critical parameter interactions, yielding accelerated convergence for over-specified experts and improved specialization across multimodal and vision tasks.
• Empirical and theoretical analyses confirm that full Laplace gating (LL) outperforms other configurations by decoupling mean–variance interactions and enhancing performance.

The Softmax–Laplace Model refers to a class of gating mechanisms for Hierarchical Mixture-of-Experts (HMoE) architectures, where “gating” networks select expert subnetworks via parametric functions. Critically, this framework distinguishes between the traditional Softmax gating function and a Laplace gating variant. Systematic analysis demonstrates that substituting Laplace gates for Softmax—in particular at both hierarchy levels—removes fundamental parameter interactions, yielding accelerated convergence for over-specified experts and improving expert specialization. These findings are theoretically established and empirically validated across multimodal, image classification, and domain generalization tasks (Nguyen et al., 3 Oct 2024).

1. Formal Definitions and Notation

Consider a two-level HMoE with real input $x\in\mathbb{R}^d$ and scalar output $y\in\mathbb{R}$. Each gating function at both hierarchy levels produces a sparse expert mixture.

• Softmax Gating (“S”): For expert $i$,

$$s_i(x) = w_i^\top x + b_i, \quad g_i(x) = \frac{\exp(s_i(x))}{\sum_j \exp(s_j(x))}$$

where $w_i \in \mathbb{R}^d$, $b_i \in \mathbb{R}$, and $g_i(x)$ is the selection weight.

• Laplace Gating (“L”): For expert $i$,

$$s_i(x) = w_i^\top x + b_i, \quad \ell_i(x) = \frac{\exp(-|s_i(x)|)}{\sum_j \exp(-|s_j(x)|)}$$

• HMoE Architecture: With $k_1$ first-level and $k_2$ second-level experts (indices $i_1$ and $i_2$), the conditional output density:

$$p(y\mid x) = \sum_{i_1=1}^{k_1} \pi^{(1)}_{i_1}(x) \sum_{i_2=1}^{k_2} \pi^{(2)}_{i_2|i_1}(x) \, \mathcal{N}(y\mid \eta_{i_1i_2}^\top x + \tau_{i_1i_2}, \nu_{i_1i_2})$$

where each $\pi^{(1)}$ and $\pi^{(2)}$ can be Softmax or Laplace, and the expert is Gaussian with learned mean and variance.

2. Theoretical Properties and Estimation Rates

Three gating configurations are distinguished:

• SS: Softmax at both levels
• SL: Softmax outer, Laplace inner
• LL: Laplace at both levels

Conditional Density Estimation

Under standard compactness and identifiability assumptions, all schemes achieve parametric conditional-density estimation rates: $$\mathbb{E}_X [h(p_{\hat G_n}(\cdot|X), p_{G_*}(\cdot|X))] = \widetilde{O}(n^{-1/2})$$ where $h$ is the squared Hellinger distance.

Expert Specialization and Voronoi Loss

A refined Voronoi-loss $\mathcal{L}(G, G_*)$ quantifies how closely fitted experts approximate true atoms:

• Exact-specified (one fitted per true):

$$\|\hat\eta_{i_1i_2} - \eta_{i_1i_2}^*\| = \widetilde{O}_P(n^{-1/2})$$

• Over-specified (multiple fitted per true):
  • SS, SL: $\widetilde{O}_P(n^{-1/r(m)})$, $r(2)=4$, $r(3)=6$
  • LL: $\widetilde{O}_P(n^{-1/4})$ for all $m\ge2$

Gating	Exact-specified	Over-specified
SS	$n^{-1/2}$	$n^{-1/r^{SS}(m)}$
SL	$n^{-1/2}$	$n^{-1/r^{SL}(m)}$
LL	$n^{-1/2}$	$n^{-1/4}$

Substituting Laplace at the inner level only (SL) does not break the mean–bias–variance interactions underlying the slow rates; only the full Laplace–Laplace (LL) configuration eliminates these and achieves accelerated over-specified convergence.

Underlying Mechanisms

• Under Softmax gating, parameter interactions are encoded in identities such as $$\partial u/\partial\eta = \partial^2 u/\partial a\,\partial\tau$$, leading to slow expert convergence rates.
• With Laplace at both levels, these identities vanish, leaving only the standard Gaussian-mean/variance interaction, and yielding the $n^{-1/4}$ rate for over-specified experts.

3. Model Implementation and Training

Computational Workflow

The core forward algorithm is as follows:

D_o, C_o, L_o = Gate_outer(x) x_outer = Dispatch(x, D_o) D_i, C_i, L_i = Gate_inner(x_outer) x_expert = Dispatch(x_outer, D_i) y_expert = Experts(x_expert) y_inner = Combine(y_expert, C_i) y_final = Combine(y_inner, C_o) Loss = task_loss(y_final) + lambda * (L_o + L_i)

• Gate_outer/Gate_inner produce soft or sparse tensors via either Softmax ($g_i$) or Laplace ($\ell_i$).
• Experts are typically small independent FFNs.
• Regularization includes batchwise expert capacity constraints and a load-balancing loss:

$$L_{\textrm{gate}} = \lambda \sum_{i=1}^E \left(\mathbb{E}_x [\pi_i(x)] - \frac{1}{E}\right)^2$$

with $\lambda \approx 0.1$.

Gradient Computation and Initialization

• For Laplace: $$\frac{\partial \ell_i}{\partial s_i} = -\operatorname{sign}(s_i)\ell_i + \ell_i\sum_j \operatorname{sign}(s_j)\ell_j$$
• Gating biases and conditional weights are zero-initialized, with small random initialization for weights.
• Training uses Adam optimizer, learning rate $1\mathrm{e}{-4}$, weight decay $1\mathrm{e}{-5}$, dropout 0.1, typically for 100 epochs.

4. Empirical Evaluation

Multimodal Fusion: MIMIC-IV

• Modalities: vital-signs, chest X-ray (DenseNet-121), clinical notes (BioClinicalBERT)
• Tasks: 48h in-hospital mortality (48-IHM), length-of-stay (LOS), 25-label phenotype (25-PHE)
• Architecture: 12 stacked two-level HMoE modules, $E_o=2, E_i=4$, residual connections

Method	48-IHM (AUROC/F1)	LOS (AUROC/F1)	25-PHE (AUROC/F1)
MoE	83.13 / 46.82	83.76 / 74.32	73.87 / 35.96
HMoE(LL)	85.59 / 47.57	86.26 / 76.07	73.81 / 35.64

HMoE (LL) outperforms all baseline methods.

Latent Domain Discovery

• Datasets: eICU (by region $\rightarrow$ domain), MIMIC-IV (by admission year), with or without CXR/notes.
• Tasks: readmission, post-discharge mortality
• Baselines: Oracle, Base, DANN, MLDG, IRM, SLDG

HMoE (SL) achieves top or near-oracle performance. Use of multimodal features (HMoE-M) further improves results.

Image Classification

• CIFAR-10/tiny-ImageNet (MoE layer): LL gating best by ~1–2% accuracy
• Vision-MoE (ViT backbone with 2 or 4 MoE layers) on CIFAR-10 / ImageNet: LL gating consistently best

Ablation Studies and Routing

• LL gating delivers more diversified expert assignments, particularly for over-specified configurations.
• Increasing number of inner experts ($E_i$) yields greater performance increases than increasing outer experts, with diminishing returns beyond $E_i=4\!-\!8$.

5. Interpretation, Limitations, and Future Directions

• The Laplace–Laplace gating combination in HMoE architectures universally accelerates over-specified expert convergence from $\widetilde{O}(n^{-1/r(m)})$ to $\widetilde{O}(n^{-1/4})$ by fully decoupling gating–expert parameter interactions.
• Empirical results in large-scale multimodal, domain generalization, and vision tasks consistently favor Laplace–Laplace over all other gating configurations.
• The hierarchical routing required for HMoE incurs additional computation and memory costs; future directions include model pruning or distillation to address this.
• The Softmax–Laplace (SL) configuration does not yield improved convergence; full Laplace gating (LL) at both levels is necessary.
• The precise scaling exponents $r^{Soft/Lap}(m)$ for larger $m$ remain an open problem closely related to algebraic geometry.
• Potential avenues include deeper HMoE hierarchies and alternative gating families, for instance, Student’s $t$ gating (Nguyen et al., 3 Oct 2024).