💬 NLP

Can Model Merging Improve Aggregation in DiLoCo?

This paper bridges the gap between model merging and distributed learning by demonstrating that adapting the Iso-C merging technique with Nesterov momentum (IsoLoCo) significantly outperforms existing DiLoCo methods in low-communication distributed training, particularly as the number of local workers increases.

Original authors: Stefan Horoi, Benjamin Thérien, Guy Wolf, Eugene Belilovsky

Published 2026-07-07

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Stefan Horoi, Benjamin Thérien, Guy Wolf, Eugene Belilovsky

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to train a giant, super-smart AI brain. To do this, you need a massive amount of computer power. Usually, companies use a "data-parallel" approach: they line up hundreds of computers side-by-side, all working on the exact same task at the exact same time, constantly shouting updates to each other. This is fast, but it requires expensive, high-speed cables connecting every single computer. If one computer is slow or the cable is weak, the whole line has to wait.

The Problem: The "Island" Approach
To save money and use weaker connections, researchers developed a method called DiLoCo. Imagine instead of a single line, you have many separate "islands" of computers.

Each island trains its own version of the AI brain independently for a while (let's say 30 steps).
They don't talk during those 30 steps.
After 30 steps, they send a summary of what they learned (a "pseudo-gradient") to a central hub.
The hub averages these summaries and updates the main brain.

The Glitch:
The paper found a flaw in this "Island" approach. As you add more islands (workers) or let them train longer before talking, the AI starts to get confused. It's like a group of people trying to draw a picture together without talking: if they work too long on their own, they start drawing different things, and when they try to combine their work, the result is messy. In technical terms, the "updates" from the different islands start to clash with each other, degrading the AI's performance.

The Inspiration: Merging Expert Models
Separately, another group of researchers was working on Model Merging. Imagine you have an expert chef who learned to make Italian food and another who learned to make Japanese food. You want one chef who can do both.

Old Way: Just average their recipes. This often fails because the ingredients clash (e.g., mixing soy sauce and heavy cream).
New Way (Task Arithmetic): Instead of averaging the whole recipe, you look at the difference between the expert's final recipe and the original base recipe. You then carefully combine just those differences.

The paper's big "Aha!" moment was realizing that DiLoCo is actually doing the same thing as Model Merging, just in a loop.

The "Base Model" is the AI brain from the last time everyone talked.
The "Islands" are the experts who trained on their own.
The "Updates" they send back are the "differences" (or task vectors).

The paper argues that the reason DiLoCo gets messy with many islands is the same reason merging recipes gets messy: Interference. The updates from different islands are fighting each other.

The Solution: IsoLoCo
The authors decided to borrow the "best recipe" from the Model Merging world to fix DiLoCo. They tested several merging techniques and found one called Iso-C (Isotropic merging) worked best.

Think of Iso-C like a "harmonizer" for the updates. When the islands send their updates back:

The harmonizer looks at the "shape" of the updates.
If one island is shouting too loudly in a specific direction (a specific mathematical direction), the harmonizer turns that volume down to match the average.
This prevents any single island from dominating or clashing with the others, creating a smoother, more balanced combined update.

They took this "harmonizer" and added a "momentum" feature (like a runner who keeps their speed up even when the road gets bumpy) to create a new method they call IsoLoCo.

The Results
The paper tested this on different sizes of AI models and with different numbers of islands (from 1 to 128).

The Winner: IsoLoCo consistently beat the original DiLoCo method.
The Gap: The more islands you added, the bigger the advantage IsoLoCo had. With 128 islands, the original method got quite messy, but IsoLoCo stayed clean and efficient.
The Speed: They also made a "fast mode" for IsoLoCo using a mathematical shortcut (Newton-Schulz iterations) that made the process much faster without losing quality.

In Summary
The paper shows that by treating distributed AI training like a "merging of experts" problem, we can use advanced math tricks to stop the computers from fighting each other. This allows us to train massive AI models across many weakly connected computers without losing performance, making large-scale AI training cheaper and more scalable.

Technical Summary: Model Merging for Aggregation in DiLoCo

Problem Statement

Distributed Low-Communication (DiLoCo) training has emerged as a method to train large language models across poorly connected compute clusters ("islands") by minimizing global synchronization. In DiLoCo, workers perform multiple local optimization steps (inner steps) before aggregating their parameter updates (pseudo-gradients) via a global optimizer. While effective at reducing communication costs, DiLoCo suffers from performance degradation as the number of workers ( $R$ ) or the number of local steps ( $H$ ) increases. This degradation occurs because local replicas diverge significantly, causing the averaged pseudo-gradients to misalign with the updates produced by tightly synchronized data-parallel training.

Concurrently, Model Merging techniques have been developed to combine independently fine-tuned models into a single model. A central challenge in merging is "task interference," where parameter updates from different tasks conflict, leading to performance drops, especially as the number of models to merge increases.

The authors identify a structural analogy between these two fields: the pseudo-gradient aggregation step in DiLoCo (with an SGD outer optimizer) is mathematically equivalent to Task Arithmetic model merging. In this analogy:

The shared global model at the start of an outer step corresponds to the pre-trained initialization.
The local worker replicas correspond to fine-tuned "expert" models.
The local updates (pseudo-gradients) correspond to "task vectors."
The number of workers corresponds to the number of tasks.

The core problem addressed is whether advanced model merging techniques, designed specifically to mitigate task interference, can be repurposed to improve the aggregation step in DiLoCo, thereby restoring performance as the number of workers or local steps increases.

Methodology

Theoretical Bridge

The paper establishes that DiLoCo's outer update with SGD is a direct application of task arithmetic. By viewing the aggregation of pseudo-gradients as a model merging problem, the authors propose substituting the standard averaging operation with sophisticated merging algorithms designed to reduce negative interference.

Experimental Evaluation of Merging Methods

The authors systematically benchmarked several state-of-the-art merging methods as drop-in replacements for DiLoCo's aggregation step. The methods evaluated included:

Sparsity-based: TIES (TrIm, Elect Sign) and DARE (Drop And REscale).
Orthogonalization-based: TSV (Task Singular Vector), Iso-C (Isotropic merging in Common subspace), and Iso-CTS.

Findings on Merging Methods:

Sparsity methods (TIES, DARE): Performed similarly to standard DiLoCo SGD only at very low sparsity levels. Performance degraded as sparsity increased, suggesting they did not offer a distinct advantage over simple averaging in this context.
TSV: Performed poorly (higher validation loss) and was computationally prohibitive, requiring multiple Singular Value Decompositions (SVDs) per step.
Iso-CTS: Outperformed the baseline but was computationally expensive and offered no significant advantage over Iso-C in the DiLoCo setting (where local tasks are short and similar).
Iso-C: Emerged as the most promising candidate. It improves performance by replacing the singular value spectrum of the aggregated pseudo-gradient with its mean, effectively flattening high-energy directions to yield a balanced common subspace. It requires only one SVD per parameter matrix, offering a favorable cost-performance trade-off.

Proposed Algorithm: IsoLoCo

Building on the success of Iso-C, the authors propose IsoLoCo (Isotropic Distributed Low-Communication). IsoLoCo adapts Iso-C for distributed training by:

Isotropic Correction: For 2D parameter matrices (weight layers), it computes the SVD of the averaged pseudo-gradient, replaces the singular values with their mean, and reconstructs the update. This regularizes the spectrum while preserving the subspace structure.
Nesterov Momentum: Unlike standard Iso-C, IsoLoCo incorporates Nesterov momentum into the outer optimizer. This is motivated by prior work showing momentum's benefit in FedOpt/DiLoCo settings. The corrected pseudo-gradient is fed into a momentum buffer to generate the final update.

Implementation Optimization

To address the computational cost of SVD, the authors developed a fast implementation of IsoLoCo using Newton–Schulz iterations. This replaces the standard SVD with iterative orthogonalization and uses an RMS proxy for the singular-value mean. This optimization yields up to two orders of magnitude faster orthogonalization with negligible loss degradation.

Key Results

Performance Across Worker Counts

Experiments were conducted on Llama-style models (178M, 512M, and 1B parameters) using the DCLM dataset.

Scaling with Workers: IsoLoCo significantly outperforms standard DiLoCo as the number of workers increases.
- At the 178M scale with $R=128$ , IsoLoCo reduced the validation loss gap relative to the data-parallel baseline by 5.8% compared to DiLoCo.
- At $R=128$ , IsoLoCo showed a 9.7% degradation relative to data-parallel training, whereas DiLoCo showed a 16.5% degradation.
- IsoLoCo even outperformed the data-parallel baseline at low worker counts ( $R \le 4$ ).
Model Size: The advantage of IsoLoCo over DiLoCo persists across model sizes (178M to 1B), though the relative improvement is slightly smaller for larger models (1.2% at 1B vs. ~2.4% at 178M), likely because larger models naturally scale better with worker counts.

Impact of Inner Steps

IsoLoCo was evaluated with increasing synchronization intervals ( $H = 30$ to $240$). It consistently outperformed DiLoCo across all intervals, demonstrating robustness to longer periods of local divergence.

Comparison with Momentum-based DiLoCo

Notably, the combination of DiLoCo SGD with Iso-C aggregation (without momentum) already outperformed the standard DiLoCo (which uses Nesterov momentum). Adding Nesterov momentum to Iso-C (creating IsoLoCo) further widened the performance gap, indicating that the merging-based aggregation and momentum mechanisms are complementary.

Compatibility with Muon

The paper notes that IsoLoCo can be composed with the Muon inner optimizer (MuLoCo). In this configuration, IsoLoCo achieved the lowest validation loss among all methods for $R \ge 32$ , while IsoLoCo with an AdamW inner optimizer matched or surpassed MuLoCo at high worker counts with significantly lower orthogonalization costs.

Significance and Claims

The paper claims that viewing low-communication distributed training through the lens of model merging provides a powerful framework for improving scalability. The primary contributions are:

Theoretical Unification: Establishing a direct analogy between DiLoCo's pseudo-gradient aggregation and task-arithmetic model merging.
Empirical Validation: Demonstrating that model merging techniques, specifically Iso-C, can serve as superior aggregation rules for distributed training, outperforming simple averaging and even momentum-based DiLoCo.
IsoLoCo: Proposing a practical, momentum-enhanced algorithm that significantly narrows the performance gap between low-communication training and data-parallel gold standards, particularly in regimes with many workers or long synchronization intervals.
Efficiency: Providing a Newton-Schulz-based implementation that makes the SVD-heavy merging process computationally feasible for iterative distributed training.

The authors conclude that merging-inspired aggregation is an effective strategy for low-communication distributed training, offering a path to scale training across geographically distributed or poorly connected clusters without sacrificing model quality.