From 9f52b70e55a7ccd800cc379a3e8e92560670ed90 Mon Sep 17 00:00:00 2001 From: Ana Gainaru Date: Tue, 21 Jul 2026 09:55:51 -0400 Subject: [PATCH 1/4] Bug fix in JVP that was calling fwd_bwd twice if there is no historic data --- src/apeiron/training/updater/jvp_reg.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/apeiron/training/updater/jvp_reg.py b/src/apeiron/training/updater/jvp_reg.py index e4b615f..7ae8291 100644 --- a/src/apeiron/training/updater/jvp_reg.py +++ b/src/apeiron/training/updater/jvp_reg.py @@ -66,8 +66,11 @@ def fwd_bwd( """ # Current gradients ### + # Deliberately current-only: base.fwd_bwd() mixes in hist_batch when + # given one, which would corrupt the tangent direction below and + # double-count hist_batch against the explicit backward pass further down. loss_cur = super().fwd_bwd( - batch, hist_batch + batch ) # fill the model gradients with the default grad info ### JVP gradients ### @@ -78,7 +81,7 @@ def fwd_bwd( ### History gradients #### if hist_batch is None: - return super().fwd_bwd(batch) + return loss_cur x_mem, y_mem = hist_batch outputs_mem = self.model(x_mem) From 91682020387e74fc4cac7dc645fcb58270bc5224 Mon Sep 17 00:00:00 2001 From: Ana Gainaru Date: Tue, 21 Jul 2026 09:56:30 -0400 Subject: [PATCH 2/4] Changes to the base updater to mix historic and new data in a balanced way --- CLAUDE.md | 10 ++++++++++ src/apeiron/config/configuration.py | 5 +++++ src/apeiron/training/updater/base.py | 17 ++++++++++++++++- 3 files changed, 31 insertions(+), 1 deletion(-) diff --git a/CLAUDE.md b/CLAUDE.md index e9bcb1c..cf88d85 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -75,6 +75,16 @@ Note: `EnsembleDetector` is recognized by the loader but raises `NotImplementedE | `kfac_online` | KFAC approximation | kfac_lambda, kfac_ema_decay | | `none` | No-op (skip CL) | (none) | +Replay of historical data is handled in `BaseUpdater.fwd_bwd()` and gated by +`[continual_learning] mix_historic_data` (default `false`). When enabled and the +harness supplies historical dataloaders, half of each stream is combined into a +single forward/backward pass, so the sample count per step stays at +`train.batch_size` whether or not mixing is on (an undersized historical batch is +topped up from the current one). `base`, `ewc_online`, and `kfac_online` inherit this. `jvp_reg` ignores the flag: it +overrides `fwd_bwd()` and mixes the two streams itself (it needs a current-only +gradient as the JVP tangent direction), so it calls `super().fwd_bwd(batch)` without +the historical batch. `none` skips training entirely. + ### Coding Conventions - Python 3.13+, type hints everywhere - Formatting: ruff format, ruff check, mypy diff --git a/src/apeiron/config/configuration.py b/src/apeiron/config/configuration.py index e79a566..160c2c8 100644 --- a/src/apeiron/config/configuration.py +++ b/src/apeiron/config/configuration.py @@ -112,6 +112,11 @@ class DataCfg: class ContinualLearningCfg: update_mode: str = "base" + # Replay: concatenate the historical batch onto the current one in + # BaseUpdater.fwd_bwd(). Has no effect when the harness provides no + # historical dataloaders, nor on jvp_reg, which mixes the two streams itself. + mix_historic_data: bool = False + # For JVP regularization jvp_lambda: float = 0.001 jvp_deltax_norm: float = 1 diff --git a/src/apeiron/training/updater/base.py b/src/apeiron/training/updater/base.py index 206476d..8a2e5aa 100644 --- a/src/apeiron/training/updater/base.py +++ b/src/apeiron/training/updater/base.py @@ -26,14 +26,29 @@ def __init__(self, cfg: Config, modelHarness: BaseModelHarness) -> None: self.cfg = cfg self.criterion: Callable[..., torch.Tensor] = modelHarness.get_criterion() self.model: nn.Module = modelHarness.model + self.mix_historic_data: bool = cfg.continual_learning.mix_historic_data def fwd_bwd( self, batch: tuple[torch.Tensor, torch.Tensor], hist_batch: Optional[tuple[torch.Tensor, torch.Tensor]] = None, ) -> float: - """Perform forward and backward pass on current batch.""" + """Perform forward and backward pass. + + When `continual_learning.mix_historic_data` is enabled and a historical + batch is supplied, half of each stream is taken so the total sample count + per step stays the same as in the un-mixed case. If the historical batch + is too small to supply its half, the remainder is taken from the current + batch rather than shrinking the step. + """ x, y = batch + if self.mix_historic_data and hist_batch is not None: + x_hist, y_hist = hist_batch + n_total = x.shape[0] + n_hist = min(n_total // 2, x_hist.shape[0]) + n_cur = n_total - n_hist + x = torch.cat([x[:n_cur], x_hist[:n_hist]], dim=0) + y = torch.cat([y[:n_cur], y_hist[:n_hist]], dim=0) outputs = self.model(x) loss = self.criterion(outputs, y) From f9ad6d5d0446c93aee426f8985a399bf65215de0 Mon Sep 17 00:00:00 2001 From: Ana Gainaru Date: Tue, 21 Jul 2026 09:56:50 -0400 Subject: [PATCH 3/4] Adding tests for mixing hisotric and new data --- tests/test_updaters.py | 67 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/tests/test_updaters.py b/tests/test_updaters.py index e9efddb..731e572 100644 --- a/tests/test_updaters.py +++ b/tests/test_updaters.py @@ -47,6 +47,73 @@ def test_hooks_are_noop(self, dummy_harness): updater.update_post_optimizer_call() updater.cl_postprocessing() + @staticmethod + def _grads(model): + return [p.grad.clone() for p in model.parameters() if p.grad is not None] + + def _mixing_updater(self, default_cfg, make_harness, enabled: bool): + cfg = replace( + default_cfg, + continual_learning=ContinualLearningCfg(mix_historic_data=enabled), + ) + harness = make_harness(cfg) + return BaseUpdater(cfg=cfg, modelHarness=harness), harness + + def test_hist_batch_ignored_by_default(self, dummy_harness): + """mix_historic_data defaults to False, so hist_batch must not change grads.""" + updater = BaseUpdater(cfg=dummy_harness.cfg, modelHarness=dummy_harness) + assert updater.mix_historic_data is False + batch = (torch.randn(4, 4), torch.randint(0, 3, (4,))) + hist = (torch.randn(4, 4), torch.randint(0, 3, (4,))) + + dummy_harness.model.zero_grad() + updater.fwd_bwd(batch) + grad_cur_only = self._grads(dummy_harness.model) + + dummy_harness.model.zero_grad() + updater.fwd_bwd(batch, hist) + grad_with_hist = self._grads(dummy_harness.model) + + for a, b in zip(grad_cur_only, grad_with_hist): + assert torch.allclose(a, b, atol=1e-6) + + def test_hist_batch_is_mixed_when_enabled(self, default_cfg, make_harness): + updater, harness = self._mixing_updater(default_cfg, make_harness, enabled=True) + batch = (torch.randn(4, 4), torch.randint(0, 3, (4,))) + hist = (torch.randn(4, 4), torch.randint(0, 3, (4,))) + + harness.model.zero_grad() + updater.fwd_bwd(batch) + grad_cur_only = self._grads(harness.model) + + harness.model.zero_grad() + updater.fwd_bwd(batch, hist) + grad_mixed = self._grads(harness.model) + + assert any( + not torch.allclose(a, b) for a, b in zip(grad_cur_only, grad_mixed) + ), "hist_batch did not influence the gradients" + + def test_hist_batch_matches_half_and_half_pass(self, default_cfg, make_harness): + """Mixing must equal one fwd/bwd over half the current plus half the hist batch.""" + updater, harness = self._mixing_updater(default_cfg, make_harness, enabled=True) + batch = (torch.randn(8, 4), torch.randint(0, 3, (8,))) + hist = (torch.randn(8, 4), torch.randint(0, 3, (8,))) + + harness.model.zero_grad() + updater.fwd_bwd(batch, hist) + grad_mixed = self._grads(harness.model) + + harness.model.zero_grad() + x = torch.cat([batch[0][:4], hist[0][:4]], dim=0) + y = torch.cat([batch[1][:4], hist[1][:4]], dim=0) + loss = harness.get_criterion()(harness.model(x), y) + (loss / harness.cfg.train.grad_accumulation_steps).backward() + grad_ref = self._grads(harness.model) + + for a, b in zip(grad_mixed, grad_ref): + assert torch.allclose(a, b, atol=1e-6) + def test_grad_accumulation_scaling(self, default_cfg, make_harness): cfg2 = replace( default_cfg, From 5016ccfa2250190c55ff3cc67e758e094b63fef9 Mon Sep 17 00:00:00 2001 From: Ana Gainaru Date: Wed, 22 Jul 2026 10:12:31 -0400 Subject: [PATCH 4/4] Change the MNIST example to concatenate all historic images in the historic dataloaders --- examples/mnist/model.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/examples/mnist/model.py b/examples/mnist/model.py index 3dd5fe9..c43d525 100644 --- a/examples/mnist/model.py +++ b/examples/mnist/model.py @@ -142,21 +142,29 @@ def get_hist_dataloaders( self, ) -> Tuple[Optional[DataLoader], Optional[DataLoader]]: """ - If no history yet: add current drift to history and return (None, None). - Else: return loaders over ConcatDataset of prior drifts, then append current drift to history. - Effective dataset length = len(aug_history) * len(full_split). + If no history yet: return (None, None). + Else: return loaders over a replay buffer spanning EVERY prior cumulative + regime (0 … k-1), not just the immediately-previous one. Each past regime + `i` is the composition ``FixedAffine(aug_history[:i])``, matching the data + that was actually streamed at window i-1. Concatenating all of them keeps + the model anchored to the earliest regimes so CL adaptation does not + specialise to the most recent drift alone. + Effective dataset length = (len(aug_history) - 1) * len(full_split). """ if self.task_counter == 1: return None, None - # Concatenate FULL train/val views for each historical drift + # One FULL train/val view per prior cumulative regime: regimes 0 … k-1. + # aug_history[-1] is the current regime, so we stop before it. train_views = [ TransformedView( - self.ds_train, x_transform=FixedAffine(self.aug_history[:-1]) + self.ds_train, x_transform=FixedAffine(self.aug_history[:i]) ) + for i in range(1, len(self.aug_history)) ] val_views = [ - TransformedView(self.ds_val, x_transform=FixedAffine(self.aug_history[:-1])) + TransformedView(self.ds_val, x_transform=FixedAffine(self.aug_history[:i])) + for i in range(1, len(self.aug_history)) ] ds_hist_train: ConcatDataset[Any] = ConcatDataset(train_views)