perf: vectorize SineCosine/MFO updates; sep-CMA-ES for high dims

SineCosine and MFO used per-element Python double loops (swarm x dim) with a .item() sync per element -- ~18000 GPU<->CPU syncs per step at d=200, which made a single run take ~65 minutes and repeatedly exhausted Kaggle sessions. Both updates are elementwise and now fully vectorized with torch ops: SineCosine drops from ~23s/step to ~0.018s/step (about 1000x), MFO similarly.

CMA-ES now switches to pycma's separable/diagonal covariance (sep-CMA-ES, O(d) memory) above diagonal_threshold=5000 dimensions, so it can run on NN-scale parameter vectors instead of OOMing on a d x d covariance matrix (a 250k-param net would need ~500 GB). Emits a warning when the fallback engages.

Full suite: 94 passed, 1 skipped.

Author: hallelx2

swarmtorch/bio_inspired/model_training/sine_cosine.py

@@ -32,20 +32,16 @@ def _update_positions(self) -> None:
         a = 2
         r1 = a - self.iteration_count * (a / max_iter)
 
-        for i in range(self.swarm_size):
-            for j in range(self.positions.shape[1]):
-                r2 = torch.rand(1, device=self.device).item() * 2 * 3.14159
-                r3 = torch.rand(1, device=self.device).item() * 2
-                r4 = torch.rand(1, device=self.device).item()
-
-                if r4 < 0.5:
-                    self.positions[i, j] = self.positions[i, j] + r1 * math.sin(
-                        r2
-                    ) * torch.abs(r3 * self.best_position[j] - self.positions[i, j])
-                else:
-                    self.positions[i, j] = self.positions[i, j] + r1 * math.cos(
-                        r2
-                    ) * torch.abs(r3 * self.best_position[j] - self.positions[i, j])
+        # Vectorized SCA update (formerly a per-element Python double loop,
+        # which forced millions of GPU<->CPU syncs and dominated wall-clock).
+        # r2 in [0, 2*pi), r3 in [0, 2), r4 in [0, 1) per element.
+        r2 = torch.rand_like(self.positions) * (2 * math.pi)
+        r3 = torch.rand_like(self.positions) * 2
+        r4 = torch.rand_like(self.positions)
+
+        target = torch.abs(r3 * self.best_position.unsqueeze(0) - self.positions)
+        trig = torch.where(r4 < 0.5, torch.sin(r2), torch.cos(r2))
+        self.positions = self.positions + r1 * trig * target
 
         self._set_params(self.best_position)
         self.iteration_count += 1
@@ -87,17 +83,26 @@ def _update_positions(self) -> None:
 
         max_iter = 1000
         t = (self.iteration_count / max_iter) * 2 - 1
-
-        for i in range(self.swarm_size):
-            for j in range(self.positions.shape[1]):
-                flame_idx = int(i * self.flames.shape[0] / self.swarm_size)
-                distance = torch.abs(self.positions[i, j] - self.flames[flame_idx, j])
-                b = 1
-                t_val = t * (1 - j / self.positions.shape[1])
-                self.positions[i, j] = (
-                    distance * math.exp(b * t_val) * math.cos(2 * 3.14159 * t_val)
-                    + self.flames[flame_idx, j]
-                )
+        b = 1.0
+        d = self.positions.shape[1]
+        n_flames = self.flames.shape[0]
+
+        # Vectorized MFO update (was a per-element Python double loop).
+        # Each particle i spirals around flame floor(i * n_flames / swarm).
+        flame_idx = (
+            torch.arange(self.swarm_size, device=self.device) * n_flames
+            // self.swarm_size
+        ).clamp(max=n_flames - 1)
+        chosen = self.flames[flame_idx]  # (swarm_size, d)
+
+        # t_val decays across dimensions: t * (1 - j/d), broadcast over particles.
+        j = torch.arange(d, device=self.device, dtype=self.positions.dtype)
+        t_val = (t * (1 - j / d)).unsqueeze(0)  # (1, d)
+
+        distance = torch.abs(self.positions - chosen)
+        self.positions = (
+            distance * torch.exp(b * t_val) * torch.cos(2 * math.pi * t_val) + chosen
+        )
 
         self._set_params(self.best_position)
         self.iteration_count += 1

swarmtorch/evolutionary/model_training/cmaes.py

@@ -55,6 +55,7 @@ def __init__(
         device: str = "cpu",
         init_strategy: str = "model",
         init_sigma: float = 0.1,
+        diagonal_threshold: int = 5000,
         **kwargs: Any,
     ) -> None:
         # CMA-ES picks its own default popsize if ``swarm_size`` is None.
@@ -70,6 +71,11 @@ def __init__(
         )
         self._user_swarm_size = swarm_size
         self.sigma0 = sigma0
+        # Above this dimension, the full d x d covariance matrix is
+        # infeasible (a 250k-param net needs ~500 GB). Switch pycma to its
+        # separable/diagonal covariance (sep-CMA-ES) automatically, which
+        # is O(d) memory and the standard high-dimensional variant.
+        self.diagonal_threshold = diagonal_threshold
         self.iteration_count = 0
         self.best_position: torch.Tensor | None = None
         self.best_fitness = torch.tensor(float("inf"), device=self.device)
@@ -99,6 +105,22 @@ def _init_swarm(self) -> None:
         if self._user_swarm_size is not None and self._user_swarm_size > 0:
             opts["popsize"] = int(self._user_swarm_size)
 
+        # High-dimensional problems (e.g. NN weight spaces) cannot hold a
+        # full covariance matrix. Use the separable/diagonal variant so the
+        # optimizer runs in O(d) instead of O(d^2) and never OOMs.
+        if d > self.diagonal_threshold:
+            import warnings
+
+            warnings.warn(
+                f"CMA-ES: d={d} exceeds diagonal_threshold="
+                f"{self.diagonal_threshold}; using separable (diagonal) "
+                f"covariance (sep-CMA-ES) to stay feasible.",
+                stacklevel=2,
+            )
+            # CMA_diagonal=True keeps the covariance diagonal for the whole
+            # run; this is the documented pycma high-dimensional mode.
+            opts["CMA_diagonal"] = True
+
         self._es = cma.CMAEvolutionStrategy(x0, sigma0, opts)
         # Resolved population size (pycma may pick its own default).
         self.swarm_size = int(self._es.popsize)