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., 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., 2024).