Can Model Merging Improve Aggregation in DiLoCo?

Abstract: Model merging techniques, which aggregate independently finetuned models into one to combine their capabilities, have become a topic of significant interest in recent years, with a broad array of methods having been proposed to tackle this problem. Simultaneously, an emerging trend in distributed learning has been the use of methods such as local SGD and DiLoCo, which greatly reduce communication costs by periodically aggregating the independently trained local models. However, these communication-efficient methods have been shown to degrade in performance relative to the FLOP-matched data-parallel gold standard as the number of independent local models grows and as the number of local training steps before global communication is increased. In this work, we draw an explicit analogy between the pseudo-gradient aggregation step in local SGD/DiLoCo and task arithmetic-based model merging, establishing a straightforward way to utilize merging methods in the context of distributed optimization. We then evaluate multiple state-of-the-art model merging methods in this setting and identify one method in particular, Iso-C, as a promising approach for improving DiLoCo. We find that DiLoCo SGD with Iso-C aggregation outperforms not only simple pseudo-gradient averaging but even the momentum-based DiLoCo, despite lacking a momentum mechanism itself. Building on this finding, we propose IsoLoCo, which adapts Iso-C for distributed training by equipping it with Nesterov momentum. Our empirical evaluations on LLM pre-training across varying numbers of local workers show that IsoLoCo significantly outperforms DiLoCo, with the gap between them widening as the number of workers increases. This advantage remains present across model sizes and inner step counts, confirming that merging-inspired aggregation is an effective strategy for low-communication distributed training.

• The paper introduces IsoLoCo, a merging-inspired aggregation technique that uses isotropic spectral correction with momentum to reduce gradient interference.
• IsoLoCo achieves up to 5.8% lower validation loss than DiLoCo, improves scaling across worker counts and model sizes, and remains efficient with long synchronizations.
• The study establishes a formal link between task-arithmetic merging and pseudo-gradient averaging, providing a principled framework for future distributed LLM optimizations.

Merging-Inspired Aggregation in Distributed Low-Communication Training

Motivation and Background

The paper "Can Model Merging Improve Aggregation in DiLoCo?" (2607.03011) introduces a formal connection between task-arithmetic-based model merging and pseudo-gradient aggregation in DiLoCo-style distributed learning. DiLoCo is a federated optimization variant tailored for low-communication training of large-scale LLMs across heterogeneous clusters, where each worker performs multiple local updates before synchronization. As the number of participating workers or inner steps increases, DiLoCo suffers degradation relative to data-parallel baselines, underscoring the importance of robust global aggregation strategies.

Model merging, originally developed for aggregating independently fine-tuned model checkpoints, leverages task vectors—parameter displacements from shared initialization—to combine expertise across domains. Task interference in merging, caused by conflicting parameter updates, parallels the divergence of local replicas in DiLoCo during multi-step local optimization. Motivated by this analogy, the paper systematically evaluates whether merging-inspired aggregation can serve as a drop-in replacement for DiLoCo’s pseudo-gradient averaging, potentially mitigating its scalability bottlenecks.

Figure 1: DiLoCo as iterative model merging; local training produces task vectors (pseudo-gradients), with aggregation mirroring task-arithmetic merging, applied iteratively.

Model Merging Approaches as Aggregation Mechanisms

Traditional DiLoCo aggregates pseudo-gradients via arithmetic averaging. However, the merging literature offers refined alternatives that explicitly address parameter interference:

• Sparsity-Based Methods: TIES and DARE prune task vectors or reweight surviving parameters based on sign or magnitude, reducing direct parameter conflict. In distributed training, these methods emulate pseudo-gradient averaging at low sparsity but lose effectiveness as sparsification increases.
• Orthogonalization-Based Methods: TSV and Iso-C operate in the spectral domain. TSV compresses task matrices via SVD and decorrelates them via whitening, but incurs excessive computational cost for iterative distributed optimization. Iso-C, by flattening the singular-value spectrum of the aggregated task matrices to their mean, regularizes dominant directions in the gradient space and achieves strong improvements with modest computational overhead.

Empirical evaluation shows that Iso-C, despite lacking momentum, outperforms both vanilla DiLoCo SGD and momentum-based DiLoCo, with 2% lower validation loss under comparable settings. Iso-C’s computational requirements—a single SVD per weight matrix—make it practical for iterative use.

IsoLoCo: Isotropic Model Merging with Momentum

Expanding on Iso-C, the authors introduce IsoLoCo, which augments isotropic spectral correction with Nesterov momentum, further enhancing gradient aggregation stability and convergence. IsoLoCo replaces the singular values of each pseudo-gradient matrix with their mean (or a root-mean-square proxy in the fast implementation), regularizes the update, and then applies momentum.

Key empirical findings include:

• Worker Count Scaling: IsoLoCo consistently outperforms DiLoCo as the number of workers increases, with the improvement widening at scale. At $R=128$, IsoLoCo achieves a 5.8% lower validation loss than DiLoCo, reducing loss degradation relative to data-parallel baselines from 16.5% (DiLoCo) to 9.7%.
• Model Size: IsoLoCo maintains superiority over DiLoCo across model sizes up to 1B parameters, although the gap narrows as model size increases due to DiLoCo’s improved scaling with larger models.
• Synchronization Intervals: IsoLoCo demonstrates resilience to increased synchronization intervals (up to $H=240$), outperforming DiLoCo in all tested settings.
• Computational Efficiency: The Newton-Schulz iterative method, used for fast orthogonalization, accelerates IsoLoCo by up to two orders of magnitude with negligible loss increase (~0.45%), yielding a 15% end-to-end speedup.

Theoretical Implications and Practical Consequences

By reframing DiLoCo as iterative model merging, the paper highlights that the challenges faced by distributed training at scale (e.g., parameter divergence, communication bottlenecks) can be addressed using merging-inspired spectral correction techniques that are agnostic to task or data heterogeneity. The isotropic spectral flattening used by IsoLoCo reduces the dominance of high-energy directions in aggregated gradients, thus mitigating negative interference effects observed in both merging and federated settings.

Pairing IsoLoCo’s outer step with Muon as the inner optimizer (MuLoCo) further improves worker-count scaling, particularly at high $R$. IsoLoCo with an AdamW inner optimizer achieves equivalent or better results than MuLoCo for large worker counts, with reduced computational complexity due to less frequent orthogonalization.

From a practical perspective, merging-inspired aggregation enables FLOP-efficient, scalable training of LLMs across poorly connected or heterogeneous clusters, bypassing some limitations inherent to synchronous data-parallelism. The approach is robust to long synchronization intervals, varying worker counts, and multiple model scales.

Ablations and Design Analysis

Ablation studies confirm that:

• Nesterov momentum modestly outperforms standard momentum.
• Applying orthogonalization to pseudo-gradients (rather than the final update) yields superior results.
• The primary gains from isotropic merging arise from flattening high-energy singular values, thus reducing their undue influence on gradient aggregation.

Future Directions

The paper suggests several avenues for further exploration:

• Theoretical characterization of convergence guarantees for merging-inspired outer steps.
• Extension to architectures beyond dense transformers (e.g., MoE, hybrid networks) and tasks outside language modeling.
• Integration with communication compression and pipelining techniques.
• Evaluation in single-worker and highly heterogeneous cluster regimes.

Conclusion

"Can Model Merging Improve Aggregation in DiLoCo?" (2607.03011) establishes a formal equivalence between task-arithmetic merging and pseudo-gradient aggregation in distributed optimization, motivating the adoption of merging-inspired spectral correction mechanisms in federated LLM training. IsoLoCo, leveraging isotropic correction and momentum, significantly improves over standard DiLoCo, enabling robust scaling with both model size and worker count while delivering computational efficiency and stable convergence. The merging perspective provides a principled framework for designing future distributed optimization algorithms, with direct implications for scalable, low-communication LLM training in heterogeneous environments.