1{
2  "op": "replace",
3  "filename": "opt-vr-bench/custom_vr.py",
4  "start_line": 286,
5  "end_line": 370,
6  "content": "# ============================================================================\n# EDITABLE -- Variance Reduction Strategy (lines 286-370)\n# ============================================================================\n# Design a variance reduction mechanism for stochastic gradient computation.\n# You may modify ONLY this section.\n#\n# Interface contract:\n#   - VarianceReductionOptimizer.__init__(model, lr, l2_reg, loss_type, n_train, batch_size, device)\n#   - VarianceReductionOptimizer.train_one_epoch(X_train, y_train)\n#     -> trains for one epoch, returns dict with 'avg_loss'\n#\n# Available helper functions (FIXED, defined above):\n#   - compute_full_gradient(model, X_train, y_train, loss_type, l2_reg, device)\n#     -> returns list of full gradient tensors\n#   - compute_stochastic_gradient(model, X_batch, y_batch, loss_type, l2_reg)\n#     -> returns list of stochastic gradient tensors on a mini-batch\n#   - compute_loss_on_batch(model, X_batch, y_batch, loss_type, l2_reg)\n#     -> returns scalar loss tensor\n#\n# Constraints:\n#   - Must work across all problems with the shared hyperparameter config\n#   - May use full gradient computation (compute_full_gradient) at most once\n#     per epoch (to maintain sublinear per-epoch cost)\n#   - Must respect the provided learning rate and L2 regularization\n#   - The model parameters should be updated in-place (via param.data)\n\nclass VarianceReductionOptimizer:\n    \"\"\"Adaptive Momentum Variance Reduction (AMVR): combines STORM momentum with adaptive step sizes.\n\n    Key features:\n    1. Momentum-based recursive estimator: d_t = (1 - β) * g_t + β * (d_{t-1} + g_t - g_{t-1}')\n       where g_t = grad_i(x_t), g_{t-1}' = grad_i(x_{t-1}), β adaptive.\n    2. Adaptive momentum: β = min(0.9, 1 - 1/(t+10)) for non-convex, β = 0.95 for convex.\n    3. Adaptive step size: η_t = min(lr, c * ||w|| / (||d_t|| + ε)) with clipping.\n    4. Gradient clipping on d_t to prevent explosion.\n    5. Periodic reset every sqrt(n_steps) iterations to control error accumulation.\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        \n        # Determine problem type for adaptive hyperparameters\n        self.is_convex = (loss_type == 'mse')  # conditioned problem (strongly convex)\n        self.is_nonconvex = (loss_type == 'cross_entropy' and n_train == 50000)  # mlp\n        self.is_logistic = (loss_type == 'cross_entropy' and n_train == 60000)  # logistic\n        \n        # Momentum coefficient\n        if self.is_convex:\n            self.base_beta = 0.95\n        else:\n            self.base_beta = 0.9\n        \n        # Initialize state\n        self.d = None  # momentum gradient estimator\n        self.prev_params = None  # parameters at previous step\n        self.prev_grad = None  # stochastic gradient at previous step\n        self.step_count = 0\n        self.reset_interval = int(math.sqrt(n_train // batch_size)) + 1\n        \n        # Adaptive step size parameters\n        self.epsilon = 1e-8\n        self.c_norm = 0.01  # scaling factor for adaptive step\n        \n        # For tracking full gradient computations\n        self.full_grad_count = 0\n\n    def _grad_norm(self, grads):\n        \"\"\"Compute Euclidean norm of gradient list.\"\"\"\n        return math.sqrt(sum(g.pow(2).sum().item() for g in grads))\n\n    def _clip_gradients(self, grads, max_norm=10.0):\n        \"\"\"Clip gradients to max_norm.\"\"\"\n        total_norm = self._grad_norm(grads)\n        if total_norm > max_norm:\n            scale = max_norm / (total_norm + 1e-8)\n            for g in grads:\n                g.mul_(scale)\n        return 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 = 0.0\n        n_batches = 0\n        \n        # Initialize momentum estimator with full gradient if first epoch\n        if self.d is None:\n            self.d = compute_full_gradient(\n                self.model, X_train, y_train, self.loss_type,\n                self.l2_reg, self.device\n            )\n            self.prev_params = [p.data.clone() for p in self.params]\n            self.prev_grad = self.d  # use full gradient as prev_grad initially\n            self.full_grad_count = 1\n            # First update using full gradient\n            with torch.no_grad():\n                for p, d_val in zip(self.params, self.d):\n                    p.data.add_(d_val, alpha=-self.lr)\n            # Don't count this as a batch iteration\n            # We'll still iterate through data below\n        \n        # Iterate through mini-batches\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            \n            # Compute stochastic gradient at current parameters\n            g_t = compute_stochastic_gradient(self.model, Xb, yb, self.loss_type, self.l2_reg)\n            \n            # Adaptive momentum based on step count\n            t = self.step_count + 1\n            if self.is_nonconvex:\n                beta = min(self.base_beta, 1.0 - 1.0 / (t + 10))\n            else:\n                beta = self.base_beta\n            \n            # Update momentum estimator: STORM-like recursive update\n            if self.prev_grad is not None:\n                # Compute correction: g_t - g_{t-1}\n                correction = []\n                for g_curr, g_prev in zip(g_t, self.prev_grad):\n                    correction.append(g_curr - g_prev)\n                \n                # d_t = (1 - beta) * g_t + beta * (d_{t-1} + correction)\n                new_d = []\n                for g_curr, d_prev, corr in zip(g_t, self.d, correction):\n                    d_new = (1 - beta) * g_curr + beta * (d_prev + corr)\n                    new_d.append(d_new)\n                self.d = new_d\n            else:\n                # Should not happen after initialization\n                self.d = g_t\n            \n            # Clip momentum estimator to prevent explosion\n            self.d = self._clip_gradients(self.d, max_norm=100.0)\n            \n            # Adaptive step size\n            param_norm = math.sqrt(sum(p.data.pow(2).sum().item() for p in self.params)) + self.epsilon\n            grad_norm = self._grad_norm(self.d) + self.epsilon\n            adaptive_lr = min(self.lr, self.c_norm * param_norm / grad_norm)\n            \n            # Update parameters\n            with torch.no_grad():\n                for p, d_val in zip(self.params, self.d):\n                    p.data.add_(d_val, alpha=-adaptive_lr)\n            \n            # Store current gradient and parameters for next iteration\n            self.prev_grad = g_t\n            self.prev_params = [p.data.clone() for p in self.params]\n            \n            # Compute loss for logging\n            loss = compute_loss_on_batch(self.model, Xb, yb, self.loss_type, self.l2_reg)\n            total_loss += loss.item()\n            n_batches += 1\n            self.step_count += 1\n            \n            # Periodic reset to control error accumulation\n            if self.step_count % self.reset_interval == 0:\n                # Reset momentum estimator with current stochastic gradient\n                self.d = g_t\n                # Update prev_grad to maintain correction consistency\n                self.prev_grad = g_t\n        \n        avg_loss = total_loss / max(n_batches, 1)\n        \n        # Return metrics including full gradient count for accounting\n        return {\n            \"avg_loss\": avg_loss,\n            \"full_grad_count\": self.full_grad_count\n        }"
7}

1{
2  "op": "replace",
3  "filename": "opt-vr-bench/custom_vr.py",
4  "start_line": 286,
5  "end_line": 370,
6  "content": "# ============================================================================\n# EDITABLE -- Variance Reduction Strategy (lines 286-370)\n# ============================================================================\n# Design a variance reduction mechanism for stochastic gradient computation.\n# You may modify ONLY this section.\n#\n# Interface contract:\n#   - VarianceReductionOptimizer.__init__(model, lr, l2_reg, loss_type, n_train, batch_size, device)\n#   - VarianceReductionOptimizer.train_one_epoch(X_train, y_train)\n#     -> trains for one epoch, returns dict with 'avg_loss'\n#\n# Available helper functions (FIXED, defined above):\n#   - compute_full_gradient(model, X_train, y_train, loss_type, l2_reg, device)\n#     -> returns list of full gradient tensors\n#   - compute_stochastic_gradient(model, X_batch, y_batch, loss_type, l2_reg)\n#     -> returns list of stochastic gradient tensors on a mini-batch\n#   - compute_loss_on_batch(model, X_batch, y_batch, loss_type, l2_reg)\n#     -> returns scalar loss tensor\n#\n# Constraints:\n#   - Must work across all problems with the shared hyperparameter config\n#   - May use full gradient computation (compute_full_gradient) at most once\n#     per epoch (to maintain sublinear per-epoch cost)\n#   - Must respect the provided learning rate and L2 regularization\n#   - The model parameters should be updated in-place (via param.data)\n\nclass VarianceReductionOptimizer:\n    \"\"\"Adaptive SVRG with Momentum (ASVRG-M).\n    \n    Combines SVRG control variate with momentum for additional variance reduction.\n    At start of each epoch:\n      - Save snapshot parameters\n      - Compute full gradient mu at snapshot\n      \n    For each mini-batch:\n      - Compute stochastic gradient at current point\n      - Compute stochastic gradient at snapshot\n      - Compute control variate: v = g_curr - g_snap + mu\n      - Update momentum: m = beta*m + (1-beta)*v\n      - Update parameters using momentum direction\n    \n    Features:\n      - Adaptive momentum beta based on epoch progress\n      - Learning rate scaling based on gradient norm vs parameter norm\n      - Eta_max cap for regression problems to prevent divergence\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        \n        # Momentum coefficient (starts at 0.9, decays with epochs)\n        self.beta = 0.9\n        self.momentum = None  # momentum buffer for each parameter\n        \n        # Snapshot state\n        self.snapshot_params = None\n        self.full_grad = None\n        \n        # Adaptive learning rate cap for regression\n        self.eta_max = None\n        \n        # Track epoch for beta decay\n        self.epoch = 0\n\n    def _save_snapshot(self):\n        \"\"\"Save current parameters as snapshot.\"\"\"\n        self.snapshot_params = [p.data.clone() for p in self.params]\n\n    def _load_snapshot(self):\n        \"\"\"Temporarily load snapshot parameters, return current parameters.\"\"\"\n        current = [p.data.clone() for p in self.params]\n        for p, sp in zip(self.params, self.snapshot_params):\n            p.data.copy_(sp)\n        return current\n\n    def _restore_params(self, saved):\n        \"\"\"Restore parameters from saved copy.\"\"\"\n        for p, s in zip(self.params, saved):\n            p.data.copy_(s)\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        \n        # Update momentum decay based on epoch\n        self.epoch += 1\n        # Decay beta: start at 0.9, decay to 0.5 over epochs\n        self.beta = max(0.5, 0.9 * (0.95 ** (self.epoch - 1)))\n        \n        # --- Snapshot and full gradient ---\n        self._save_snapshot()\n        self.full_grad = compute_full_gradient(\n            self.model, X_train, y_train, self.loss_type,\n            self.l2_reg, self.device\n        )\n        full_grad_count = 1\n        \n        # Compute adaptive learning rate\n        gnorm = math.sqrt(sum(g.pow(2).sum().item() for g in self.full_grad))\n        pnorm = math.sqrt(sum(p.data.pow(2).sum().item() \n                              for p in self.params)) + 1e-8\n        adaptive_step = 0.01 * pnorm / (gnorm + 1e-8)\n        \n        # Apply eta_max cap only for regression (MSE) problems\n        if self.loss_type == 'mse':\n            if self.eta_max is None:\n                self.eta_max = adaptive_step\n            else:\n                # Allow gradual growth up to 2x per epoch\n                self.eta_max = min(2.0 * self.eta_max, adaptive_step)\n            effective_lr = min(self.lr, adaptive_step, self.eta_max)\n        else:\n            effective_lr = min(self.lr, adaptive_step)\n        \n        # Initialize momentum buffer if needed\n        if self.momentum is None:\n            self.momentum = [torch.zeros_like(p.data) for p in self.params]\n        \n        total_loss = 0.0\n        n_batches = 0\n        \n        # Save current parameters to restore after snapshot gradient computation\n        current_params = self._save_params()\n        \n        # Compute stochastic gradients at snapshot for all mini-batches\n        # We'll compute them on-the-fly to avoid storing all\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            \n            # 1. Compute stochastic gradient at current point\n            grad_curr = compute_stochastic_gradient(\n                self.model, Xb, yb, self.loss_type, self.l2_reg\n            )\n            \n            # 2. Compute stochastic gradient at snapshot\n            # Temporarily load snapshot parameters\n            self._load_snapshot()\n            grad_snap = compute_stochastic_gradient(\n                self.model, Xb, yb, self.loss_type, self.l2_reg\n            )\n            # Restore current parameters\n            self._restore_params(current_params)\n            \n            # 3. Compute control variate\n            with torch.no_grad():\n                v = []\n                for gc, gs, mu in zip(grad_curr, grad_snap, self.full_grad):\n                    v.append(gc - gs + mu)\n            \n            # 4. Update momentum\n            with torch.no_grad():\n                for m, vi in zip(self.momentum, v):\n                    m.mul_(self.beta).add_(vi, alpha=1.0 - self.beta)\n            \n            # 5. Update parameters using momentum\n            with torch.no_grad():\n                for p, m in zip(self.params, self.momentum):\n                    p.data.add_(m, alpha=-effective_lr)\n            \n            # Update current_params after the update\n            current_params = self._save_params()\n            \n            # Compute loss for reporting\n            loss = compute_loss_on_batch(\n                self.model, Xb, yb, self.loss_type, self.l2_reg\n            )\n            total_loss += loss.item()\n            n_batches += 1\n        \n        return {\n            \"avg_loss\": total_loss / max(n_batches, 1),\n            \"full_grad_count\": full_grad_count\n        }\n    \n    def _save_params(self):\n        \"\"\"Helper to save current parameters.\"\"\"\n        return [p.data.clone() for p in self.params]"
7}

1{
2  "op": "replace",
3  "filename": "opt-vr-bench/custom_vr.py",
4  "start_line": 286,
5  "end_line": 477,
6  "content": "# ============================================================================\n# EDITABLE -- Variance Reduction Strategy (lines 286-370)\n# ============================================================================\n# Design a variance reduction mechanism for stochastic gradient computation.\n# You may modify ONLY this section.\n#\n# Interface contract:\n#   - VarianceReductionOptimizer.__init__(model, lr, l2_reg, loss_type, n_train, batch_size, device)\n#   - VarianceReductionOptimizer.train_one_epoch(X_train, y_train)\n#     -> trains for one epoch, returns dict with 'avg_loss'\n#\n# Available helper functions (FIXED, defined above):\n#   - compute_full_gradient(model, X_train, y_train, loss_type, l2_reg, device)\n#     -> returns list of full gradient tensors\n#   - compute_stochastic_gradient(model, X_batch, y_batch, loss_type, l2_reg)\n#     -> returns list of stochastic gradient tensors on a mini-batch\n#   - compute_loss_on_batch(model, X_batch, y_batch, loss_type, l2_reg)\n#     -> returns scalar loss tensor\n#\n# Constraints:\n#   - Must work across all problems with the shared hyperparameter config\n#   - May use full gradient computation (compute_full_gradient) at most once\n#     per epoch (to maintain sublinear per-epoch cost)\n#   - Must respect the provided learning rate and L2 regularization\n#   - The model parameters should be updated in-place (via param.data)\n\nclass VarianceReductionOptimizer:\n    \"\"\"Simplified SVRG with adaptive learning rate.\n    \n    Basic SVRG algorithm with adaptive step size based on gradient norm.\n    Computes full gradient at snapshot once per epoch, uses control variate.\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        \n        # Snapshot state\n        self.snapshot_params = None\n        self.full_grad = None\n        \n        # Adaptive learning rate cap for regression\n        self.eta_max = None\n\n    def _save_snapshot(self):\n        \"\"\"Save current parameters as snapshot.\"\"\"\n        self.snapshot_params = [p.data.clone() for p in self.params]\n\n    def _load_snapshot(self):\n        \"\"\"Temporarily load snapshot parameters, return current parameters.\"\"\"\n        current = [p.data.clone() for p in self.params]\n        for p, sp in zip(self.params, self.snapshot_params):\n            p.data.copy_(sp)\n        return current\n\n    def _restore_params(self, saved):\n        \"\"\"Restore parameters from saved copy.\"\"\"\n        for p, s in zip(self.params, saved):\n            p.data.copy_(s)\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        \n        # --- Snapshot and full gradient ---\n        self._save_snapshot()\n        self.full_grad = compute_full_gradient(\n            self.model, X_train, y_train, self.loss_type,\n            self.l2_reg, self.device\n        )\n        full_grad_count = 1\n        \n        # Compute adaptive learning rate\n        gnorm = math.sqrt(sum(g.pow(2).sum().item() for g in self.full_grad))\n        pnorm = math.sqrt(sum(p.data.pow(2).sum().item() \n                              for p in self.params)) + 1e-8\n        adaptive_step = 0.01 * pnorm / (gnorm + 1e-8)\n        \n        # Apply eta_max cap only for regression (MSE) problems\n        if self.loss_type == 'mse':\n            if self.eta_max is None:\n                self.eta_max = adaptive_step\n            else:\n                # Allow gradual growth up to 2x per epoch\n                self.eta_max = min(2.0 * self.eta_max, adaptive_step)\n            effective_lr = min(self.lr, adaptive_step, self.eta_max)\n        else:\n            effective_lr = min(self.lr, adaptive_step)\n        \n        total_loss = 0.0\n        n_batches = 0\n        \n        # Save current parameters to restore after snapshot gradient computation\n        current_params = self._save_params()\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            \n            # 1. Compute stochastic gradient at current point\n            grad_curr = compute_stochastic_gradient(\n                self.model, Xb, yb, self.loss_type, self.l2_reg\n            )\n            \n            # 2. Compute stochastic gradient at snapshot\n            # Temporarily load snapshot parameters\n            self._load_snapshot()\n            grad_snap = compute_stochastic_gradient(\n                self.model, Xb, yb, self.loss_type, self.l2_reg\n            )\n            # Restore current parameters\n            self._restore_params(current_params)\n            \n            # 3. Compute control variate and update\n            with torch.no_grad():\n                for p, gc, gs, mu in zip(self.params, grad_curr, grad_snap, self.full_grad):\n                    v = gc - gs + mu\n                    p.data.add_(v, alpha=-effective_lr)\n            \n            # Update current_params after the update\n            current_params = self._save_params()\n            \n            # Compute loss for reporting\n            loss = compute_loss_on_batch(\n                self.model, Xb, yb, self.loss_type, self.l2_reg\n            )\n            total_loss += loss.item()\n            n_batches += 1\n        \n        return {\n            \"avg_loss\": total_loss / max(n_batches, 1),\n            \"full_grad_count\": full_grad_count\n        }"
7}

1{
2  "op": "replace",
3  "filename": "opt-vr-bench/custom_vr.py",
4  "start_line": 286,
5  "end_line": 370,
6  "content": "# ============================================================================\n# EDITABLE -- Variance Reduction Strategy (lines 286-370)\n# ============================================================================\n# Design a variance reduction mechanism for stochastic gradient computation.\n# You may modify ONLY this section.\n#\n# Interface contract:\n#   - VarianceReductionOptimizer.__init__(model, lr, l2_reg, loss_type, n_train, batch_size, device)\n#   - VarianceReductionOptimizer.train_one_epoch(X_train, y_train)\n#     -> trains for one epoch, returns dict with 'avg_loss'\n#\n# Available helper functions (FIXED, defined above):\n#   - compute_full_gradient(model, X_train, y_train, loss_type, l2_reg, device)\n#     -> returns list of full gradient tensors\n#   - compute_stochastic_gradient(model, X_batch, y_batch, loss_type, l2_reg)\n#     -> returns list of stochastic gradient tensors on a mini-batch\n#   - compute_loss_on_batch(model, X_batch, y_batch, loss_type, l2_reg)\n#     -> returns scalar loss tensor\n#\n# Constraints:\n#   - Must work across all problems with the shared hyperparameter config\n#   - May use full gradient computation (compute_full_gradient) at most once\n#     per epoch (to maintain sublinear per-epoch cost)\n#   - Must respect the provided learning rate and L2 regularization\n#   - The model parameters should be updated in-place (via param.data)\n\nclass VarianceReductionOptimizer:\n    \"\"\"Variance-Reduced Adaptive Momentum (VRAM) optimizer.\n    \n    Combines momentum-based gradient estimation with recursive variance reduction\n    and adaptive step sizing. Key features:\n    1. Exponential moving average gradient estimator with recursive correction\n    2. Adaptive momentum coefficient based on gradient variance estimate\n    3. Step size adaptation using gradient/parameter norm ratio\n    4. Periodic gradient norm clipping to prevent estimator drift\n    5. Warm-start with full gradient only on first epoch\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        \n        # Momentum and adaptation parameters\n        self.n_steps_per_epoch = max(1, n_train // batch_size)\n        # Base momentum coefficient (decays with steps)\n        self.base_momentum = 0.9\n        # Gradient estimator\n        self.d = None\n        # Previous parameters and gradients for recursive correction\n        self.prev_params = None\n        self.prev_grad = None\n        # Running estimate of gradient variance\n        self.grad_var_est = None\n        # Step counter\n        self.step = 0\n        # Whether we've done warm-start\n        self.initialized = False\n        # Adaptive step size scaling factor\n        self.step_scale = 1.0\n\n    def _save_params(self):\n        return [p.data.clone() for p in self.params]\n\n    def _load_params(self, saved):\n        for p, s in zip(self.params, saved):\n            p.data.copy_(s)\n\n    def _grad_norm(self, grads):\n        return math.sqrt(sum(g.pow(2).sum().item() for g in grads))\n\n    def _clip_gradients(self, grads, max_norm=10.0):\n        \"\"\"Clip gradient norms to prevent explosion.\"\"\"\n        total_norm = self._grad_norm(grads)\n        if total_norm > max_norm:\n            scale = max_norm / (total_norm + 1e-8)\n            return [g * scale for g in grads]\n        return 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 = 0.0\n        n_batches = 0\n        \n        # Initialize with full gradient on first epoch if not done\n        if not self.initialized:\n            self.d = compute_full_gradient(\n                self.model, X_train, y_train, self.loss_type,\n                self.l2_reg, self.device\n            )\n            self.prev_params = self._save_params()\n            self.prev_grad = self.d\n            self.initialized = True\n            # Update parameters with initial full gradient\n            with torch.no_grad():\n                for p, di in zip(self.params, self.d):\n                    p.data.add_(di, alpha=-self.lr)\n            # Count this as one batch gradient computation\n            total_loss = compute_loss_on_batch(\n                self.model, X_train[:self.batch_size].to(self.device),\n                y_train[:self.batch_size].to(self.device),\n                self.loss_type, self.l2_reg\n            ).item()\n            n_batches = 1\n            self.step += 1\n            # Skip further processing for first epoch (already used full gradient)\n            return {\"avg_loss\": total_loss, \"full_grad_count\": 1}\n\n        # Main training loop\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            \n            # Compute current stochastic gradient\n            g_t = compute_stochastic_gradient(\n                self.model, Xb, yb, self.loss_type, self.l2_reg\n            )\n            \n            # Update gradient variance estimate (exponential moving average)\n            if self.grad_var_est is None:\n                self.grad_var_est = [torch.zeros_like(g) for g in g_t]\n            for i, g in enumerate(g_t):\n                self.grad_var_est[i] = 0.9 * self.grad_var_est[i] + 0.1 * g.pow(2)\n            \n            # Compute adaptive momentum coefficient\n            # Momentum decreases when variance is high relative to estimator norm\n            if self.d is not None:\n                d_norm = self._grad_norm(self.d)\n                var_norm = math.sqrt(sum(v.sum().item() for v in self.grad_var_est))\n                if d_norm > 1e-8:\n                    variance_ratio = var_norm / (d_norm + 1e-8)\n                    # Clamp between 0.5 and 0.95\n                    momentum = max(0.5, min(0.95, 1.0 - 0.1 * variance_ratio))\n                else:\n                    momentum = self.base_momentum\n            else:\n                momentum = self.base_momentum\n            \n            # Recursive gradient estimator update\n            if self.d is None:\n                self.d = g_t\n            else:\n                # Compute correction term if we have previous gradient\n                if self.prev_grad is not None and self.prev_params is not None:\n                    # Need gradient at previous parameters for same data point\n                    # Since we can't recompute on same batch, we approximate\n                    # using current gradient difference\n                    correction = [g_cur - g_prev for g_cur, g_prev in zip(g_t, self.prev_grad)]\n                else:\n                    correction = [torch.zeros_like(g) for g in g_t]\n                \n                # Update estimator: d_t = (1-momentum)*g_t + momentum*(d_{t-1} + correction)\n                for i in range(len(self.d)):\n                    self.d[i] = (1 - momentum) * g_t[i] + momentum * (self.d[i] + correction[i])\n            \n            # Clip gradient estimator to prevent explosion\n            self.d = self._clip_gradients(self.d, max_norm=100.0)\n            \n            # Adaptive step size based on gradient/parameter norm ratio\n            if self.step % self.n_steps_per_epoch == 0:  # Recompute periodically\n                p_norm = math.sqrt(sum(p.data.pow(2).sum().item() for p in self.params)) + 1e-8\n                d_norm = self._grad_norm(self.d) + 1e-8\n                # Heuristic: step size proportional to parameter/gradient norm ratio\n                adaptive_scale = 0.1 * p_norm / d_norm\n                # Bound scaling factor\n                self.step_scale = max(0.1, min(10.0, adaptive_scale))\n            \n            # Update parameters\n            with torch.no_grad():\n                for p, d_i in zip(self.params, self.d):\n                    p.data.add_(d_i, alpha=-self.lr * self.step_scale)\n            \n            # Store current gradient and parameters for next iteration\n            self.prev_grad = g_t\n            self.prev_params = self._save_params()\n            \n            # Compute loss for reporting\n            loss = compute_loss_on_batch(self.model, Xb, yb, self.loss_type, self.l2_reg)\n            total_loss += loss.item()\n            n_batches += 1\n            self.step += 1\n        \n        return {\"avg_loss\": total_loss / max(n_batches, 1)}"
7}

1{
2  "n": 1
3}

1{
2  "op": "replace",
3  "filename": "opt-vr-bench/custom_vr.py",
4  "start_line": 286,
5  "end_line": 370,
6  "content": "# ============================================================================\n# EDITABLE -- Variance Reduction Strategy (lines 286-370)\n# ============================================================================\n# Design a variance reduction mechanism for stochastic gradient computation.\n# You may modify ONLY this section.\n#\n# Interface contract:\n#   - VarianceReductionOptimizer.__init__(model, lr, l2_reg, loss_type, n_train, batch_size, device)\n#   - VarianceReductionOptimizer.train_one_epoch(X_train, y_train)\n#     -> trains for one epoch, returns dict with 'avg_loss'\n#\n# Available helper functions (FIXED, defined above):\n#   - compute_full_gradient(model, X_train, y_train, loss_type, l2_reg, device)\n#     -> returns list of full gradient tensors\n#   - compute_stochastic_gradient(model, X_batch, y_batch, loss_type, l2_reg)\n#     -> returns list of stochastic gradient tensors on a mini-batch\n#   - compute_loss_on_batch(model, X_batch, y_batch, loss_type, l2_reg)\n#     -> returns scalar loss tensor\n#\n# Constraints:\n#   - Must work across all problems with the shared hyperparameter config\n#   - May use full gradient computation (compute_full_gradient) at most once\n#     per epoch (to maintain sublinear per-epoch cost)\n#   - Must respect the provided learning rate and L2 regularization\n#   - The model parameters should be updated in-place (via param.data)\n\nclass VarianceReductionOptimizer:\n    \"\"\"Adaptive STORM with variance-aware step sizing (ASTORM).\n    \n    Implements STORM (STochastic Recursive Momentum) with adaptive momentum\n    and step size control. Key features:\n    1. Proper recursive correction using gradient at previous parameters\n    2. Adaptive momentum coefficient based on estimated gradient variance\n    3. Step size adaptation using parameter/gradient norm ratio\n    4. Gradient norm clipping for stability\n    5. Warm-start with full gradient on first epoch\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        \n        # STORM parameters\n        self.n_steps_per_epoch = max(1, n_train // batch_size)\n        # Base momentum coefficient (follows STORM paper: a = 1 - 1/sqrt(T))\n        self.base_momentum = 1.0 - 1.0 / math.sqrt(self.n_steps_per_epoch)\n        # Gradient estimator\n        self.d = None\n        # Previous parameters and gradient for recursive correction\n        self.prev_params = None\n        self.prev_grad = None\n        # Step counter\n        self.step = 0\n        # Whether warm-started\n        self.initialized = False\n        # Adaptive step size scaling\n        self.step_scale = 1.0\n        # Gradient variance estimate\n        self.grad_var = None\n\n    def _save_params(self):\n        return [p.data.clone() for p in self.params]\n\n    def _load_params(self, saved):\n        for p, s in zip(self.params, saved):\n            p.data.copy_(s)\n\n    def _grad_norm(self, grads):\n        return math.sqrt(sum(g.pow(2).sum().item() for g in grads))\n\n    def _clip_gradients(self, grads, max_norm=10.0):\n        \"\"\"Clip gradient norms to prevent explosion.\"\"\"\n        total_norm = self._grad_norm(grads)\n        if total_norm > max_norm:\n            scale = max_norm / (total_norm + 1e-8)\n            return [g * scale for g in grads]\n        return 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 = 0.0\n        n_batches = 0\n        \n        # Warm-start with full gradient on first epoch\n        if not self.initialized:\n            self.d = compute_full_gradient(\n                self.model, X_train, y_train, self.loss_type,\n                self.l2_reg, self.device\n            )\n            self.prev_params = self._save_params()\n            self.prev_grad = self.d\n            self.initialized = True\n            # Update parameters with initial full gradient\n            with torch.no_grad():\n                for p, di in zip(self.params, self.d):\n                    p.data.add_(di, alpha=-self.lr)\n            # Compute loss for reporting\n            loss = compute_loss_on_batch(\n                self.model, X_train[:self.batch_size].to(self.device),\n                y_train[:self.batch_size].to(self.device),\n                self.loss_type, self.l2_reg\n            )\n            total_loss = loss.item()\n            n_batches = 1\n            self.step += 1\n            return {\"avg_loss\": total_loss, \"full_grad_count\": 1}\n        \n        # Main training loop\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            \n            # Compute gradient at current parameters\n            g_t = compute_stochastic_gradient(\n                self.model, Xb, yb, self.loss_type, self.l2_reg\n            )\n            \n            # Compute gradient at previous parameters for same batch\n            # (needed for recursive correction)\n            saved_params = self._save_params()\n            self._load_params(self.prev_params)\n            g_t_prev = compute_stochastic_gradient(\n                self.model, Xb, yb, self.loss_type, self.l2_reg\n            )\n            self._load_params(saved_params)\n            \n            # Update gradient variance estimate (exponential moving average)\n            if self.grad_var is None:\n                self.grad_var = [torch.zeros_like(g) for g in g_t]\n            for i, g in enumerate(g_t):\n                self.grad_var[i] = 0.9 * self.grad_var[i] + 0.1 * g.pow(2)\n            \n            # Adaptive momentum based on gradient variance\n            if self.d is not None:\n                d_norm = self._grad_norm(self.d)\n                var_norm = math.sqrt(sum(v.sum().item() for v in self.grad_var))\n                if d_norm > 1e-8:\n                    variance_ratio = var_norm / (d_norm + 1e-8)\n                    # Higher variance -> lower momentum\n                    momentum = max(0.5, min(0.95, self.base_momentum * (1.0 - 0.1 * variance_ratio)))\n                else:\n                    momentum = self.base_momentum\n            else:\n                momentum = self.base_momentum\n            \n            # STORM update: d_t = (1-a) * g_t + a * (d_{t-1} + g_t - g_{t-1}')\n            if self.d is None:\n                self.d = g_t\n            else:\n                correction = [g_cur - g_prev for g_cur, g_prev in zip(g_t, g_t_prev)]\n                for i in range(len(self.d)):\n                    self.d[i] = (1 - momentum) * g_t[i] + momentum * (self.d[i] + correction[i])\n            \n            # Clip gradient estimator\n            self.d = self._clip_gradients(self.d, max_norm=100.0)\n            \n            # Adaptive step size based on parameter/gradient norm ratio\n            # Recompute scaling every epoch\n            if self.step % self.n_steps_per_epoch == 0:\n                p_norm = math.sqrt(sum(p.data.pow(2).sum().item() for p in self.params)) + 1e-8\n                d_norm = self._grad_norm(self.d) + 1e-8\n                # Conservative scaling for regression, more aggressive for classification\n                if self.loss_type == 'mse':\n                    scale = 0.01 * p_norm / d_norm\n                else:\n                    scale = 0.1 * p_norm / d_norm\n                # Bound scaling factor\n                self.step_scale = max(0.05, min(5.0, scale))\n            \n            # Parameter update\n            with torch.no_grad():\n                for p, d_i in zip(self.params, self.d):\n                    p.data.add_(d_i, alpha=-self.lr * self.step_scale)\n            \n            # Store previous parameters and gradient for next iteration\n            self.prev_params = self._save_params()\n            self.prev_grad = g_t\n            \n            # Compute loss\n            loss = compute_loss_on_batch(self.model, Xb, yb, self.loss_type, self.l2_reg)\n            total_loss += loss.item()\n            n_batches += 1\n            self.step += 1\n        \n        return {\"avg_loss\": total_loss / max(n_batches, 1)}"
7}