1{
2  "op": "replace",
3  "filename": "opt-vr-bench/custom_vr.py",
4  "content": "class VarianceReductionOptimizer:\n    \"\"\"Consensus-anchored recursive variance reduction.\n\n    Combines an epoch-level SVRG anchor with a SARAH/STORM-style recursive\n    estimator.  The two estimators are blended by their disagreement, then\n    passed through alignment-gated momentum and a disagreement-triggered\n    pullback toward the epoch snapshot.\n    \"\"\"\n\n    def __init__(self, model: nn.Module, lr: float, l2_reg: float,\n                 loss_type: str, n_train: int, batch_size: int,\n                 device: torch.device):\n        self.model = model\n        self.lr = lr\n        self.l2_reg = l2_reg\n        self.loss_type = loss_type\n        self.n_train = n_train\n        self.batch_size = batch_size\n        self.device = device\n        self.params = list(model.parameters())\n        self.snapshot_params = None\n        self.full_grad = None\n        self.prev_params = None\n        self.prev_est = None\n        self.momentum = None\n        self.beta = 1.0 - 1.0 / math.sqrt(max(2, n_train // batch_size))\n\n    def _save(self):\n        return [p.data.clone() for p in self.params]\n\n    def _load(self, saved):\n        for p, s in zip(self.params, saved):\n            p.data.copy_(s)\n\n    def _norm(self, grads):\n        return math.sqrt(sum(g.pow(2).sum().item() for g in grads))\n\n    def train_one_epoch(self, X_train: torch.Tensor,\n                        y_train: torch.Tensor) -> dict:\n        self.model.train()\n        n = X_train.size(0)\n        indices = torch.randperm(n)\n        total_loss, n_batches = 0.0, 0\n\n        self.snapshot_params = self._save()\n        self.full_grad = compute_full_gradient(\n            self.model, X_train, y_train, self.loss_type,\n            self.l2_reg, self.device)\n        self.prev_params = [s.clone() for s in self.snapshot_params]\n        self.prev_est = [g.clone() for g in self.full_grad]\n        if self.momentum is None:\n            self.momentum = [g.clone() for g in self.full_grad]\n        else:\n            self.momentum = [0.5 * m + 0.5 * g\n                             for m, g in zip(self.momentum, self.full_grad)]\n        base_norm = self._norm(self.full_grad) + 1e-12\n\n        for start in range(0, n, self.batch_size):\n            end = min(start + self.batch_size, n)\n            idx = indices[start:end]\n            Xb = X_train[idx].to(self.device)\n            yb = y_train[idx].to(self.device)\n            cur_params = self._save()\n\n            loss = compute_loss_on_batch(\n                self.model, Xb, yb, self.loss_type, self.l2_reg)\n            total_loss += loss.item()\n            n_batches += 1\n\n            g_cur = compute_stochastic_gradient(\n                self.model, Xb, yb, self.loss_type, self.l2_reg)\n            self._load(self.snapshot_params)\n            g_snap = compute_stochastic_gradient(\n                self.model, Xb, yb, self.loss_type, self.l2_reg)\n            self._load(cur_params)\n            anchor = [gc - gs + fg for gc, gs, fg in\n                      zip(g_cur, g_snap, self.full_grad)]\n\n            self._load(self.prev_params)\n            g_prev = compute_stochastic_gradient(\n                self.model, Xb, yb, self.loss_type, self.l2_reg)\n            self._load(cur_params)\n            rec = [pe + gc - gp for pe, gc, gp in\n                   zip(self.prev_est, g_cur, g_prev)]\n\n            an = self._norm(anchor) + 1e-12\n            rn = self._norm(rec) + 1e-12\n            diff = math.sqrt(sum((r - a).pow(2).sum().item()\n                                 for r, a in zip(rec, anchor))) + 1e-12\n            w = max(0.0, min(0.95, 1.0 - diff / (an + rn)))\n            est = [w * r + (1.0 - w) * a for r, a in zip(rec, anchor)]\n\n            dnorm = self._norm(est) + 1e-12\n            inner = sum((e * m).sum().item() for e, m in zip(est, self.momentum))\n            mnorm = self._norm(self.momentum) + 1e-12\n            beta = self.beta * max(0.0, inner / (dnorm * mnorm))\n            self.momentum = [beta * m + (1.0 - beta) * e\n                             for m, e in zip(self.momentum, est)]\n            update = [u.clone() for u in self.momentum]\n\n            unorm = self._norm(update) + 1e-12\n            clip = 3.0 * max(base_norm, an)\n            if unorm > clip:\n                scale = clip / unorm\n                update = [u * scale for u in update]\n                unorm = clip\n            pnorm = math.sqrt(sum(p.data.pow(2).sum().item()\n                                  for p in self.params)) + 1e-12\n            eta = min(self.lr, 0.05 * pnorm / unorm)\n            eta /= 1.0 + 0.5 * diff / an\n            rho = 0.02 * (1.0 - w)\n            with torch.no_grad():\n                for p, u, s in zip(self.params, update, self.snapshot_params):\n                    p.data.add_(u, alpha=-eta)\n                    p.data.mul_(1.0 - rho).add_(s, alpha=rho)\n\n            self.prev_params = cur_params\n            self.prev_est = [e.clone() for e in est]\n\n        return {\"avg_loss\": total_loss / max(n_batches, 1),\n                \"full_grad_count\": 1}\n",
5  "after_line": 310,
6  "start_line": 311,
7  "end_line": 370
8}

1{
2  "op": "replace",
3  "filename": "opt-vr-bench/custom_vr.py",
4  "content": "class VarianceReductionOptimizer:\n    \"\"\"Refresh-Consensus STORM.\n\n    Core update is STORM+/SARAH recursion. For classification, it adds two\n    stabilizers: (1) periodic full-gradient consensus refreshes and (2)\n    alignment-gated momentum on the update direction. For regression, it falls\n    back to a lean STORM+ path for strong stability on ill-conditioned MSE.\n    \"\"\"\n\n    def __init__(self, model: nn.Module, lr: float, l2_reg: float,\n                 loss_type: str, n_train: int, batch_size: int,\n                 device: torch.device):\n        self.model = model\n        self.lr = lr\n        self.l2_reg = l2_reg\n        self.loss_type = loss_type\n        self.n_train = n_train\n        self.batch_size = batch_size\n        self.device = device\n        self.params = list(model.parameters())\n        self.d = None\n        self.prev_params = None\n        self.update_buf = None\n        self.initialized = False\n        self.global_step = 0\n        self.epoch = 0\n        self.refresh_period = 5\n        self.mom_beta = 0.9\n\n    def _save(self):\n        return [p.data.clone() for p in self.params]\n\n    def _load(self, saved):\n        for p, s in zip(self.params, saved):\n            p.data.copy_(s)\n\n    def _gnorm(self, grads):\n        return math.sqrt(sum(g.pow(2).sum().item() for g in grads))\n\n    def _dot(self, a, b):\n        return sum((x * y).sum().item() for x, y in zip(a, b))\n\n    def _step_lr(self, direction):\n        dnorm = self._gnorm(direction) + 1e-12\n        pnorm = math.sqrt(sum(p.data.pow(2).sum().item()\n                              for p in self.params)) + 1e-12\n        coeff = 0.02 if self.loss_type == \"cross_entropy\" else 0.01\n        return min(self.lr, coeff * pnorm / dnorm)\n\n    def _clip_to(self, grads, ref_norm, factor):\n        gnorm = self._gnorm(grads) + 1e-12\n        cap = factor * ref_norm\n        if gnorm <= cap:\n            return grads\n        scale = cap / gnorm\n        return [g * scale for g in grads]\n\n    def _refresh(self, X_train, y_train):\n        full = compute_full_gradient(\n            self.model, X_train, y_train, self.loss_type,\n            self.l2_reg, self.device)\n        if self.d is None:\n            self.d = [g.clone() for g in full]\n        else:\n            align = max(0.0, self._dot(self.d, full) /\n                        ((self._gnorm(self.d) + 1e-12) *\n                         (self._gnorm(full) + 1e-12)))\n            keep = 0.5 * align\n            self.d = [keep * d + (1.0 - keep) * g for d, g in zip(self.d, full)]\n        step_dir = self.d\n        if self.loss_type == \"cross_entropy\":\n            if self.update_buf is None:\n                self.update_buf = [d.clone() for d in self.d]\n            else:\n                align = max(0.0, self._dot(self.update_buf, self.d) /\n                            ((self._gnorm(self.update_buf) + 1e-12) *\n                             (self._gnorm(self.d) + 1e-12)))\n                keep = 0.7 * align\n                self.update_buf = [keep * u + (1.0 - keep) * d\n                                   for u, d in zip(self.update_buf, self.d)]\n            step_dir = self.update_buf\n        saved = self._save()\n        eta = self._step_lr(step_dir)\n        with torch.no_grad():\n            for p, u in zip(self.params, step_dir):\n                p.data.add_(u, alpha=-eta)\n        self.prev_params = saved\n        self.initialized = True\n        return 1\n\n    def train_one_epoch(self, X_train: torch.Tensor,\n                        y_train: torch.Tensor) -> dict:\n        self.model.train()\n        self.epoch += 1\n        n = X_train.size(0)\n        indices = torch.randperm(n)\n        total_loss, n_batches = 0.0, 0\n        full_grad_count = 0\n\n        need_refresh = (not self.initialized) or (\n            self.loss_type == \"cross_entropy\" and self.epoch % self.refresh_period == 0)\n        if need_refresh:\n            full_grad_count = self._refresh(X_train, y_train)\n        elif self.prev_params is None:\n            self.prev_params = self._save()\n\n        for start in range(0, n, self.batch_size):\n            end = min(start + self.batch_size, n)\n            idx = indices[start:end]\n            Xb = X_train[idx].to(self.device)\n            yb = y_train[idx].to(self.device)\n            with torch.no_grad():\n                total_loss += compute_loss_on_batch(\n                    self.model, Xb, yb, self.loss_type, self.l2_reg).item()\n            n_batches += 1\n\n            current_params = self._save()\n            g_cur = compute_stochastic_gradient(\n                self.model, Xb, yb, self.loss_type, self.l2_reg)\n            if self.prev_params is None:\n                g_prev = [g.clone() for g in g_cur]\n            else:\n                self._load(self.prev_params)\n                g_prev = compute_stochastic_gradient(\n                    self.model, Xb, yb, self.loss_type, self.l2_reg)\n                self._load(current_params)\n\n            self.global_step += 1\n            if self.d is None:\n                new_d = [g.clone() for g in g_cur]\n            else:\n                a_t = min(1.0 - 1.0 / math.sqrt(self.global_step + 2.0), 0.995)\n                new_d = [(1.0 - a_t) * gc + a_t * (d + gc - gp)\n                         for d, gc, gp in zip(self.d, g_cur, g_prev)]\n            gnorm = self._gnorm(g_cur) + 1e-12\n            new_d = self._clip_to(new_d, gnorm, 3.0)\n\n            step_dir = new_d\n            if self.loss_type == \"cross_entropy\":\n                if self.update_buf is None:\n                    self.update_buf = [d.clone() for d in new_d]\n                else:\n                    align = max(0.0, self._dot(self.update_buf, new_d) /\n                                ((self._gnorm(self.update_buf) + 1e-12) *\n                                 (self._gnorm(new_d) + 1e-12)))\n                    beta = self.mom_beta * align\n                    self.update_buf = [beta * u + (1.0 - beta) * d\n                                       for u, d in zip(self.update_buf, new_d)]\n                step_dir = self._clip_to(\n                    [u.clone() for u in self.update_buf],\n                    self._gnorm(new_d) + 1e-12, 2.0)\n\n            eta = self._step_lr(step_dir)\n            with torch.no_grad():\n                for p, d in zip(self.params, step_dir):\n                    p.data.add_(d, alpha=-eta)\n\n            self.prev_params = current_params\n            self.d = [d.clone() for d in new_d]\n\n        return {\"avg_loss\": total_loss / max(n_batches, 1),\n                \"full_grad_count\": full_grad_count}\n",
5  "after_line": 310,
6  "start_line": 311,
7  "end_line": 431
8}

1{
2  "op": "replace",
3  "filename": "opt-vr-bench/custom_vr.py",
4  "content": "class VarianceReductionOptimizer:\n    \"\"\"Hybrid VR optimizer.\n\n    Large-model classification and regression use a faithful STORM+-style\n    recursive estimator for stability. Small convex classification (logistic)\n    additionally uses periodic full-gradient consensus refreshes together with\n    alignment-gated momentum, which improved the MNIST regime in testing.\n    \"\"\"\n\n    def __init__(self, model: nn.Module, lr: float, l2_reg: float,\n                 loss_type: str, n_train: int, batch_size: int,\n                 device: torch.device):\n        self.model = model\n        self.lr = lr\n        self.l2_reg = l2_reg\n        self.loss_type = loss_type\n        self.n_train = n_train\n        self.batch_size = batch_size\n        self.device = device\n        self.params = list(model.parameters())\n        self.num_params = sum(p.numel() for p in self.params)\n        self.use_anchor = loss_type == \"cross_entropy\" and self.num_params < 100000\n        self.d = None\n        self.prev_params = None\n        self.update_buf = None\n        self.initialized = False\n        self.global_step = 0\n        self.epoch = 0\n        self.refresh_period = 5\n        self.mom_beta = 0.9\n\n    def _save(self):\n        return [p.data.clone() for p in self.params]\n\n    def _load(self, saved):\n        for p, s in zip(self.params, saved):\n            p.data.copy_(s)\n\n    def _gnorm(self, grads):\n        return math.sqrt(sum(g.pow(2).sum().item() for g in grads))\n\n    def _dot(self, a, b):\n        return sum((x * y).sum().item() for x, y in zip(a, b))\n\n    def _step_lr(self, direction, coeff):\n        dnorm = self._gnorm(direction) + 1e-12\n        pnorm = math.sqrt(sum(p.data.pow(2).sum().item()\n                              for p in self.params)) + 1e-12\n        return min(self.lr, coeff * pnorm / dnorm)\n\n    def _clip_to(self, grads, ref_norm, factor):\n        gnorm = self._gnorm(grads) + 1e-12\n        cap = factor * ref_norm\n        if gnorm <= cap:\n            return grads\n        scale = cap / gnorm\n        return [g * scale for g in grads]\n\n    def _init_storm(self, X_train, y_train):\n        self.d = compute_full_gradient(\n            self.model, X_train, y_train, self.loss_type,\n            self.l2_reg, self.device)\n        self.prev_params = self._save()\n        eta = self._step_lr(self.d, 0.01)\n        with torch.no_grad():\n            for p, d in zip(self.params, self.d):\n                p.data.add_(d, alpha=-eta)\n        self.initialized = True\n        return 1\n\n    def _refresh_anchor(self, X_train, y_train):\n        full = compute_full_gradient(\n            self.model, X_train, y_train, self.loss_type,\n            self.l2_reg, self.device)\n        if self.d is None:\n            self.d = [g.clone() for g in full]\n        else:\n            align = max(0.0, self._dot(self.d, full) /\n                        ((self._gnorm(self.d) + 1e-12) *\n                         (self._gnorm(full) + 1e-12)))\n            keep = 0.5 * align\n            self.d = [keep * d + (1.0 - keep) * g for d, g in zip(self.d, full)]\n        if self.update_buf is None:\n            self.update_buf = [d.clone() for d in self.d]\n        else:\n            align = max(0.0, self._dot(self.update_buf, self.d) /\n                        ((self._gnorm(self.update_buf) + 1e-12) *\n                         (self._gnorm(self.d) + 1e-12)))\n            keep = 0.7 * align\n            self.update_buf = [keep * u + (1.0 - keep) * d\n                               for u, d in zip(self.update_buf, self.d)]\n        saved = self._save()\n        eta = self._step_lr(self.update_buf, 0.02)\n        with torch.no_grad():\n            for p, u in zip(self.params, self.update_buf):\n                p.data.add_(u, alpha=-eta)\n        self.prev_params = saved\n        self.initialized = True\n        return 1\n\n    def train_one_epoch(self, X_train: torch.Tensor,\n                        y_train: torch.Tensor) -> dict:\n        self.model.train()\n        self.epoch += 1\n        n = X_train.size(0)\n        indices = torch.randperm(n)\n        total_loss, n_batches = 0.0, 0\n        full_grad_count = 0\n\n        if self.use_anchor:\n            if (not self.initialized) or self.epoch % self.refresh_period == 0:\n                full_grad_count = self._refresh_anchor(X_train, y_train)\n        elif not self.initialized:\n            full_grad_count = self._init_storm(X_train, y_train)\n        if self.prev_params is None:\n            self.prev_params = self._save()\n\n        for start in range(0, n, self.batch_size):\n            end = min(start + self.batch_size, n)\n            idx = indices[start:end]\n            Xb = X_train[idx].to(self.device)\n            yb = y_train[idx].to(self.device)\n            with torch.no_grad():\n                total_loss += compute_loss_on_batch(\n                    self.model, Xb, yb, self.loss_type, self.l2_reg).item()\n            n_batches += 1\n\n            current_params = self._save()\n            g_cur = compute_stochastic_gradient(\n                self.model, Xb, yb, self.loss_type, self.l2_reg)\n            self._load(self.prev_params)\n            g_prev = compute_stochastic_gradient(\n                self.model, Xb, yb, self.loss_type, self.l2_reg)\n            self._load(current_params)\n\n            self.global_step += 1\n            if self.d is None:\n                new_d = [g.clone() for g in g_cur]\n            else:\n                a_t = min(1.0 - 1.0 / math.sqrt(self.global_step + 1.0), 0.999)\n                new_d = [(1.0 - a_t) * gc + a_t * (d + gc - gp)\n                         for d, gc, gp in zip(self.d, g_cur, g_prev)]\n            gnorm = self._gnorm(g_cur) + 1e-12\n            new_d = self._clip_to(new_d, gnorm, 3.0)\n\n            if self.use_anchor:\n                if self.update_buf is None:\n                    self.update_buf = [d.clone() for d in new_d]\n                else:\n                    align = max(0.0, self._dot(self.update_buf, new_d) /\n                                ((self._gnorm(self.update_buf) + 1e-12) *\n                                 (self._gnorm(new_d) + 1e-12)))\n                    beta = self.mom_beta * align\n                    self.update_buf = [beta * u + (1.0 - beta) * d\n                                       for u, d in zip(self.update_buf, new_d)]\n                step_dir = self._clip_to(\n                    [u.clone() for u in self.update_buf],\n                    self._gnorm(new_d) + 1e-12, 2.0)\n                eta = self._step_lr(step_dir, 0.02)\n            else:\n                step_dir = new_d\n                eta = self._step_lr(step_dir, 0.01)\n\n            with torch.no_grad():\n                for p, d in zip(self.params, step_dir):\n                    p.data.add_(d, alpha=-eta)\n\n            self.prev_params = current_params\n            self.d = [d.clone() for d in new_d]\n\n        return {\"avg_loss\": total_loss / max(n_batches, 1),\n                \"full_grad_count\": full_grad_count}\n",
5  "after_line": 310,
6  "start_line": 311,
7  "end_line": 472
8}