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

Thank you for your review and detailed comments! You can find below our detailed response to your questions and concerns. Please let us know if you have any further questions or suggestions.

"The idea of TTMM is sound, and the experiments also verify the effectiveness of TTMM on Wikipedia and GitHub (python) in perplexity. It is likely to be a practical paradigm for test-time approaches."

Thank you!

Evaluation on non-perplexity tasks

Thank you for this feedback! Based on this feedback and the feedback of reviewer Ltqr, we evaluated TTMM on MMLU for the rebuttal. We have included the results in the main comment above. In summary, we find that TTMM slightly improves MMLU accuracy.

MoEs with larger numbers of experts

Thank you for mentioning larger MoEs such as DeepSeek-V3. For perspective, we find improvement in using 100 experts for comparatively small individual domains, such as English Wikipedia or Python code from GitHub. In contrast, DeepSeek-V3 is trained across many broad domains. We believe therefore that our results indicate that further increasing the number of experts in LLMs is a promising direction. Doing this, however, in a standard MoE-based architecture is challenging since all experts need to be in GPU memory simultaneously, which requires substantial infrastructure for an increased number of experts. We propose and study the alternative of training all individual experts in an embarrassingly parallel way, and then linking them post-hoc via a router. This provides the potential advantage of requiring substantially less resources. We very much appreciate your suggestion of scaling to larger LLMs, which we also believe to be an exciting direction for future work.

Selection of number of experts and number of active experts

We agree that the selection of these is non-trivial and therefore we will add an expanded section discussing this choice to the final paper.

Choosing the number of experts ($K$): The total number of experts has almost no direct effect on latency. However, more experts require a large amount of static CPU memory. Further, increasing the number of experts also increases the separation of knowledge/skills between experts. This then requires more active experts to achieve the same performance, which also increases latency. For this reason, we agree that selecting $K$ for optimal placement on the Pareto frontier of latency and performance is non-trivial. We observe that choosing 100 experts is also indicated by the ``elbow method'', common in clustering (cf. Figure 6 (middle)).

Choosing the number of active experts ($N$): The number of active experts is slightly more straightforward to select. Note that TTMM chooses the number of active experts adaptively depending on the prompt. The average number of active experts is determined by the temperature parameter $\beta$, which we ablate in Figure 5. We suggest selecting this parameter on a holdout dataset (which one can interpret as a form of ``training'' the router), and find that performance is relatively robust to the exact value of $\beta$.

Discussion of improvement of TTMM over TTT on Wikipedia

We agree with you. We were also surprised to see TTMM outperform TTT on the Wikipedia dataset. We expect that this may be because TTMM can incorporate knowledge from a larger set of experts, each trained on a few thousands of data points. In contrast, TTT is only training on the 100 most related neighbors. If there are more than 100 related data points, this may miss some information. We added a mention of this to the discussion of our results.

We hope to have addressed your remaining concerns. We would greatly appreciate it if you could reconsider the score based on our response and extended results.

Dear reviewer,

We wanted to kindly remind you that the discussion period on OpenReview will close in a few days. We hope to have addressed your remaining concerns with our response to your review and would be happy to discuss if you have any further questions.

Latency of TTMM on our hardware

Thank you for your question. We have measured the average latency of TTMM on our hardware (Wikipedia dataset, Llama 3.2-1B). The results are summarized below. We will include these numbers in the appendix. Please note that latency is expected to be lower on newer clusters with higher CPU-GPU bandwidth.

Train-time: 125h

Embedding generation-time: 2h

Test-time (avg): 15.7ms (Note that this is not generating tokens, but only evaluating their perplexity)

Thank you again for your time and effort in reviewing our work!

Thank you for your review and detailed comments! You can find below our detailed response to your questions and concerns. Please let us know if you have any further questions or suggestions.

"This method significantly reduces the real-time overhead during test-time while improving the performance (perplexity) of language models. This is a desirable characteristic for applications, and the method could be particularly useful for applications that require real-time performance."

Thank you!

Title: "Essentially Free Test-Time Training via Model Merging"

TTMM indeed requires storing the LoRAs on CPU memory. We agree that this is a limitation of our proposed method, and we added the following to the first paragraph of the discussion section: "A limitation of TTMM is that it requires additional CPU memory for storing all experts." Nevertheless, we would like to emphasize here that TTMM does not incur significant latency beyond standard inference, while improving performance.

Evaluation on non-perplexity tasks

Thank you for this feedback! Based on this feedback and the feedback of reviewer tuLa, we evaluated TTMM on MMLU for the rebuttal. We have included the results in the main comment above. In summary, we find that TTMM slightly improves MMLU accuracy.

Scaling the number of experts

We expect that scaling the total number of experts (for a fixed dataset) increases the separation of knowledge/skills between experts, which then requires more experts to be active during inference to achieve the same performance. Therefore, in terms of the Pareto frontier of latency and performance, fewer total experts might be beneficial since they may require fewer active experts to match performance.

We will clarify this in our discussion.

Model merging interference

We find that as the number of active experts is increased (significantly larger than 10), the performance gain of TTMM plateaus. We think that this is most likely due to overlapping signals propagated through the individual expert models during inference, which interfere in the merged model, as also studied in a large body of work in model merging [e.g., 1, 2]. This is further supported by our observation that ensembling the logits of experts outperforms model merging (albeit at a much higher latency). While our study focuses on how to best weight the active experts based on their relevance to the test task, we fully agree that studying interference is an important direction for future work. One interesting approach for further study might be to explicitly guide the training of the individual experts towards solutions that have small interference.

Analysis of clustering method

Next to bisecting k-means, we also clustered the Wikipedia dataset using regular k-means. K-means resulted in an unbalanced clustering, with one cluster capturing $>1/4$ of data points and many clusters capturing $<100$ data points. This observation motivated our decision to use bisecting k-means.

Selection of embeddings

Thank you for this question! In an attempt to answer this question, we evaluated different embedding models for our new evaluation on MMLU. We find that TTMM outperforms the base model regardless of the choice of embedding model. In particular, the largest Qwen2-1.5B embedding model performs worse than the 100M MPNet model. We estimate that this might be because finding a decent clustering might be significantly easier than solving the questions, since the former only requires a broad understanding of the meaning of some keywords while the latter requires some ability of reasoning.

Single-step assumption in theoretical analysis

Thank you for this question! We adopted this assumption purely to simplify notation. Indeed, with slightly more complex notation to keep track of the individual steps of gradient descent, our proof of Proposition C.1 also proves the recursion $\|\theta_{x^\star}^{(t+1)} - \theta'^{(t+1)}\| \leq \|\theta_{x^\star}^{(t)} - \theta'^{(t)}\| + \eta G (\text{diam}(\mathcal{D}_{x^\star}) + \text{diam}(\mathcal{D}'))$ (cf. lines 535-536). Unrolling this recursion proves the same bound as in Proposition C.1 for $T$ steps of gradient descent, with an additional factor $T$ in the bound.

We agree with you that this was an oversimplification, and we have changed the proposition in the paper to the variant with $T$-step gradient descent. Thank you for this feedback!

References:

1. Yadav et al. Ties-merging: Resolving interference when merging models. NeurIPS, 2023.
2. Yang et al. AdaMerging: Adaptive Model Merging for Multi-Task Learning. ICLR, 2024.

Effect of LoRA

Thank you for this question! We adopt LoRA primarily because CPU memory and loading from CPU memory is a central bottleneck for TTMM. With Llama-3.2-1B as base and 100 experts, LoRAs take around 17 GB while full-parameter experts would require 247 GB (2.47 GB per expert). Moreover, at test-time, active experts have to be loaded into GPU memory. Loading 10 LoRAs is only around 1.7 GB, while loading 10 full-parameter experts is 24.7 GB. This would substantially increase latency.

Discussion of improvement of TTMM over TTT on Wikipedia

We agree with you. We were also surprised to see TTMM outperform TTT on the Wikipedia dataset. We expect that this may be because TTMM can incorporate knowledge from a larger set of experts, each trained on a few thousands of data points. In contrast, TTT is only training on the 100 most related neighbors. If there are more than 100 related data points, this may miss some information. We added a mention of this to the discussion of our results.

Prefix length for GitHub

Below is an ablation of the prefix length for Python code, showing the perplexity of TTMM and the base model, respectively.

Since code does not typically summarize the contents of a file in the first few tokens, the prefix must be longer than for Wikipedia to provide sufficient information on the content of the suffix. Code typically begins with a series of import statements or copyright comments, which alone are not sufficiently informative for the selection of experts. To address this limitation of our evaluation, we added an evaluation on MMLU (see above), which follows a natural structure of prompt & response.

We hope to have addressed your remaining concerns. We would greatly appreciate it if you could reconsider the score based on our response and extended results.

Dear reviewer,

We wanted to kindly remind you that the discussion period on OpenReview will close in a few days. We hope to have addressed your remaining concerns with our response to your review and would be happy to discuss if you have any further questions.

Thank you again for your time and effort in reviewing our work!

Thank you for detailed comments to my review! They were basically useful to gain my understanding of the proposed method, and some questions are resolved appropriately by the authors' comment.

The comment "Effect of LoRA" doesn't resolve one of my question: "Is there any consideration of whether synthesizing a large number of trained LoRAs can maintain the same model expressiveness as tuning the original parameter space directly?", because the answer only mentioned about the efficiency. This is unfortunate, but is relatively trivial considering that the comments were overall appropriate.

According to the consideration of the rebuttal, I'd like to give a stronger recommendation for this paper.

Thank you for your review and detailed comments! You can find below our detailed response to your questions and concerns. Please let us know if you have any further questions or suggestions.

Comparison of latency vs performance between TTMM and TTT

Thank you for highlighting the importance of discussing the Pareto frontier of latency and performance. Figure 1 shows that TTMM has only marginally higher latency than standard inference, while significantly improving performance. In contrast, Figure 1 shows that TTT has significantly higher latency. We further discuss latency in Section 3.1 and find that TTMM has a constant overhead of around 115ms (with 10 active experts), while TTT has a constant overhead of around 15s, over 100x larger.

The exact Pareto frontier varies depending on the generation length. Figure 1 plots this for 100 generated tokens and Figure 4 (left) shows how the latency changes for other generation lengths. We do not show TTT in Figure 4 (left) since the constant overhead of TTT before generating any token is larger than the time to generate 1k tokens with TTMM (and 10 active experts).

We summarize this in the following table, which we will also add to the paper. Here, $T$ denotes the time for generation. On the hardware we use to measure latency (NVIDIA RTX 4090), generation takes around 6ms/token.

Figure 4b

Thank you for this question! Figure 4 (middle) shows the latency of TTMM at varying numbers of active experts. At 100 active experts, the loading of experts from CPU memory dominates latency. For this reason, we focus our performance evaluation on fewer active experts (1, 3, and 10), where the latency of TTMM is small. Nevertheless, even at 100 active experts, the 1s latency of TTMM is still substantially faster than the roughly 15s latency of TTT.

Training of the router

Your suggestion to train the router is great! We believe that training the router on holdout training data is likely to further improve performance. While training the router is common practice with MoEs, we did not train the router since this would significantly increase the cost and complexity of training if experts do not all fit into GPU memory simultaneously. We find that even when we train each expert in an embarrassingly parallel way, TTMM achieves good performance. Based on our findings, we find your suggestion to be an exciting direction for future work.

Discussion of the number of experts

Thank you for this question!

Considerations for selecting the number of clusters ($K$): The total number of experts has almost no impact on latency. It primarily affects the required CPU memory for storing all experts. Provided sufficient memory for storing the experts, $K$ can be scaled to billions of experts without significantly increasing latency, similarly to retrieval systems of documents. We select $K$ using the elbow method based on the bisecting k-means clustering.

Considerations for selecting the number of active clusters ($N$): The number of active experts has the strongest effect on latency, since all active experts have to be loaded from CPU into GPU memory before inference. We find that until $10$ active experts, this latency is minimal.

Dear reviewer,

We wanted to kindly remind you that the discussion period on OpenReview will close in a few days. We hope to have addressed your remaining concerns with our response to your review. We would greatly appreciate it if you reconsider the score based on our response and extended results.

Thank you again for your time and effort in reviewing our work!