Add reEWC

.pre-commit-config.yaml

@@ -59,7 +59,7 @@ repos:
     hooks:
       - id: codespell
         stages: [pre-commit, commit-msg]
-        args: ["--ignore-words-list", "Commun"]
+        args: ["--ignore-words-list", "Commun,Mater"]
         exclude: |
           (?x)(
               ^example_inputs/data/|

CHANGELOG.md

@@ -3,6 +3,7 @@ All notable changes to this project will be documented in this file.
 
 ## [0.12.2.dev]
 ### Added
+- reEWC fine-tuning with forgetting prevention for single-modal models: optional experience replay (`rehearsal`, `load_memory_path`, `mem_batch_size`, `mem_ratio`) and an Elastic Weight Consolidation penalty from a precomputed Fisher matrix (`continue.fisher_information`, `continue.opt_params`, `continue.ewc_lambda`), plus a `cosineannealingwarmuplr` scheduler
 - Support OpenEquivariance
 - Per-atom stress (atomic virial) support in LAMMPS pair_e3gnn and ASE calculator
 - `compute_atomic_virial` option in `SevenNetCalculator`
@@ -14,6 +15,9 @@ All notable changes to this project will be documented in this file.
 - LAMMPS pair_e3gnn refactored to use pair-wise force (dE/dr) instead of position-based gradient.
 - Deploy no longer replaces force_output with ForceStressOutput; force/stress computed in LAMMPS C++ side.
 
+### Fixed
+- Load FlashTP-saved checkpoints (e.g. SevenNet-Nano) when FlashTP is unavailable by falling back to the e3nn backend, so they work for inference and fine-tuning without FlashTP installed. An explicit `enable_flash=True` still fails loud.
+
 ## [0.12.1]
 ### Fixed
 - FlashTP with LAMMPS parallel in torch

README.md

@@ -15,6 +15,7 @@ Full documentation, including **installation**, **usage**, and **pretrained mode
  - GPU-parallelized molecular dynamics with LAMMPS
  - CUDA-accelerated D3 (van der Waals) dispersion
  - Multi-fidelity training for combining multiple databases with different calculation settings
+ - Fine-tuning with forgetting prevention (experience replay + Elastic Weight Consolidation) for continual learning
  - [Tensor product accelerators](https://sevennet.readthedocs.io/en/latest/user_guide/accelerator.html)
 
 
@@ -71,3 +72,16 @@ If you utilize the pretrained model SevenNet-Omni or multi-task training strateg
 	year = {2026},
 }
 ```
+
+If you utilize the reEWC forgetting-aware fine-tuning strategy for continual learning of pretrained universal machine-learning interatomic potentials, please cite the following paper:
+```bib
+@article{kim_efficient_2026,
+	title = {An Efficient Forgetting-Aware Fine-Tuning Framework for Pretrained Universal Machine-Learning Interatomic Potentials},
+	volume = {12},
+	doi = {10.1038/s41524-025-01895-w},
+	number = {26},
+	journal = {npj Comput. Mater.},
+	author = {Kim, Jisu and Lee, Jiho and Oh, Sangmin and Park, Yutack and Hwang, Seungwoo and Han, Seungwu and Kang, Sungwoo and Kang, Youngho},
+	year = {2026},
+}
+```

docs/source/index.rst

@@ -16,6 +16,7 @@ SevenNet (Scalable EquiVariance-Enabled Neural Network) is a graph neural networ
 * GPU-parallelized molecular dynamics with LAMMPS
 * CUDA-accelerated D3 (van der Waals) dispersion
 * Multi-fidelity training for combining multiple databases with different calculation settings
+* Fine-tuning with forgetting prevention (Experience replay + Elastic Weight Consolidation) for continual learning
 
 
 Installation

docs/source/user_guide/index.rst

@@ -26,5 +26,6 @@ SevenNet offers various pretrained models, MD engines (ASE, LAMMPS), and user in
    ase_calculator
    torchsim
    cli
+   reewc
    d3
    note_book

docs/source/user_guide/reewc.md

@@ -0,0 +1,63 @@
+# Forgetting-prevented (Continual-learning) fine-tuning (reEWC)
+
+Fine-tuning a pretrained model on a target system improves accuracy there, but the
+model can lose accuracy on the original training domain (catastrophic forgetting).
+reEWC mitigates this with two complementary mechanisms that can be used together or
+separately:
+
+- **Experience replay (rehearsal)** -- replay an old-task "memory" set each training
+  step so the model keeps fitting it while learning the target data.
+- **Elastic Weight Consolidation (EWC)** -- add a penalty
+  `lambda/2 * sum_i F_i (theta_i - theta*_i)^2` that anchors parameters to their
+  pre-fine-tuning values `theta*`, weighted by a precomputed Fisher matrix `F`.
+
+reEWC is for **single-modal** models (e.g. SevenNet-0, SevenNet-Nano). Multi-fidelity
+(modal) models are not supported yet.
+
+## Getting started
+
+A ready-to-edit input with both mechanisms is available as a preset:
+
+```bash
+sevenn preset reewc > input.yaml
+```
+
+The preset documents every key inline. Replay lives in the `data:` block
+(`rehearsal`, `load_memory_path`, `mem_batch_size`, `mem_ratio`) and EWC in the
+`train.continue:` block (`fisher_information`, `opt_params`, `ewc_lambda`). Every
+reEWC key is optional; when none are set, training is unchanged. Remove the replay
+block or the EWC keys to run only one mechanism. Run training as usual:
+
+```bash
+sevenn train input.yaml -s
+```
+
+## Fisher information and reference parameters
+
+`fisher_information` and `opt_params` are **precomputed and consumed** -- SevenNet
+does not estimate the Fisher matrix. Both are `torch.save`d dictionaries keyed by
+parameter name; `opt_params` is the parameter set of the checkpoint before
+fine-tuning. They must satisfy:
+
+- `fisher_information` and `opt_params` cover the **same parameter names** with the
+  **same shapes** (they are a matched pair).
+- Names that overlap with the model's trainable parameters must have **matching
+  shapes**; a mismatch is an error (usually an incompatible checkpoint or SevenNet
+  version).
+- At least one name must overlap with the model; no overlap is an error.
+- A trainable parameter without a Fisher entry is **left unconstrained** and a
+  warning is emitted, so partial-coverage Fisher matrices are allowed but visible.
+
+`ewc_lambda` must be `> 0`, and EWC requires both `fisher_information` and
+`opt_params` to be set.
+
+## Notes
+
+- Replay supports `dataset_type: 'graph'` (the default) only.
+- reEWC does not support distributed (DDP) training.
+- `load_memory_path` is reserved for replay: setting it without `rehearsal: True`
+  raises an error.
+- When replay is enabled, the memory set is evaluated each epoch and logged as a
+  `memoryset` column group in `lc.csv`, alongside `trainset` and `validset`.
+- A `cosineannealingwarmuplr` scheduler (cosine annealing with warm-up restarts,
+  used for the reEWC paper work) is also available for fine-tuning.

example_inputs/training/input_full.yaml

@@ -73,6 +73,11 @@ train:
         #checkpoint: 'checkpoint_best.pth'         # Checkpoint of pre-trained model or a model want to continue training.
         #reset_optimizer: False                    # Set True for fine-tuning
         #reset_scheduler: False                    # Set True for fine-tuning
+        # reEWC (single-modal models only): add an Elastic Weight Consolidation penalty from a
+        # precomputed Fisher matrix and reference parameters to preserve prior-task accuracy.
+        #fisher_information: './fisher.pt'         # dict {param_name: tensor} of precomputed Fisher information
+        #opt_params: './opt_params.pt'             # dict {param_name: tensor} of reference (pre-finetune) parameters
+        #ewc_lambda: 100000                        # EWC penalty weight (must be > 0 when fisher/opt are given)
 
 data:
     batch_size: 4                                 # Per GPU batch size.
@@ -91,3 +96,10 @@ data:
     load_trainset_path: ['./structure_list']  # Example of using ase as data_format, support multiple files and expansion(*)
     #load_validset_path: ['./valid.extxyz']
     #load_testset_path:  ['./sevenn_data/mydata.pt']  # Graph can be preprocessed using `sevenn_graph_build` and accessible like this
+
+    # reEWC rehearsal (experience replay, single-modal models only): replay an old-task memory set
+    # each training step to mitigate catastrophic forgetting while fine-tuning on the target data.
+    #rehearsal: False                              # Set True to enable replay
+    #load_memory_path: ['./memory.extxyz']         # memory (old-task) set; this key is reserved for rehearsal
+    #mem_batch_size: 8                             # batch size for the replayed memory set
+    #mem_ratio: 1                                  # fraction (0, 1] of the memory set to use

sevenn/_const.py

@@ -212,6 +212,9 @@ def model_defaults(config):
     KEY.USE_MODAL_WISE_SCALE: False,
     KEY.SHIFT: 'per_atom_energy_mean',
     KEY.SCALE: 'force_rms',
+    KEY.REHEARSAL: False,
+    KEY.MEM_BATCH_SIZE: 0,
+    KEY.MEM_RATIO: 1,
     # KEY.DATA_SHUFFLE: True,
     # KEY.DATA_WEIGHT: False,
     # KEY.DATA_MODALITY: False,
@@ -233,6 +236,11 @@ def model_defaults(config):
     KEY.SCALE: lambda x: type(x) in [float, list] or x in IMPLEMENTED_SCALE,
     KEY.USE_MODAL_WISE_SHIFT: bool,
     KEY.USE_MODAL_WISE_SCALE: bool,
+    KEY.REHEARSAL: lambda x: isinstance(x, bool),
+    KEY.MEM_BATCH_SIZE: lambda x: isinstance(x, int) and not isinstance(x, bool),
+    KEY.MEM_RATIO: lambda x: isinstance(x, (int, float))
+    and not isinstance(x, bool)
+    and 0 < x <= 1,
     # KEY.DATA_SHUFFLE: bool,
     KEY.COMPUTE_STATISTICS: bool,
     # KEY.DATA_WEIGHT: bool,

sevenn/_keys.py

@@ -103,6 +103,10 @@
 LOAD_TRAINSET = 'load_trainset_path'
 LOAD_VALIDSET = 'load_validset_path'
 LOAD_TESTSET = 'load_testset_path'
+LOAD_MEMORY_PATH = 'load_memory_path'  # reEWC rehearsal memory set
+REHEARSAL = 'rehearsal'
+MEM_BATCH_SIZE = 'mem_batch_size'
+MEM_RATIO = 'mem_ratio'
 FORMAT_OUTPUTS = 'format_outputs_for_ase'
 COMPUTE_STATISTICS = 'compute_statistics'
 DATASET_TYPE = 'dataset_type'
@@ -135,6 +139,9 @@
 USE_STATISTIC_VALUES_FOR_CP_MODAL_ONLY = (
     'use_statistic_values_for_cp_modal_only'
 )
+OPT_PARAMS = 'opt_params'  # reEWC: reference (optimal) params pickle
+FISHER = 'fisher_information'  # reEWC: precomputed Fisher information pickle
+EWC_LAMBDA = 'ewc_lambda'  # reEWC: EWC penalty weight
 
 CSV_LOG = 'csv_log'
 

sevenn/checkpoint.py

@@ -327,11 +327,29 @@ def build_model(
         enable_cueq = cp_using_cueq if enable_cueq is None else enable_cueq
 
         cp_using_flash = self.config.get(KEY.USE_FLASH_TP, False)
+        flash_requested_explicitly = enable_flash is True
         enable_flash = cp_using_flash if enable_flash is None else enable_flash
 
         cp_using_oeq = self.config.get(KEY.USE_OEQ, False)
         enable_oeq = cp_using_oeq if enable_oeq is None else enable_oeq
 
+        # FlashTP-saved checkpoints must still load where FlashTP is unavailable.
+        if enable_flash:
+            from sevenn.nn.flash_helper import is_flash_available
+
+            if not is_flash_available():
+                if flash_requested_explicitly or _flash_lammps:
+                    raise ValueError(
+                        'FlashTP was requested but is not available (package '
+                        'not installed or no GPU available).'
+                    )
+                warnings.warn(
+                    'FlashTP is unavailable; loading the checkpoint with the '
+                    'e3nn backend instead.',
+                    UserWarning,
+                )
+                enable_flash = False
+
         if sum([enable_cueq, enable_flash, enable_oeq]) > 1:
             raise ValueError('Only one TP accelerator can be enabled.')
 

sevenn/main/sevenn_preset.py

@@ -27,6 +27,7 @@ def add_args(parser):
             'base',
             'multi_modal',
             'mf_ompa_fine_tune',
+            'reewc',
         ],
         help=preset_help
     )

sevenn/presets/reewc.yaml

@@ -0,0 +1,90 @@
+# Example input.yaml for forgetting-prevented fine-tuning (reEWC).
+# Replay and EWC are independent; keep one block, both, or neither.
+# reEWC is for single-modal models (e.g. SevenNet-0, SevenNet-Nano).
+
+model:  # keep consistent with the checkpoint being fine-tuned
+    chemical_species: 'Auto'
+    cutoff: 5.0
+    channel: 128
+    is_parity: False
+    lmax: 2
+    num_convolution_layer: 5
+    irreps_manual:
+        - "128x0e"
+        - "128x0e+64x1e+32x2e"
+        - "128x0e+64x1e+32x2e"
+        - "128x0e+64x1e+32x2e"
+        - "128x0e+64x1e+32x2e"
+        - "128x0e"
+
+    weight_nn_hidden_neurons: [64, 64]
+    radial_basis:
+        radial_basis_name: 'bessel'
+        bessel_basis_num: 8
+    cutoff_function:
+        cutoff_function_name: 'XPLOR'
+        cutoff_on: 4.5
+    self_connection_type: 'linear'
+
+    train_shift_scale: False
+    train_denominator: False
+
+train:
+    random_seed: 1
+    is_train_stress: True
+    epoch: 100
+
+    loss: 'Huber'
+    loss_param:
+        delta: 0.01
+
+    optimizer: 'adam'
+    optim_param:
+        lr: 0.004
+    # cosineannealingwarmuplr (cosine annealing with warm-up restarts) was used
+    # for the reEWC work; exponentiallr also works.
+    scheduler: 'exponentiallr'
+    scheduler_param:
+        gamma: 0.99
+
+    force_loss_weight: 1.0
+    stress_loss_weight: 0.01
+
+    per_epoch: 10
+
+    error_record:
+        - ['Energy', 'RMSE']
+        - ['Force', 'RMSE']
+        - ['Stress', 'RMSE']
+        - ['TotalLoss', 'None']
+
+    continue:
+        reset_optimizer: True
+        reset_scheduler: True
+        reset_epoch: True
+        checkpoint: 'SevenNet-0_11July2024'
+
+        # EWC: anchor parameters to their pre-fine-tuning values via a
+        # precomputed Fisher matrix. fisher_information and opt_params are
+        # torch.save'd dicts {param_name: tensor} matching the model's
+        # trainable parameters; opt_params is the checkpoint before fine-tuning.
+        # All three keys are required together; remove them to disable EWC.
+        fisher_information: './fisher.pt'
+        opt_params: './opt_params.pt'
+        ewc_lambda: 100000        # EWC penalty weight (> 0)
+
+data:
+    batch_size: 4
+    data_divide_ratio: 0.1
+    data_format_args:
+        index: ':'
+
+    load_trainset_path: ['./target_train.extxyz']
+    load_validset_path: ['./valid.extxyz']
+
+    # Replay (experience replay): replay an old-task memory set each step so the
+    # model keeps fitting it. Remove this block to disable replay.
+    rehearsal: True
+    load_memory_path: ['./memory.extxyz']  # requires rehearsal: True
+    mem_batch_size: 8
+    mem_ratio: 1                            # fraction (0, 1] of the memory set

sevenn/scripts/processing_epoch.py

@@ -41,6 +41,12 @@ def processing_epoch_v2(
         config, trainer.loss_functions
     )
     recorders = {k: deepcopy(recorder) for k in loaders}
+    # reEWC: log the replayed memory set as a separate 'memoryset' column group.
+    memory_recorder = (
+        deepcopy(recorder)
+        if getattr(trainer, 'memory_loader', None) is not None
+        else None
+    )
 
     best_val = float('inf')
     best_key = None
@@ -58,6 +64,8 @@ def processing_epoch_v2(
         head = ['epoch', 'lr']
         for k, rec in recorders.items():
             head.extend(list(rec.get_dct(prefix=k)))
+        if memory_recorder is not None:
+            head.extend(list(memory_recorder.get_dct(prefix='memoryset')))
         with open(csv_path, 'w') as f:
             f.write(','.join(head) + '\n')
 
@@ -88,9 +96,17 @@ def processing_epoch_v2(
                 loader.sampler.set_epoch(epoch)
 
             rec = recorders[k]
-            trainer.run_one_epoch(loader, is_train, rec)
+            trainer.run_one_epoch(
+                loader,
+                is_train,
+                rec,
+                memory_error_recorder=memory_recorder if is_train else None,
+            )
             csv_dct.update(rec.get_dct(prefix=k))
             errors[k] = rec.epoch_forward()
+        if memory_recorder is not None:
+            csv_dct.update(memory_recorder.get_dct(prefix='memoryset'))
+            errors['memoryset'] = memory_recorder.epoch_forward()
         log.write_full_table(list(errors.values()), list(errors))
         trainer.scheduler_step(best_val)