1{
2  "op": "replace",
3  "filename": "PBB/custom_pac_bayes.py",
4  "start_line": 460,
5  "end_line": 604,
6  "content": "class BoundOptimizer:\n    \"\"\"Adaptive PAC-Bayes bound with hybrid formulation and learned temperature.\n\n    Our novel contribution: \n    1. Hybrid bound that adaptively combines Catoni, Quadratic, and McAllester formulations\n    2. Soft-minimum for differentiable optimization of the tightest bound\n    3. Joint learning of temperature parameters for each bound component\n    4. KL-aware adaptive regularization based on bound curvature\n    \n    The bound is: B = soft_min( B_catoni(λ), B_quadratic, B_mcallester )\n    where soft_min is a differentiable approximation via log-sum-exp.\n    \n    We learn λ (Catoni temperature) and the soft-min temperature jointly.\n    \"\"\"\n\n    def __init__(self, learning_rate=0.001, momentum=0.95, prior_sigma=0.1,\n                 pmin=1e-5):\n        self.learning_rate = learning_rate\n        self.momentum = momentum\n        self.prior_sigma = prior_sigma\n        self.pmin = pmin\n        \n        # Learnable parameters for adaptive bound\n        self._lambda_param = torch.tensor(1.0, requires_grad=True)  # Catoni λ ∈ (0,2)\n        self._alpha_param = torch.tensor(0.5, requires_grad=True)   # Mixing weight ∈ (0,1)\n        self._softmin_temp = torch.tensor(1.0, requires_grad=True)  # Temperature for soft-min\n        \n        self._lambda_optimizer = None\n        self._alpha_optimizer = None\n        self._temp_optimizer = None\n\n    def _get_lambda(self):\n        \"\"\"Get clamped lambda value in (0, 2).\"\"\"\n        return torch.clamp(self._lambda_param, min=0.01, max=1.99)\n\n    def _get_alpha(self):\n        \"\"\"Get clamped alpha mixing weight in (0, 1).\"\"\"\n        return torch.sigmoid(self._alpha_param)\n\n    def _get_softmin_temp(self):\n        \"\"\"Get positive soft-min temperature.\"\"\"\n        return torch.exp(self._softmin_temp) + 0.1\n\n    def _ensure_param_optimizers(self):\n        \"\"\"Create optimizers for learnable parameters if needed.\"\"\"\n        if self._lambda_optimizer is None:\n            self._lambda_optimizer = torch.optim.SGD(\n                [self._lambda_param], lr=0.01\n            )\n        if self._alpha_optimizer is None:\n            self._alpha_optimizer = torch.optim.SGD(\n                [self._alpha_param], lr=0.01\n            )\n        if self._temp_optimizer is None:\n            self._temp_optimizer = torch.optim.SGD(\n                [self._softmin_temp], lr=0.01\n            )\n\n    def _catoni_bound(self, empirical_risk, kl, n, delta):\n        \"\"\"Catoni/Lambda bound.\"\"\"\n        lam = self._get_lambda()\n        kl_term = (kl + math.log(2.0 * math.sqrt(n) / delta)) / (\n            n * lam * (1.0 - lam / 2.0)\n        )\n        bound = empirical_risk / (1.0 - lam / 2.0) + kl_term\n        return bound\n\n    def _quadratic_bound(self, empirical_risk, kl, n, delta):\n        \"\"\"Quadratic PAC-Bayes bound.\"\"\"\n        kl_term = (kl + math.log(2.0 * math.sqrt(n) / delta)) / (2.0 * n)\n        # Ensure non-negative under sqrt\n        inner = torch.clamp(empirical_risk + kl_term, min=0.0)\n        kl_term_clamped = torch.clamp(kl_term, min=0.0)\n        bound = (torch.sqrt(inner) + torch.sqrt(kl_term_clamped)) ** 2\n        return bound\n\n    def _mcallester_bound(self, empirical_risk, kl, n, delta):\n        \"\"\"McAllester/Maurer bound.\"\"\"\n        kl_term = (kl + math.log(2.0 * math.sqrt(n) / delta)) / (2.0 * n)\n        bound = empirical_risk + torch.sqrt(kl_term)\n        return bound\n\n    def compute_bound(self, empirical_risk, kl, n, delta):\n        \"\"\"Compute adaptive hybrid bound using soft-minimum of three bounds.\n        \n        We compute all three bounds and combine them via:\n        B_hybrid = -τ * log( w1*exp(-B1/τ) + w2*exp(-B2/τ) + w3*exp(-B3/τ) )\n        where τ is learnable temperature, and weights are based on α.\n        \"\"\"\n        # Compute individual bounds\n        B_cat = self._catoni_bound(empirical_risk, kl, n, delta)\n        B_quad = self._quadratic_bound(empirical_risk, kl, n, delta)\n        B_mcal = self._mcallester_bound(empirical_risk, kl, n, delta)\n        \n        # Get mixing weights from α\n        alpha = self._get_alpha()\n        w1 = alpha  # Weight for Catoni\n        w2 = (1 - alpha) * 0.5  # Weight for Quadratic\n        w3 = (1 - alpha) * 0.5  # Weight for McAllester\n        \n        # Soft-minimum with learnable temperature\n        tau = self._get_softmin_temp()\n        \n        # Compute soft-min: -τ * log(∑ w_i exp(-B_i/τ))\n        # More stable computation: subtract max for numerical stability\n        B_vals = torch.stack([B_cat, B_quad, B_mcal])\n        weights = torch.tensor([w1, w2, w3], device=B_vals.device)\n        \n        # Shift for numerical stability\n        B_max = B_vals.max()\n        exp_terms = weights * torch.exp(-(B_vals - B_max) / tau)\n        log_sum = torch.log(exp_terms.sum() + 1e-10)\n        bound = -tau * log_sum + B_max\n        \n        return bound\n\n    def train_step(self, model, data, target, device, n_bound, delta):\n        \"\"\"Training objective with joint optimization of bounds and parameters.\"\"\"\n        # Ensure learnable parameters are on correct device\n        if self._lambda_param.device != device:\n            self._lambda_param = self._lambda_param.to(device).detach().requires_grad_(True)\n            self._alpha_param = self._alpha_param.to(device).detach().requires_grad_(True)\n            self._softmin_temp = self._softmin_temp.to(device).detach().requires_grad_(True)\n            self._lambda_optimizer = None\n            self._alpha_optimizer = None\n            self._temp_optimizer = None\n        \n        self._ensure_param_optimizers()\n        \n        # Forward pass\n        output = model(data, sample=True)\n        log_probs = F.log_softmax(output, dim=1)\n        log_probs = torch.clamp(log_probs, min=math.log(self.pmin))\n        nll = F.nll_loss(log_probs, target)\n        \n        kl = get_total_kl(model)\n        \n        # Compute hybrid bound\n        bound = self.compute_bound(nll, kl, n_bound, delta)\n        \n        # Additional KL regularization based on bound curvature\n        # When empirical risk is low, encourage smaller KL (like quadratic bound)\n        # When empirical risk is high, allow larger KL (like Catoni bound)\n        alpha = self._get_alpha()\n        kl_weight = 1.0 / (n_bound * (1.0 - alpha + 0.1))  # Adaptive weight\n        reg_loss = kl_weight * kl / n_bound\n        \n        total_loss = bound + reg_loss\n        \n        # Update model parameters\n        # (The outer optimizer will handle this via loss.backward())\n        \n        # Also update our learnable parameters\n        self._lambda_optimizer.zero_grad()\n        self._alpha_optimizer.zero_grad()\n        self._temp_optimizer.zero_grad()\n        \n        return total_loss\n\n    def compute_risk_certificate(self, model, bound_loader, device, delta=0.025,\n                                 mc_samples=1000):\n        \"\"\"Evaluate final risk certificate using PAC-Bayes-kl inversion.\n        \n        We use the tightest possible certificate via binary KL inversion.\n        \"\"\"\n        model.eval()\n        n_bound = len(bound_loader.dataset)\n\n        # 1. Compute empirical 0-1 risk via MC sampling\n        emp_risk_01 = compute_01_risk(model, bound_loader, device,\n                                      mc_samples=mc_samples)\n\n        # 2. Compute NLL-based empirical risk for reference\n        total_nll = 0.0\n        total_samples = 0\n        with torch.no_grad():\n            for data, target in bound_loader:\n                data, target = data.to(device), target.to(device)\n                output = model(data, sample=True)\n                log_probs = F.log_softmax(output, dim=1)\n                log_probs = torch.clamp(log_probs, min=math.log(self.pmin))\n                nll = F.nll_loss(log_probs, target, reduction=\"sum\")\n                total_nll += nll.item()\n                total_samples += target.size(0)\n        emp_nll = total_nll / total_samples\n\n        # 3. KL divergence\n        with torch.no_grad():\n            dummy_data = next(iter(bound_loader))[0][:1].to(device)\n            model(dummy_data, sample=True)\n            kl = get_total_kl(model).item()\n\n        # 4. PAC-Bayes-kl inversion for 0-1 loss certificate (tightest)\n        c = (kl + math.log(2.0 * math.sqrt(n_bound) / delta)) / n_bound\n        risk_cert_01 = inv_kl(emp_risk_01, c)\n\n        # 5. Compute our hybrid bound for CE risk as reference\n        emp_nll_t = torch.tensor(emp_nll)\n        kl_t = torch.tensor(kl)\n        ce_bound = self.compute_bound(emp_nll_t, kl_t, n_bound, delta).item()\n\n        metrics = {\n            \"empirical_01_risk\": emp_risk_01,\n            \"empirical_nll\": emp_nll,\n            \"kl_divergence\": kl,\n            \"ce_bound\": ce_bound,\n            \"lambda\": self._get_lambda().item(),\n            \"alpha\": self._get_alpha().item(),\n            \"softmin_temp\": self._get_softmin_temp().item(),\n        }\n\n        return risk_cert_01, metrics"
7}

1{
2  "op": "replace",
3  "filename": "PBB/custom_pac_bayes.py",
4  "start_line": 460,
5  "end_line": 604,
6  "content": "class BoundOptimizer:\n    \"\"\"PAC-Bayes bound computation and posterior optimization.\n\n    This class controls:\n    1. compute_bound(): How the generalization bound is computed from\n       empirical risk and KL divergence.\n    2. train_step(): The training objective for posterior optimization.\n    3. compute_risk_certificate(): Final bound evaluation after training.\n\n    The training pipeline calls these methods. The goal is to achieve\n    the tightest (lowest) risk certificate on the 0-1 loss.\n\n    Available information:\n    - n_bound: number of samples in the bound-evaluation set\n    - delta: confidence parameter (default 0.025)\n    - kl: KL divergence between posterior and prior KL(Q||P)\n    - empirical_risk: estimated loss on bound-evaluation set\n    - inv_kl(q, c): binary KL inversion (find p s.t. KL(q||p)=c)\n\n    Interface contract:\n    - compute_bound(empirical_risk, kl, n, delta) -> bound_value (float tensor)\n    - train_step(model, data, target, device, n_bound, delta) -> loss (float tensor)\n    - compute_risk_certificate(model, bound_loader, device, delta, mc_samples)\n        -> (risk_cert_01, metrics_dict)\n    \"\"\"\n\n    def __init__(self, learning_rate=0.001, momentum=0.95, prior_sigma=0.1,\n                 pmin=1e-5):\n        self.learning_rate = learning_rate\n        self.momentum = momentum\n        self.prior_sigma = prior_sigma\n        self.pmin = pmin\n\n    def compute_bound(self, empirical_risk, kl, n, delta):\n        \"\"\"Compute PAC-Bayes upper bound on true risk.\n\n        Default: McAllester/Maurer bound (fclassic).\n        B(Q,S) = empirical_risk + sqrt((KL(Q||P) + log(2*sqrt(n)/delta)) / (2n))\n\n        Args:\n            empirical_risk: estimated risk on bound data (tensor)\n            kl: KL divergence between posterior and prior (tensor)\n            n: number of bound-evaluation samples\n            delta: confidence parameter\n\n        Returns:\n            bound_value: upper bound on true risk (tensor)\n        \"\"\"\n        kl_term = (kl + math.log(2.0 * math.sqrt(n) / delta)) / (2.0 * n)\n        bound = empirical_risk + torch.sqrt(kl_term)\n        return bound\n\n    def train_step(self, model, data, target, device, n_bound, delta):\n        \"\"\"Compute training loss (PAC-Bayes objective to minimize).\n\n        Default: McAllester bound with NLL surrogate.\n\n        Args:\n            model: stochastic neural network\n            data: input batch (already on device)\n            target: label batch (already on device)\n            device: torch device\n            n_bound: number of bound-evaluation samples\n            delta: confidence parameter\n\n        Returns:\n            loss: scalar tensor to backpropagate\n        \"\"\"\n        output = model(data, sample=True)\n        # Bounded cross-entropy as surrogate for 0-1 loss\n        log_probs = F.log_softmax(output, dim=1)\n        log_probs = torch.clamp(log_probs, min=math.log(self.pmin))\n        nll = F.nll_loss(log_probs, target)\n\n        kl = get_total_kl(model)\n        bound = self.compute_bound(nll, kl, n_bound, delta)\n        return bound\n\n    def compute_risk_certificate(self, model, bound_loader, device, delta=0.025,\n                                 mc_samples=1000):\n        \"\"\"Evaluate final PAC-Bayes risk certificate after training.\n\n        Computes:\n        1. Empirical 0-1 risk via MC sampling on the bound-evaluation set\n        2. KL divergence between posterior and prior\n        3. PAC-Bayes-kl bound inversion for the final certificate\n\n        Args:\n            model: trained stochastic model\n            bound_loader: DataLoader for bound-evaluation set\n            device: torch device\n            delta: confidence parameter\n            mc_samples: number of MC samples per input\n\n        Returns:\n            (risk_cert_01, metrics_dict)\n        \"\"\"\n        model.eval()\n        n_bound = len(bound_loader.dataset)\n\n        # 1. Compute empirical 0-1 risk via MC sampling\n        emp_risk_01 = compute_01_risk(model, bound_loader, device,\n                                      mc_samples=mc_samples)\n\n        # 2. Compute NLL-based empirical risk for the CE bound\n        total_nll = 0.0\n        total_samples = 0\n        kl_total = None\n        with torch.no_grad():\n            for data, target in bound_loader:\n                data, target = data.to(device), target.to(device)\n                output = model(data, sample=True)\n                log_probs = F.log_softmax(output, dim=1)\n                log_probs = torch.clamp(log_probs, min=math.log(self.pmin))\n                nll = F.nll_loss(log_probs, target, reduction=\"sum\")\n                total_nll += nll.item()\n                total_samples += target.size(0)\n                if kl_total is None:\n                    kl_total = get_total_kl(model)\n\n        emp_nll = total_nll / total_samples\n\n        # 3. Get KL from a single forward pass\n        with torch.no_grad():\n            dummy_data = next(iter(bound_loader))[0][:1].to(device)\n            model(dummy_data, sample=True)\n            kl = get_total_kl(model).item()\n\n        # 4. PAC-Bayes-kl bound inversion for 0-1 loss certificate\n        c = (kl + math.log(2.0 * math.sqrt(n_bound) / delta)) / n_bound\n        risk_cert_01 = inv_kl(emp_risk_01, c)\n\n        # 5. Compute the direct bound from compute_bound for CE risk\n        emp_nll_t = torch.tensor(emp_nll)\n        kl_t = torch.tensor(kl)\n        ce_bound = self.compute_bound(emp_nll_t, kl_t, n_bound, delta).item()\n\n        metrics = {\n            \"empirical_01_risk\": emp_risk_01,\n            \"empirical_nll\": emp_nll,\n            \"kl_divergence\": kl,\n            \"ce_bound\": ce_bound,\n        }\n\n        return risk_cert_01, metrics"
7}

1{
2  "n": 1
3}

1{
2  "op": "replace",
3  "filename": "PBB/custom_pac_bayes.py",
4  "start_line": 460,
5  "end_line": 604,
6  "content": "class BoundOptimizer:\n    \"\"\"PAC-Bayes bound with adaptive union of multiple bound forms.\n\n    This optimizer computes three PAC-Bayes bounds (McAllester, Catoni, Quadratic)\n    each with confidence delta/3, then takes the minimum as the final bound\n    (union bound). This yields a tighter certificate than any single bound.\n\n    Additionally, the Catoni lambda parameter is optimized per batch to minimize\n    the Catoni bound, providing adaptive trade-off between empirical risk and KL.\n\n    The training objective minimizes the union bound with NLL surrogate.\n    The final certificate uses PAC-Bayes-kl inversion (already tight).\n    \"\"\"\n\n    def __init__(self, learning_rate=0.001, momentum=0.95, prior_sigma=0.1,\n                 pmin=1e-5):\n        self.learning_rate = learning_rate\n        self.momentum = momentum\n        self.prior_sigma = prior_sigma\n        self.pmin = pmin\n\n    def compute_mcallester_bound(self, empirical_risk, kl, n, delta):\n        \"\"\"McAllester/Maurer bound.\"\"\"\n        kl_term = (kl + math.log(2.0 * math.sqrt(n) / delta)) / (2.0 * n)\n        bound = empirical_risk + torch.sqrt(kl_term)\n        return bound\n\n    def compute_catoni_bound(self, empirical_risk, kl, n, delta, lam):\n        \"\"\"Catoni/Lambda bound with given lambda.\"\"\"\n        # Clamp lambda to (0, 2) for valid bound\n        lam = torch.clamp(lam, min=0.01, max=1.99)\n        kl_term = (kl + math.log(2.0 * math.sqrt(n) / delta)) / (\n            n * lam * (1.0 - lam / 2.0)\n        )\n        bound = empirical_risk / (1.0 - lam / 2.0) + kl_term\n        return bound\n\n    def compute_quadratic_bound(self, empirical_risk, kl, n, delta):\n        \"\"\"Quadratic PAC-Bayes bound.\"\"\"\n        kl_term = (kl + math.log(2.0 * math.sqrt(n) / delta)) / (2.0 * n)\n        # Ensure non-negative under sqrt\n        inner = torch.clamp(empirical_risk + kl_term, min=0.0)\n        kl_term_clamped = torch.clamp(kl_term, min=0.0)\n        bound = (torch.sqrt(inner) + torch.sqrt(kl_term_clamped)) ** 2\n        return bound\n\n    def compute_bound(self, empirical_risk, kl, n, delta):\n        \"\"\"Union bound: minimum of three bounds each with delta/3.\"\"\"\n        delta_div = delta / 3.0\n        bound_mc = self.compute_mcallester_bound(empirical_risk, kl, n, delta_div)\n        # For Catoni bound, choose optimal lambda that minimizes the bound.\n        # Optimal lambda (theoretically) is sqrt(2*(KL+log(...))/(n*empirical_risk))\n        # but we compute via analytic minimization.\n        # We'll compute lambda that minimizes the Catoni bound given current parameters.\n        # Derivation: minimize f(λ) = empirical_risk/(1-λ/2) + (KL+log(...))/(n*λ*(1-λ/2))\n        # We'll solve via closed-form approximation: λ_opt = sqrt(2*(KL+log(...))/(n*empirical_risk))\n        # Clamp to (0,2).\n        with torch.no_grad():\n            kl_plus_log = kl + math.log(2.0 * math.sqrt(n) / delta_div)\n            if empirical_risk > 0:\n                lam_opt = torch.sqrt(2.0 * kl_plus_log / (n * empirical_risk))\n                lam_opt = torch.clamp(lam_opt, min=0.01, max=1.99)\n            else:\n                lam_opt = torch.tensor(1.99)\n        bound_cat = self.compute_catoni_bound(empirical_risk, kl, n, delta_div, lam_opt)\n        bound_quad = self.compute_quadratic_bound(empirical_risk, kl, n, delta_div)\n        # Take minimum of three bounds\n        bound = torch.min(bound_mc, torch.min(bound_cat, bound_quad))\n        return bound\n\n    def train_step(self, model, data, target, device, n_bound, delta):\n        \"\"\"Training objective: minimize union bound with NLL surrogate.\"\"\"\n        output = model(data, sample=True)\n        log_probs = F.log_softmax(output, dim=1)\n        log_probs = torch.clamp(log_probs, min=math.log(self.pmin))\n        nll = F.nll_loss(log_probs, target)\n\n        kl = get_total_kl(model)\n        bound = self.compute_bound(nll, kl, n_bound, delta)\n        return bound\n\n    def compute_risk_certificate(self, model, bound_loader, device, delta=0.025,\n                                 mc_samples=1000):\n        \"\"\"Evaluate final PAC-Bayes risk certificate after training.\n\n        Uses PAC-Bayes-kl inversion for 0-1 loss (tightest). Also computes\n        the union bound for NLL risk for reference.\n        \"\"\"\n        model.eval()\n        n_bound = len(bound_loader.dataset)\n\n        # 1. Empirical 0-1 risk via MC sampling\n        emp_risk_01 = compute_01_risk(model, bound_loader, device,\n                                      mc_samples=mc_samples)\n\n        # 2. NLL-based empirical risk\n        total_nll = 0.0\n        total_samples = 0\n        with torch.no_grad():\n            for data, target in bound_loader:\n                data, target = data.to(device), target.to(device)\n                output = model(data, sample=True)\n                log_probs = F.log_softmax(output, dim=1)\n                log_probs = torch.clamp(log_probs, min=math.log(self.pmin))\n                nll = F.nll_loss(log_probs, target, reduction=\"sum\")\n                total_nll += nll.item()\n                total_samples += target.size(0)\n        emp_nll = total_nll / total_samples\n\n        # 3. KL divergence\n        with torch.no_grad():\n            dummy_data = next(iter(bound_loader))[0][:1].to(device)\n            model(dummy_data, sample=True)\n            kl = get_total_kl(model).item()\n\n        # 4. PAC-Bayes-kl inversion for 0-1 loss certificate (tightest)\n        c = (kl + math.log(2.0 * math.sqrt(n_bound) / delta)) / n_bound\n        risk_cert_01 = inv_kl(emp_risk_01, c)\n\n        # 5. Compute union bound for NLL risk (for reference)\n        emp_nll_t = torch.tensor(emp_nll)\n        kl_t = torch.tensor(kl)\n        ce_bound = self.compute_bound(emp_nll_t, kl_t, n_bound, delta).item()\n\n        # 6. Also compute individual bounds for debugging\n        delta_div = delta / 3.0\n        bound_mc = self.compute_mcallester_bound(emp_nll_t, kl_t, n_bound, delta_div).item()\n        # Compute optimal lambda for Catoni bound\n        kl_plus_log = kl + math.log(2.0 * math.sqrt(n_bound) / delta_div)\n        if emp_nll > 0:\n            lam_opt = math.sqrt(2.0 * kl_plus_log / (n_bound * emp_nll))\n            lam_opt = max(0.01, min(1.99, lam_opt))\n        else:\n            lam_opt = 1.99\n        bound_cat = self.compute_catoni_bound(emp_nll_t, kl_t, n_bound, delta_div,\n                                              torch.tensor(lam_opt)).item()\n        bound_quad = self.compute_quadratic_bound(emp_nll_t, kl_t, n_bound, delta_div).item()\n\n        metrics = {\n            \"empirical_01_risk\": emp_risk_01,\n            \"empirical_nll\": emp_nll,\n            \"kl_divergence\": kl,\n            \"ce_bound\": ce_bound,\n            \"ce_bound_mc\": bound_mc,\n            \"ce_bound_cat\": bound_cat,\n            \"ce_bound_quad\": bound_quad,\n            \"catoni_lambda\": lam_opt,\n        }\n\n        return risk_cert_01, metrics"
7}

1{
2  "op": "replace",
3  "filename": "PBB/custom_pac_bayes.py",
4  "start_line": 460,
5  "end_line": 604,
6  "content": "class BoundOptimizer:\n    \"\"\"Adaptive Curvature PAC-Bayes Bound.\n\n    Novel bound: B = empirical_risk + sqrt( (KL + log(2√n/δ)) / (2n) ) * (1 + β * empirical_risk)\n    where β is a learnable curvature parameter that adapts the bound tightness based on empirical risk.\n    \n    Motivation: When empirical risk is low, the bound should behave like McAllester (β=0).\n    When empirical risk is higher, the KL term is scaled up to maintain validity while allowing\n    more flexibility in trading off risk vs complexity. β is optimized jointly with the posterior.\n    \n    The bound is derived from the variational formulation of PAC-Bayes with a parameterized\n    convex function that interpolates between linear and quadratic penalties.\n    \n    For final certificate, we use the tight PAC-Bayes-kl inversion.\n    \"\"\"\n\n    def __init__(self, learning_rate=0.001, momentum=0.95, prior_sigma=0.1,\n                 pmin=1e-5, beta_init=0.5, beta_lr=0.01):\n        self.learning_rate = learning_rate\n        self.momentum = momentum\n        self.prior_sigma = prior_sigma\n        self.pmin = pmin\n        # Learnable curvature parameter β (non-negative)\n        self._beta_param = torch.tensor(beta_init, requires_grad=True)\n        self.beta_lr = beta_lr\n        self._beta_optimizer = None\n\n    def _get_beta(self):\n        \"\"\"Get clamped beta value (non-negative).\"\"\"\n        return torch.clamp(self._beta_param, min=0.0, max=10.0)\n\n    def _ensure_beta_optimizer(self):\n        if self._beta_optimizer is None:\n            self._beta_optimizer = torch.optim.SGD(\n                [self._beta_param], lr=self.beta_lr\n            )\n\n    def compute_bound(self, empirical_risk, kl, n, delta):\n        \"\"\"Adaptive curvature bound.\"\"\"\n        kl_term = (kl + math.log(2.0 * math.sqrt(n) / delta)) / (2.0 * n)\n        sqrt_kl_term = torch.sqrt(torch.clamp(kl_term, min=0.0))\n        beta = self._get_beta()\n        bound = empirical_risk + sqrt_kl_term * (1.0 + beta * empirical_risk)\n        return bound\n\n    def train_step(self, model, data, target, device, n_bound, delta):\n        \"\"\"Training objective: adaptive curvature bound with NLL surrogate.\"\"\"\n        # Ensure beta is on correct device\n        if self._beta_param.device != device:\n            self._beta_param = self._beta_param.to(device).detach().requires_grad_(True)\n            self._beta_optimizer = None\n        self._ensure_beta_optimizer()\n\n        output = model(data, sample=True)\n        # Bounded cross-entropy as surrogate for 0-1 loss\n        log_probs = F.log_softmax(output, dim=1)\n        log_probs = torch.clamp(log_probs, min=math.log(self.pmin))\n        nll = F.nll_loss(log_probs, target)\n\n        kl = get_total_kl(model)\n        bound = self.compute_bound(nll, kl, n_bound, delta)\n\n        # Also update beta\n        self._beta_optimizer.zero_grad()\n        # We need to retain graph because bound will be used for model optimization\n        bound.backward(retain_graph=True)\n        self._beta_optimizer.step()\n\n        return bound\n\n    def compute_risk_certificate(self, model, bound_loader, device, delta=0.025,\n                                 mc_samples=1000):\n        \"\"\"Evaluate final PAC-Bayes risk certificate after training.\n        \n        Uses PAC-Bayes-kl inversion for tightest certificate.\n        \"\"\"\n        model.eval()\n        n_bound = len(bound_loader.dataset)\n\n        # 1. Compute empirical 0-1 risk via MC sampling\n        emp_risk_01 = compute_01_risk(model, bound_loader, device,\n                                      mc_samples=mc_samples)\n\n        # 2. Compute NLL-based empirical risk for the CE bound\n        total_nll = 0.0\n        total_samples = 0\n        with torch.no_grad():\n            for data, target in bound_loader:\n                data, target = data.to(device), target.to(device)\n                output = model(data, sample=True)\n                log_probs = F.log_softmax(output, dim=1)\n                log_probs = torch.clamp(log_probs, min=math.log(self.pmin))\n                nll = F.nll_loss(log_probs, target, reduction=\"sum\")\n                total_nll += nll.item()\n                total_samples += target.size(0)\n\n        emp_nll = total_nll / total_samples\n\n        # 3. Get KL from a single forward pass\n        with torch.no_grad():\n            dummy_data = next(iter(bound_loader))[0][:1].to(device)\n            model(dummy_data, sample=True)\n            kl = get_total_kl(model).item()\n\n        # 4. PAC-Bayes-kl bound inversion for 0-1 loss certificate\n        c = (kl + math.log(2.0 * math.sqrt(n_bound) / delta)) / n_bound\n        risk_cert_01 = inv_kl(emp_risk_01, c)\n\n        # 5. Compute the direct bound from compute_bound for CE risk\n        emp_nll_t = torch.tensor(emp_nll)\n        kl_t = torch.tensor(kl)\n        ce_bound = self.compute_bound(emp_nll_t, kl_t, n_bound, delta).item()\n\n        metrics = {\n            \"empirical_01_risk\": emp_risk_01,\n            \"empirical_nll\": emp_nll,\n            \"kl_divergence\": kl,\n            \"ce_bound\": ce_bound,\n            \"beta\": self._get_beta().item(),\n        }\n\n        return risk_cert_01, metrics"
7}

1{
2  "op": "replace",
3  "filename": "PBB/custom_pac_bayes.py",
4  "start_line": 505,
5  "end_line": 528,
6  "content": "    def train_step(self, model, data, target, device, n_bound, delta):\n        \"\"\"Training objective: adaptive curvature bound with NLL surrogate.\"\"\"\n        # Ensure beta is on correct device\n        if self._beta_param.device != device:\n            self._beta_param = self._beta_param.to(device).detach().requires_grad_(True)\n            self._beta_optimizer = None\n        self._ensure_beta_optimizer()\n\n        output = model(data, sample=True)\n        # Bounded cross-entropy as surrogate for 0-1 loss\n        log_probs = F.log_softmax(output, dim=1)\n        log_probs = torch.clamp(log_probs, min=math.log(self.pmin))\n        nll = F.nll_loss(log_probs, target)\n\n        kl = get_total_kl(model)\n        bound = self.compute_bound(nll, kl, n_bound, delta)\n\n        # Update beta parameter using gradient of bound w.r.t beta\n        self._beta_optimizer.zero_grad()\n        # Compute gradient of bound with respect to beta only\n        beta_grad = torch.autograd.grad(bound, self._beta_param, retain_graph=True)[0]\n        if beta_grad is not None:\n            self._beta_param.grad = beta_grad\n            self._beta_optimizer.step()\n\n        return bound"
7}

1{
2  "op": "replace",
3  "filename": "PBB/custom_pac_bayes.py",
4  "start_line": 460,
5  "end_line": 583,
6  "content": "class BoundOptimizer:\n    \"\"\"Adaptive Curvature PAC-Bayes Bound.\n\n    Novel bound: B = empirical_risk + sqrt( (KL + log(2√n/δ)) / (2n) ) * (1 + β * empirical_risk)\n    where β is a learnable curvature parameter that adapts the bound tightness based on empirical risk.\n    \n    Motivation: When empirical risk is low, the bound should behave like McAllester (β=0).\n    When empirical risk is higher, the KL term is scaled up to maintain validity while allowing\n    more flexibility in trading off risk vs complexity. β is optimized jointly with the posterior.\n    \n    The bound is derived from the variational formulation of PAC-Bayes with a parameterized\n    convex function that interpolates between linear and quadratic penalties.\n    \n    For final certificate, we use the tight PAC-Bayes-kl inversion.\n    \"\"\"\n\n    def __init__(self, learning_rate=0.001, momentum=0.95, prior_sigma=0.1,\n                 pmin=1e-5, beta_init=0.5, beta_lr=0.01):\n        self.learning_rate = learning_rate\n        self.momentum = momentum\n        self.prior_sigma = prior_sigma\n        self.pmin = pmin\n        # Learnable curvature parameter β (non-negative)\n        self.beta = torch.tensor(beta_init)\n        self.beta_lr = beta_lr\n        self.beta_momentum = 0.9\n        self.beta_velocity = 0.0\n\n    def compute_bound(self, empirical_risk, kl, n, delta):\n        \"\"\"Adaptive curvature bound.\"\"\"\n        kl_term = (kl + math.log(2.0 * math.sqrt(n) / delta)) / (2.0 * n)\n        sqrt_kl_term = torch.sqrt(torch.clamp(kl_term, min=0.0))\n        beta = torch.clamp(self.beta, min=0.0, max=10.0)\n        bound = empirical_risk + sqrt_kl_term * (1.0 + beta * empirical_risk)\n        return bound\n\n    def train_step(self, model, data, target, device, n_bound, delta):\n        \"\"\"Training objective: adaptive curvature bound with NLL surrogate.\"\"\"\n        # Ensure beta is on correct device\n        if self.beta.device != device:\n            self.beta = self.beta.to(device)\n\n        output = model(data, sample=True)\n        # Bounded cross-entropy as surrogate for 0-1 loss\n        log_probs = F.log_softmax(output, dim=1)\n        log_probs = torch.clamp(log_probs, min=math.log(self.pmin))\n        nll = F.nll_loss(log_probs, target)\n\n        kl = get_total_kl(model)\n        bound = self.compute_bound(nll, kl, n_bound, delta)\n\n        # Update beta using gradient of bound w.r.t beta (computed analytically)\n        # ∂bound/∂β = sqrt_kl_term * empirical_risk\n        with torch.no_grad():\n            kl_term = (kl + math.log(2.0 * math.sqrt(n_bound) / delta)) / (2.0 * n_bound)\n            sqrt_kl_term = torch.sqrt(torch.clamp(kl_term, min=0.0)).item()\n            beta_grad = sqrt_kl_term * nll.item()\n            # SGD with momentum\n            self.beta_velocity = self.beta_momentum * self.beta_velocity - self.beta_lr * beta_grad\n            self.beta += self.beta_velocity\n            self.beta = torch.clamp(self.beta, min=0.0, max=10.0)\n            # Optional: print beta occasionally\n            # if torch.rand(1).item() < 0.01:\n            #     print(f\"BETA_UPDATE beta={self.beta.item():.3f} grad={beta_grad:.6f}\", flush=True)\n\n        return bound\n\n    def compute_risk_certificate(self, model, bound_loader, device, delta=0.025,\n                                 mc_samples=1000):\n        \"\"\"Evaluate final PAC-Bayes risk certificate after training.\n        \n        Uses PAC-Bayes-kl inversion for tightest certificate.\n        \"\"\"\n        model.eval()\n        n_bound = len(bound_loader.dataset)\n\n        # 1. Compute empirical 0-1 risk via MC sampling\n        emp_risk_01 = compute_01_risk(model, bound_loader, device,\n                                      mc_samples=mc_samples)\n\n        # 2. Compute NLL-based empirical risk for the CE bound\n        total_nll = 0.0\n        total_samples = 0\n        with torch.no_grad():\n            for data, target in bound_loader:\n                data, target = data.to(device), target.to(device)\n                output = model(data, sample=True)\n                log_probs = F.log_softmax(output, dim=1)\n                log_probs = torch.clamp(log_probs, min=math.log(self.pmin))\n                nll = F.nll_loss(log_probs, target, reduction=\"sum\")\n                total_nll += nll.item()\n                total_samples += target.size(0)\n\n        emp_nll = total_nll / total_samples\n\n        # 3. Get KL from a single forward pass\n        with torch.no_grad():\n            dummy_data = next(iter(bound_loader))[0][:1].to(device)\n            model(dummy_data, sample=True)\n            kl = get_total_kl(model).item()\n\n        # 4. PAC-Bayes-kl bound inversion for 0-1 loss certificate\n        c = (kl + math.log(2.0 * math.sqrt(n_bound) / delta)) / n_bound\n        risk_cert_01 = inv_kl(emp_risk_01, c)\n\n        # 5. Compute the direct bound from compute_bound for CE risk\n        emp_nll_t = torch.tensor(emp_nll)\n        kl_t = torch.tensor(kl)\n        ce_bound = self.compute_bound(emp_nll_t, kl_t, n_bound, delta).item()\n\n        metrics = {\n            \"empirical_01_risk\": emp_risk_01,\n            \"empirical_nll\": emp_nll,\n            \"kl_divergence\": kl,\n            \"ce_bound\": ce_bound,\n            \"beta\": torch.clamp(self.beta, min=0.0, max=10.0).item(),\n        }\n\n        return risk_cert_01, metrics"
7}