Local Mixtures of Experts: Essentially Free Test-Time Training via Model Merging

TL;DR

Test-Time Model Merging (TTMM) scales the mixture-of-experts paradigm to orders of magnitude more experts by (i) clustering training data into many local neighborhoods at train time, (ii) fitting one small LoRA adapter per cluster, and (iii) at test time, dynamically selecting a sparse subset of adapters and averaging their parameters into a single task-specific model before decoding. TTMM is derived as an amortized approximation of test-time training (TTT) — the paper shows formally that in the limit of many experts, TTMM approaches the TTT objective — but pays nearly zero test-time overhead (roughly 20 tokens’ worth of latency for a 1B base model with 10 active experts, versus ~seconds for TTT gradient steps). On Wikipedia and GitHub-Python next-token perplexity with Llama-3.2-1B and Qwen-2.5-1.5B, TTMM with 10 active experts approaches TTT’s accuracy and beats a single fine-tuned multitask model.

Key claims

TTMM approximates TTT: for a Lipschitz network with smooth loss and single-step gradient descent adaptation, an expert trained on any cluster containing at least one $k$ -NN of the prompt embedding is close in loss to the TTT model trained on the $k$ -NN set [Prop. 3.1 / §3.2].
Merging model parameters (single forward pass) roughly matches merging model predictions (ensembling, multiple forward passes) on both eval corpora, at a fraction of the test-time compute [Table 1, Insight 1].
Merging coefficients computed by sparse softmax cross-attention between the prompt embedding and cluster centroids significantly outperform uniform weighting across active experts; on Wikipedia, uniform weighting removes the entire benefit of merging more than one expert [Fig. 5, Insight 2].
Scaling from 1 to ~100 experts materially improves perplexity over a single multitask-fine-tuned model, but scaling further to ~1000 experts saturates on the two evaluation corpora — attributed to either redundant clusters or increased interference at merge time [Insight 3, Fig. 6].
Local experts substantially outperform both the base model and other experts on their own cluster’s holdout data, confirming that the clustering + local-LoRA step is doing real work beyond generic fine-tuning [Fig. 6 (left), Insight 4].
Test-time overhead decomposition for Llama-3.2-1B with 10 active experts: ~30 ms select + ~60 ms CPU→GPU load + ~10 ms merge ≈ 100 ms total, versus ~seconds for a few dozen TTT gradient steps — a claimed >100× test-time speedup over TTT at similar accuracy [Fig. 4, §3.1].
Einsum-based batched merging of the LoRA low-rank factors is ~1000× faster than a naive per-expert for-loop with no extra memory [Fig. 7, §3.1].

Method

At train time, TTMM applies bisecting k-means to the (mean-pooled sentence-embedding) representations of all training documents, producing a partition of the corpus into $E$ clusters. For each cluster, a small LoRA adapter is trained on the base LLM for one epoch on that cluster’s documents. Cluster centroids (mean of the normalized document embeddings, then re-normalized) are cached.

At test time, given a prompt with embedding $q$ , TTMM computes merging coefficients via a sparse cross-attention over cluster centroids: $w_i \propto \exp(q \cdot c_i / \tau)$ then keep only $w_i > \varepsilon / E$ and renormalize. The surviving experts (typically 3–10 out of hundreds) are loaded from CPU to GPU, their LoRA low-rank factors are combined into a single merged LoRA via a batched einsum, and decoding proceeds with the merged model — a single forward path, no ensembling. The paper shows the RBF-kernel form of this weighting and connects it to classical local learning (Nadaraya–Watson, RBF networks). Scaling to hundreds of experts on a 1B base only requires ~MB of LoRA storage per cluster, kept on CPU.

The formal analysis (§3.2, Prop. C.1) breaks the TTT ↔ TTMM approximation into three steps: (1) train on clusters rather than every possible $k$ -NN subset, (2) summarize clusters by centroids for retrieval, (3) merge multiple experts to compensate for the imperfect cluster fit.

Results

Wikipedia perplexity (Llama-3.2-1B base 9.243): single fine-tune 8.116; TTMM-1 expert 7.740; TTMM-3 experts (merging) 7.571; TTMM-10 (merging) 7.510; TTT lower bound 7.474. Merging within 0.05 PPL of ensembling at each active-expert count [Table 1].
GitHub Python perplexity (Llama-3.2-1B base 2.750): single fine-tune 2.601; TTMM-1 2.559; TTMM-3 (merging) 2.519; TTMM-10 (merging) 2.492; TTT 2.394. Same qualitative pattern on Qwen-2.5-1.5B [Table 1].
Sparsity sweep: $\varepsilon = 0.05$ gives the best accuracy / test-time-cost tradeoff and is used elsewhere [§4 Insight 2, Fig. 5].
Expert count scaling: significant gains from $E \in \{1, \sim 100\}$ ; no additional gain from $E = 1000$ on these corpora [Insight 3, Fig. 6].
Cluster locality: an expert’s cluster-holdout PPL is substantially lower than both the base model and other experts’ — evidence that specialization is real, not an averaging illusion [Fig. 6 (left)].

Why it’s interesting

Sits in the intersection of three wiki threads that have so far been treated separately: (1) it’s a direct amortization of TTT and belongs in the Long context as weights cluster next to End-to-End Test-Time Training for Long Context (SGD on backbone), Doc-to-LoRA: Learning to Instantly Internalize Contexts (hypernetwork emits document-LoRA), and Self-Adapting Language Models (RL policy generates finetuning data) — TTMM is the pre-trained expert pool + merge-on-demand pole of the same design space, cheaper at test time than any of them because there are no gradient steps and no hypernetwork forward pass. (2) It’s a scale-up of the MoE paradigm to two orders of magnitude more experts than mainstream MoE routers handle, with merging replacing conditional routing — adjacent to the router-redesign thread in MoE Routing Design but reached from the test-time-training direction rather than the routing-collapse-and-LBL direction. (3) The learned-basis-of-LoRAs framing complements Spanning the Visual Analogy Space with a Weight Basis of LoRAs (LoRWeB): both argue “many small LoRAs + a router beats one big LoRA,” but TTMM’s router is a train-time centroid cache, and its “basis” is grown by clustering the training corpus, not learned end-to-end. The paper’s central negative result — merging saturates around 100–1000 experts on these corpora — is the first quantitative datapoint that “TTT-quality via merging” hits a data-diversity ceiling before it hits a compute ceiling.