Magnitude-Modulated Equivariant Adapter for Parameter-Efficient Fine-Tuning of Equivariant Graph Neural Networks

AI summary

Overview

Research area: Machine learning for atomistic simulation, specifically parameter-efficient fine-tuning (PEFT) of equivariant graph neural networks (GNNs) used as machine-learning interatomic potentials.

Technical level: Intermediate. The paper assumes familiarity with equivariance, irreducible representations, and spherical-harmonic-based tensor features, though the central idea is a simple per-channel scaling.

Scope: The paper proposes and evaluates MMEA, a symmetry-preserving adapter that fine-tunes pretrained equivariant GNNs by modulating feature magnitudes with scalar gates, and reports results on the rMD17, 3BPA, and AcAc molecular benchmarks.

What This Paper Is About

Pretrained equivariant graph neural networks based on spherical harmonics can approximate expensive ab-initio quantum chemistry calculations accurately, but adapting them to a new molecule or chemical environment still requires fine-tuning on task-specific data. Ordinary PEFT methods such as Adapters and LoRA break the rotational symmetry that makes these models work, and the existing equivariant alternative, ELoRA, still leaves enough freedom inside each tensor order to disturb the pretrained feature distribution. The goal of this paper is to fine-tune such models with even fewer trainable parameters while strictly preserving equivariance and improving energy and force predictions.

Key Contributions

1. A new equivariant PEFT method. The authors introduce MMEA (Magnitude-Modulated Equivariant Adapter), which applies dynamic magnitude modulation on a per-channel and per-multiplicity basis to features in equivariant GNNs based on spherical harmonics.

2. A formal equivariance argument. The paper proves that MMEA's adapter commutes with the SO(3) group action, so the adapted features remain strictly equivariant. Because all gains are computed only from the scalar (order ℓ = 0) component, which is group-invariant, the gains themselves are unchanged under rotation.

3. State-of-the-art fine-tuning results with fewer parameters. Evaluations on multiple molecular benchmarks (rMD17, 3BPA, AcAc) report that MMEA surpasses previous methods, including ELoRA and full-parameter fine-tuning, while training fewer parameters.

4. An implementation in e3nn. MMEA is integrated into the widely used e3nn framework, and code is released at https://github.com/CLaSLoVe/MMEA.

Main Findings

• rMD17 (few-shot, 50 training samples): MMEA outperformed ELoRA with an additional average improvement of approximately 6.6% on energy and 8.7% on forces. Overall, MMEA compressed average energy and force MAEs by 6–8%, with 10–25% improvements on high-error molecules. On Aspirin, energy MAE dropped from 9.7 (Full) to 8.0 (ELoRA, ↓18%) to 7.3 (MMEA, ↓25%), and force MAE from 23.9 to 18.3 (↓23%) to 16.4 (↓31%).

• 3BPA (trained at 300 K, tested at higher temperatures): MMEA gave a modest improvement over ELoRA at 300 K and consistently lower errors at 600 K and 1200 K. At 1200 K, energy RMSE was 29.8 (Scratch), 20.3 (Full), 17.6 (ELoRA), 17.1 (MMEA); force RMSE was 62.0 (Scratch), 48.7 (Full), 42.0 (ELoRA), 39.7 (MMEA). On the dihedral slice test, energy RMSE was 7.8 (Scratch), 7.3 (Full), 5.9 (ELoRA), 5.6 (MMEA), and force RMSE 16.5, 12.3, 11.4, 10.6 respectively.

• AcAc (trained at 300 K, tested at 300 K and 600 K): At 300 K, MMEA achieved the lowest energy and force RMSE. At 600 K, energy RMSE was 4.6 (Scratch), 5.8 (Full), 3.9 (ELoRA), 3.6 (MMEA, r = 16), 3.2 (MMEA, r = 32); force RMSE was 22.4 (Scratch), 16.4 (Full), 13.6 (ELoRA), 13.2 (MMEA, r = 16), 13.0 (MMEA, r = 32).

• Parameter efficiency: On rMD17, full fine-tuning trains 751,896 parameters (100.0% of full FT). ELoRA at r = 16 trains 175,880 (23.4%). MMEA at r = 16 trains 151,354 (20.1%, about 85% of ELoRA's budget), and MMEA at r = 32 trains 201,258 (26.7%).

• Ablation on rMD17-Aspirin (MAE, energy / forces): Full 9.7 / 23.9; MMEA 7.3 / 16.4; without nonlinear activation 7.6 / 16.4; without input-head reuse 9.2 / 16.7; without scalar modulation 12.9 / 30.5; without high-order modulation 8.3 / 16.6; shared high-order modulation 7.6 / 16.6; Readout-only tuning 23.8 / 36.8 (updating 0.3% of parameters); a conventional Adapter 11.0 / 26.3. Both Readout and Adapter underperformed Full, ELoRA, and MMEA.

• Rank sensitivity: Increasing the bottleneck rank initially improves accuracy on rMD17-Aspirin but eventually degrades it, indicating a trade-off between underfitting at low rank and overfitting at high rank. The optimal rank is described as a task-specific hyperparameter.

• Limitations observed: When the target system deviates substantially from the pretraining distribution, MMEA underperformed ELoRA. Using MACE-MP (pretrained on inorganic compounds) to predict organic compounds, full fine-tuning was best — Aspirin energy MAE 11.4 (Full) vs 12.5 (ELoRA) vs 13.3 (MMEA); Aspirin forces 29.4 vs 31.1 vs 31.9; Uracil energy 3.4 vs 3.5 vs 3.6; Uracil forces 19.7 vs 20.7 vs 21.4. Separately, unlike ELoRA, MMEA weights cannot be merged back into the backbone, so inference incurs slight additional computational overhead.

• Not reported in the available content: The training efficiency experiment in Appendix B.1 is truncated, so its results, and the claimed computational overhead and deviation analyses in the appendix, are not reported in the text provided.

Methodology in Plain English

Equivariant GNNs represent each atom's features as a collection of tensors organized by order ℓ; each order comes with a set of parallel channels called multiplicities. Standard adapters mix these channels through a bottleneck, which blends different irreducible representations and destroys the symmetry property. MMEA instead leaves the channels untouched and only rescales them.

Concretely, the method takes just the scalar part of the node feature (order ℓ = 0) and passes it through a small two-layer MLP: a bottleneck projection of dimension r using a SiLU activation, followed by an expansion that produces one scalar gain per multiplicity copy for every tensor order. Because the scalar part of a feature does not change under rotation, these gains are automatically rotation-invariant. Each higher-order channel is then multiplied by its gain through a function that is either 1 + x (residual scaling) or e^x (positive scaling). The scalar order itself receives an additive shift rather than a multiplication. The authors then prove that applying the adapter before or after a rotation gives the same result, so equivariance is preserved exactly.

The experiments use MACE-OFF as the frozen pretrained backbone (with MACE hyperparameters: correlation order 3, cutoff radius 5.0, max order for input 3, max order for product layer 1, 128 hidden channels, 2 interaction layers, 8 radial basis functions). Only equivariant linear layers are adapted with MMEA; tensor product layers and some scalar networks keep the original ELoRA configuration. Experiments use rank r = 16 unless stated, three runs each, PyTorch on Ubuntu 22.04.5 LTS with a single NVIDIA A100 GPU. For rMD17, 50 configurations are randomly sampled from the official training set so all methods see identical data; 3BPA and AcAc results for full fine-tuning and ELoRA are taken from prior work.

Why This Matters

Impact on research: The work challenges the assumption that fine-tuning equivariant models requires mixing channels within a tensor order. It shows that, in many practical settings, scaling channel magnitudes is enough to adapt an equivariant model to a new chemical environment, and it offers a symmetry-preserving adapter that is both lighter than ELoRA and (on the reported benchmarks) more accurate. It also frames an explicit decision rule: train from scratch or fine-tune fully when the distribution gap is large, and use ELoRA or MMEA when the gap is moderate or small.

Real-world applications:

• Molecular dynamics simulations that need quantum-level accuracy at a fraction of DFT cost, where a pretrained potential must be adapted to a new molecule or temperature regime.
• High-throughput materials and drug-discovery screening, where thousands of candidate configurations make the cubic scaling of DFT prohibitive.
• Reactive chemistry studies, such as the AcAc hydrogen-transfer and double-bond rotation modes, which test extrapolation along internal coordinates.
• Deploying foundation interatomic potentials on small task-specific datasets when high-quality quantum data is expensive to obtain.

Industry relevance: Smaller trainable parameter budgets (20.1% of full fine-tuning on rMD17) reduce memory and storage costs for task-specific model variants, and the e3nn integration lowers the barrier to adoption for groups already using that ecosystem. The inability to merge adapter weights into the backbone is a practical caveat for latency-sensitive deployment.

Future Directions

1. Efficient equivariant methods for Fully Connected Tensor Product layers. Because tensor products involve multiple inputs, how to gate them effectively remains open; the authors report that gating based on the output did not show significant benefits.
2. A generalizable unified fine-tuning method. A promising direction is combining the strengths of full fine-tuning, ELoRA, and MMEA into one method that performs well across data distributions that differ from the pretraining distribution.
3. Improving fine-tuning efficiency. The authors state that fine-tuning convergence currently remains relatively slow and call for schemes that converge faster.
4. Handling and diagnosing the distribution-gap regime. The paper shows MMEA degrades when the target deviates strongly from pretraining and points to training from scratch or full fine-tuning in that case, leaving the question of how to detect and adapt across that boundary.

Target Audience

Researchers and practitioners in machine learning for atomistic simulation, particularly those working with equivariant GNNs and interatomic potentials such as the MACE family. It is also relevant to PEFT researchers interested in symmetry-preserving adaptation beyond NLP and computer vision, and to engineers deploying pretrained molecular foundation models who need lightweight, task-specific fine-tuning within the e3nn ecosystem.

Authors’ abstract

Pretrained equivariant graph neural networks based on spherical harmonics offer efficient and accurate alternatives to computationally expensive ab-initio methods, yet adapting them to new tasks and chemical environments still requires fine-tuning. Conventional parameter-efficient fine-tuning (PEFT) techniques, such as Adapters and LoRA, typically break symmetry, making them incompatible with those equivariant architectures. ELoRA, recently proposed, is the first equivariant PEFT method. It achieves improved parameter efficiency and performance on many benchmarks. However, the relatively high degrees of freedom it retains within each tensor order can still perturb pretrained feature distributions and ultimately degrade performance. To address this, we present Magnitude-Modulated Equivariant Adapter (MMEA), a novel equivariant fine-tuning method which employs lightweight scalar gating to modulate feature magnitudes on a per-order and per-multiplicity basis. We demonstrate that MMEA preserves strict equivariance and, across multiple benchmarks, consistently improves energy and force predictions to state-of-the-art levels while training fewer parameters than competing approaches. These results suggest that, in many practical scenarios, modulating channel magnitudes is sufficient to adapt equivariant models to new chemical environments without breaking symmetry, pointing toward a new paradigm for equivariant PEFT design.