---
title: Amortized Model Ensembling (AME)
url: https://www.emergentmind.com/topics/amortized-model-ensembling-ame
type: topic
---

# Amortized Model Ensembling (AME)

Amortized Model Ensembling (AME) is a meta-optimization paradigm for synthesizing the benefits of model ensembling while reducing the computational or memory cost over conventional pointwise ensemble averaging. AME encompasses both stochastic inference-time ensembling methods—such as the Mixture-model-like Ensemble (ME) for autoregressive large language models—and parameter-space aggregation strategies—such as gradient-based neural averaging—unified by the principle of amortizing the ensemble computation over time, space, or iterations. AME enables the construction of ensemble-level predictions or solutions at significantly reduced inference or integration cost, and generalizes classic uniform model soup and mixture-of-experts frameworks in both theoretical and practical settings [2605.00419][2508.14832].

## 1. Formalization and Unification Framework

AME is instantiated in two complementary settings:  (a) output-space stochastic ensembling, where a sequence of forward passes or samples is amortized across base models (e.g., ME for LLMs), and (b) parameter-space aggregation, where a single set of model weights is synthesized from multiple experts via pseudogradient meta-optimization (data-free neural averaging).

### Output-space AME: Mixture-model-like Ensemble (ME)
Given fine-tuned models $M_1, \ldots, M_K$ with next-token distributions $p_k(x_t \mid x_{<t})$, the conventional ensemble distribution is:
$$
p_{ens}(x_t \mid x_{<t}) = \frac{1}{K} \sum_{k=1}^K p_k(x_t \mid x_{<t}).
$$
ME instead draws, at each token step, a single model $M_m$ (with probability $\pi_m = 1/K$) and samples $x_t \sim p_m(\cdot\,|\,x_{<t})$. The marginal next-token distribution under ME is provably equal to $p_{ens}$ [2605.00419].

### Parameter-space AME: Neural Averaging via Meta-Optimization
Given $N$ pretrained DNNs with weights $x_1, \ldots, x_N \in \mathbb{R}^d$, AME defines “neural averaging” by performing stochastic gradient-based aggregation over a quadratic proxy loss,
$$
f(x; \xi) = \frac{1}{2}\left(\|x - \xi\|^2 - \|\xi\|^2\right), \quad \xi \sim \mathcal{D},
$$
with each $x_i$ viewed as an independent draw from the weight distribution $\mathcal{D}$. AME generalizes model soup, optimizer-augmented ensembling, and meta-adaptive aggregation [2508.14832].

## 2. Mathematical Guarantees and Equivalence Properties

### Output Equivalence (ME)
Sampling each token via ME yields the same marginal distribution as the explicit ensemble:
$$
\Pr_{ME}[x_t = y | x_{<t}] = \sum_{m=1}^K \pi_m \, p_m(y | x_{<t}) = p_{ens}(y | x_{<t}),
$$
establishing the strict equivalence of ME’s amortized sampling to full ensemble sampling [2605.00419].

### Parameter Aggregation Guarantees (AME)
AME admits the following theoretical properties [2508.14832]:
- Plain gradient descent ensembling (with learning rate $\eta_i=1$, amplification $\zeta_i=1$) over all $x_i$ recovers uniform model soup: $w_N = \frac{1}{N} \sum_{i=1}^N x_i$.
- More general convex combinations and adaptive soups are achievable via tailored learning rates and pseudogradient scaling.
- Under suitable decaying learning rates ($\eta_t = O(t^{\alpha}),\, \alpha<-1$) and bounded $\zeta_t$, AME with AdaGrad or Adam converges with the ensemble parameter trajectory contained in the convex hull of the ingredients.
- Model soup converges in probability to the expected value of $\mathcal{D}$.

## 3. Algorithmic Formulations and Implementation Details

### ME for LLMs (Pseudocode)
```
Algorithm Mixture-model-like Ensemble (ME)
Input: Base models M₁,…,M_K; uniform πₘ = 1/K; initial prefix x_{<t}
S ← x_{<0}
while not end-of-sequence do
    m ∼ Cat(π₁,…,π_K)
    P ← Mₘ(· | S)
    x_t ∼ P
    S ← S ⧺ x_t
end while
return S
```
Lazy KV cache synchronization ensures that each model’s cache is updated only when selected; missing tokens since last selection are “prefilled” in a batched pass [2605.00419].

### AME in Parameter Space (Pseudocode)
```
Algorithm AME (any optimizer)
Input: ingredient weights {x₁…x_N}, pivot x_pivot, schedules ηᵢ, ζᵢ, optimizer
w ← x_pivot
for epoch = 1…E:
    for each i in some order:
        g ← ζᵢ · (w − xᵢ) / N
        Ĝ ← OptimizerUpdate(g)
        w ← w − ηᵢ · Ĝ
return w
```
Optimizer and schedule selection enables adaptation to task, stability, and exploration of the weight simplex [2508.14832].

## 4. Computational Trade-offs and Empirical Results

### Inference-Time Speedup (ME)
For a $K$-model ensemble and $T$-token sequence:
- **Conventional explicit ensemble:** $K \cdot T$ forward passes; $\mathcal{O}(K \cdot F)$ cost per token.
- **Mixture-model-like ensemble:** $T$ forward passes; $\mathcal{O}(F)$ cost per token.
Observed speedups for ME versus conventional ensembling are $1.78\times$–$2.68\times$ on NVIDIA H100, RTX 3090, A100, and V100 for two- and three-model LLM ensembles [2605.00419].

### Test Performance Parity
ME statistically matches conventional ensemble decoding within $\pm$1 point. Example entries:

| Task     | Best Single | CE (k=5)  | ME (k=5) |
|----------|-------------|-----------|----------|
| GSM8K    | 79.77       | 83.14     | 82.97    |
| MMLU     | 66.75       | 66.05     | 65.61    |
| BBH      | 51.94       | 52.74     | 53.04    |
| ARC      | 81.81       | 81.14     | 81.12    |

### Parameter-space AME: Out-of-domain and Robustness
AME outperforms both best single expert and uniform model soup in OOD regimes. For 50% OOD on CIFAR-10:
- Best Expert: $49.65\%$
- Model Soup: $96.64\%$
- AME (Adam): $97.09\%$ [2508.14832].

## 5. Connections to Routing, Model Soup, and AME Principles

AME provides a unifying perspective over various ensembling, mixture-of-experts, and meta-aggregation methods:
- In output space, ME is a special case of token-level routing where routing decisions are made uniformly at random, rather than via a learned selector $R(x_{<t})$.
- Model soup is realized as a single epoch of gradient descent ensembling in AME; the “pivoted pseudogradients” and adaptivity enable nonconvex or prioritized mixtures.
- The cost of ensembling (full $K$-fold computation) is amortized: ME achieves ensemble-level output statistics at only $1/K$ the per-token cost; AME synthesizes a single model representing ensemble knowledge without storing all experts [2605.00419][2508.14832].

## 6. Limitations and Research Directions

### Known Limitations
- ME is strictly valid only for random-sample decoding, not deterministic greedy/argmax decoding, where $\mathrm{argmax}(\mathrm{average}) \neq$ random argmax.
- Memory cost remains $K \times$ model size; lazy KV cache synchronization can strain memory for large $K$.
- Performance benefits plateau with highly correlated or redundant base models.
- In AME, theoretical analysis of hyperparameter-induced “phase transitions,” calibration, and layer/block-wise aggregation are unresolved [2605.00419][2508.14832].

### Open Directions
- Non-uniform weights in ME (by altering $\pi_m$) for quality/efficiency tradeoff.
- Routing and learned randomization merge ME and mixture-of-experts, possibly via differentiable routers balancing latency and accuracy.
- Adaptive $K$: dynamically adjust ensemble size based on computational budget.
- Greedy and beam heuristics: approximate ME for non-sampling decoders.
- In parameter-space AME, extensions include “FedSoup” for federated aggregation, calibration analyses, and neural averaging of inputs or architectural blocks.

## 7. Applications and Impact

AME underpins substantial new efficiency/quality tradeoffs for large model deployment and federated data-free model synthesis:
- For LLMs, ME achieves nearly identical generation quality to conventional ensemble decoding with $1.8\times$–$2.7\times$ speedup on current hardware.
- Parameter-space AME produces robust neural aggregates for OOD generalization, memory preservation, and prototype synthesis, exceeding prior uniform or greedy soup baselines [2605.00419][2508.14832].
- These methods provide a foundation for scalable, privacy-preserving, and adaptive model sharing or deployment.

AME stands as a mathematically principled, empirically validated, and extensible framework that amortizes ensemble-level performance into practical, computationally tractable algorithms across both weight- and output-space modalities [2605.00419][2508.14832].

Source: https://www.emergentmind.com/topics/amortized-model-ensembling-ame