Adding the option of mixing current data with historic data for CL

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

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)

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

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)

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)

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,