GFFMERGE: Efficient Merging of Graph Neural Force Fields and Beyond

Abstract: Graph Neural Networks (GNNs) have revolutionized Neural Force Fields for atomistic simulations, achieving near-quantum accuracy at reduced cost, yet adapting these models to new chemical systems requires expensive retraining of foundation models. Inspired by model merging in vision and language processing, we introduce GFFMERGE, the first principled framework for closed-form model merging in GNNs. We exploit the linear structure of message-passing layers and formulate merging as a convex embedding-alignment problem with an analytical solution. Through the first systematic benchmarking of model merging for GNNs, we show that existing methods designed for vision and language catastrophically fail on force field regression, while GFFMERGE recovers performance approaching gold standard joint training. Across molecular (MD17, MD22), solid-state (LiPS20), and large-scale graph benchmarks, GFFMERGE and GNNMERGE (its generic GNN counterpart) achieve 5-27$\times$ speedups while enabling modular composition of specialized models. Remarkably, our closed-form solution alone outperforms all baseline methods before fine-tuning and provides superior initialization for faster, data-efficient convergence.

• The paper introduces a closed-form analytical merging technique that aligns linear message-passing blocks to replicate joint training performance.
• The paper demonstrates that GFFMERGE achieves force and energy MAEs within 5–15% of fully retrained models while reducing computational costs by 5×–27×.
• The paper validates the framework across diverse datasets including MD17, MD22, and LiPS20, proving stability and scalability for GNN-based force fields.

Efficient Model Merging for Graph Neural Force Fields: A Detailed Analysis of GFFMERGE

Introduction and Motivation

Graph Neural Network (GNN)-based force fields exhibit near-quantum accuracy for modeling interatomic potentials in atomistic simulations, yet retraining or fine-tuning such models for new domains remains computationally intensive, intractable for evolving chemical data regimes, and often redundant when base models are already available. In vision and language domains, model merging—a post-hoc combination of fine-tuned models—has emerged as a solution to similar scalability issues. However, existing merging methods are incompatible with continuous force field regression, where requirements for smoothness, numerical stability, and tight coupling of predictions present unique challenges absent in discrete prediction tasks.

The paper "GFFMERGE: Efficient Merging of Graph Neural Force Fields and Beyond" (2606.03232) introduces the first principled, closed-form framework for merging GNN-based force fields and more generally, GNN models. GFFMERGE addresses the pressing need for modular, scalable composition and transfer of GNN expertise across chemical domains, bypassing the computational bottlenecks of conventional joint retraining.

Theoretical Framework and Methodology

The core insight exploited by GFFMERGE is the near-universal presence of linear parameter blocks within message-passing GNNs. The authors reconceptualize the merging problem as an embedding alignment task in the space of node representations. For a set of specialist models $\{\Theta_1, \ldots, \Theta_n\}$, each fine-tuned on a distinct domain $D_i$, the goal is to construct a merged model $\Theta_M$ whose predictions approximate those of the source domains on their respective support. Formally, this is cast as minimizing the predictive discrepancy over unioned domains.

GFFMERGE proceeds in two stages:

1. Closed-Form Linear Merging: By decoupling the network at the level of linear blocks (e.g., weights in MLPs or GNN message-passing layers), the authors derive a convex, closed-form solution for optimal alignment of node embeddings at intermediate layers. This avoids full-network backpropagation across merged models, enabling parallel, analytical merging without data access or gradient-based optimization.
2. Targeted Fine-Tuning: To correct residual nonlinear misalignments, only the top layers (e.g., energy/force heads, final interaction blocks) are refined using limited data and a small number of epochs. This refinement is intentionally lightweight, incurring minimal additional cost.

This approach stands in contrast to the heuristic parameter-averaging, Fisher-matrix blending, or permutation-based schemes that proved inadequate in GNN regression, as detailed in the empirical section.

Empirical Evaluation

Predictive Fidelity and Robustness

The paper provides the first systematic benchmark suite for model merging in GNN force fields. Evaluation domains span:

• Molecular systems (MD17, MD22) capturing local and long-range interactions.
• Solid-state crystal electrolytes (LiPS20) representing challenging bulk dynamics.

Accuracy is assessed on force MAE and energy MAE. Across all benchmarks, GFFMERGE delivers performance within 5–15% of gold-standard joint fine-tuning (Joint FT), exhibiting negligible degradation even as the number of merged domains increases (e.g., 5-task mixes). Competing baselines (weight averaging, Fisher, EMR, TIES) are consistently outperformed, typically suffering 2–5× higher errors, and in some cases completely destabilizing force predictions in complex regimes.

Key empirical findings:

• GFFMERGE considerably narrows the performance gap to joint training. For example, on MD17 and MD22, force MAE and energy MAE errors are within 1.1–2× of gold standard, while baselines fail catastrophically in multi-domain settings (see Tables 1–2 in the paper).
• Physical stability metrics (MD rollouts, cumulative energy/force violations) confirm that GFFMERGE produces models with comparable stability to jointly trained models, outperforming baselines on both structural and dynamical diagnostics.

Computational Efficiency

A major contribution of GFFMERGE is the elimination of redundant retraining costs. The hybrid approach—closed-form for 90%+ of parameters, fine-tuning for a thin nonlinear head—yields speedups of 5×–27× compared to full joint fine-tuning, as measured by wall-clock time on diverse architectures (M3GNet, Orb) and datasets. This translates to orders-of-magnitude energy and resource savings for practical deployment.

The merged model is not only faster to train but also shows much improved data and epoch efficiency in downstream adaptation. Rapid convergence is enabled by favorable weight-space initialization provided by closed-form merging, requiring as few as 250 samples or <20 epochs for downstream specialization.

Generalization Beyond Force Fields

GFFMERGE’s architectural generality is demonstrated by its extension to generic GNN architectures and tasks:

• GCN, GraphSAGE, GAT, NodeFormer, and graph transformers: The embedding alignment-based merging applies to node classification and link prediction without modification.
• Heterogeneous architectures and dimensions: Merging is possible even across dissimilar GNN variants or layers with mismatched dimensions via joint block alignment and adaptive padding.
• Scaling to extreme dataset sizes: On large OGBN datasets (MAG240M, Papers100M), GNNMerge reduces memory requirements by $>300\times$ and training time by $>1000\times$ relative to data-based joint retraining.

Implications and Future Directions

Practical Impact

GFFMERGE transforms modular learning and specialization in graph-based atomistic modeling, overcoming practical barriers to rapid iteration and deployment:

• Modular reuse: Domain-specific GNNs may be trivially assembled as blocks to create unified force fields for novel compositions, without expensive retraining or original data.
• Continual and federated learning: New domains can be incorporated swiftly, making GNN-based scientific computing feasible in settings with evolving or distributed data.
• Efficient resource utilization: Substantial computational and energy savings are realized, making high-fidelity simulation more accessible and environmentally sustainable.

Theoretical Significance

The closed-form analytical merging guarantees stable and interpretable solutions exploiting the linear structure of message-passing GNNs—an approach that directly addresses instability issues seen in older merging methods. This framework clarifies conditions under which model merging is mathematically sound and physically meaningful in continuous domains, providing a foundation for future theoretical advances.

Future Developments

Areas for further research include:

• Dynamic/incremental merging for lifelong and online learning settings, allowing seamless integration of new chemistry or materials knowledge.
• Robustness to distribution shift: While GFFMERGE empirically handles OOD and rMD17 (high-accuracy CCSD(T)) data well, formal guarantees or adaptive correction mechanisms for extreme domain shifts remain open.
• Extension to advanced equivariant or higher-order GNNs: The alignment framework should generalize further as new invariant or equivariant GNN architectures for atomistic systems are developed.

Conclusion

GFFMERGE (2606.03232) provides a theoretically grounded, computationally efficient, and empirically validated framework for merging GNN-based force fields and generic GNNs. Leveraging closed-form convex optimization over linear message-passing blocks and lightweight targeted adaptation, GFFMERGE matches joint training accuracy within a small margin while reducing cost by an order of magnitude. Its modular paradigm unlocks scalable model composition for applications in computational chemistry, materials science, and large-scale graph learning—laying the groundwork for future modular AI systems in scientific domains.