import json import math import random from pathlib import Path import numpy as np import torch SEED = 474 random.seed(SEED) np.random.seed(SEED) torch.manual_seed(SEED) def device_or_cpu(): if torch.cuda.is_available(): try: torch.zeros(1, device="cuda") return "cuda" except Exception: pass return "cpu" class MonotoneSplineGate(torch.nn.Module): """delta(x)=s(|x|)x, with increasing knot values in [0,1].""" def __init__(self, knots, init_identity=False): super().__init__() self.register_buffer("knots", torch.tensor(knots, dtype=torch.float32)) k = len(knots) # Positive increments in unconstrained space guarantee monotonicity. if init_identity: self.bias = torch.nn.Parameter(torch.tensor(5.0)) self.raw_inc = torch.nn.Parameter(torch.full((k - 1,), -5.0)) else: self.bias = torch.nn.Parameter(torch.tensor(-1.0)) self.raw_inc = torch.nn.Parameter(torch.full((k - 1,), -1.0)) def knot_values(self): # sigmoid(bias + cumulative positive increments), hence monotone. increments = torch.nn.functional.softplus(self.raw_inc) logits = self.bias + torch.cat([torch.zeros(1, device=increments.device), torch.cumsum(increments, dim=0)]) return torch.sigmoid(logits) def spline(self, magnitude): k = self.knots.to(magnitude.device) vals = self.knot_values() # endpoint extrapolation is constant, as required for a bounded gate. idx = torch.bucketize(magnitude.detach().reshape(-1), k[1:-1]).reshape(magnitude.shape) idx = idx.clamp(0, len(k) - 2) lo, hi = k[idx], k[idx + 1] frac = ((magnitude - lo) / (hi - lo).clamp_min(1e-6)).clamp(0, 1) out = vals[idx] * (1 - frac) + vals[idx + 1] * frac out = torch.where(magnitude >= k[-1], vals[-1], out) return out def slope(self, magnitude): k = self.knots.to(magnitude.device) vals = self.knot_values() idx = torch.bucketize(magnitude.detach().reshape(-1), k[1:-1]).reshape(magnitude.shape) idx = idx.clamp(0, len(k) - 2) out = (vals[idx + 1] - vals[idx]) / (k[idx + 1] - k[idx]).clamp_min(1e-6) return torch.where(magnitude >= k[-1], torch.zeros_like(out), out) def forward(self, x): return self.spline(x.abs()) * x def divergence(self, x): m = x.abs() return (self.spline(m) + m * self.slope(m)).sum(dim=-1) def sure(self, x): d = x.shape[-1] return ((self(x) - x).square().sum(dim=-1) + 2 * self.divergence(x) - d).mean() def finite_difference_divergence(gate, x, eps=1e-4): # Explicit Jacobian trace for a small tensor, used only as a math check. vals = [] for j in range(x.numel()): xp, xm = x.clone(), x.clone() xp.reshape(-1)[j] += eps xm.reshape(-1)[j] -= eps vals.append(((gate(xp) - gate(xm)).reshape(-1)[j] / (2 * eps)).item()) return sum(vals) def fit_gate(x, steps=250): gate = MonotoneSplineGate(np.linspace(0, 4, 9)).to(x.device) opt = torch.optim.Adam(gate.parameters(), lr=0.04) for step in range(steps): opt.zero_grad() # Mild high-magnitude anchor prevents the unconstrained solution from # attenuating every coordinate when the batch is nearly all noise. loss = gate.sure(x) + 0.01 * (1 - gate.knot_values()[-1]).square() loss.backward() torch.nn.utils.clip_grad_norm_(gate.parameters(), 5.0) opt.step() return gate def fixed_soft(x, threshold=1.0): return torch.sign(x) * torch.relu(x.abs() - threshold) def risk(est, theta): return (est - theta).square().mean().item() def run(device): # Core calculus check: analytic divergence agrees with finite differences. check_gate = MonotoneSplineGate(np.linspace(0, 4, 9)).to(device) check_x = torch.tensor([[0.31, 1.17, 2.63, 3.71]], device=device) div_formula = check_gate.divergence(check_x).item() div_fd = finite_difference_divergence(check_gate, check_x).item() if False else finite_difference_divergence(check_gate, check_x) # SURE unbiasedness check for a fixed smooth spline under Z=theta+N(0,I). theta = torch.tensor([0.0, 0.5, 1.5, 3.0], device=device).repeat(256, 1) z = theta + torch.randn_like(theta) sure = ((check_gate(z) - z).square().sum(1) + 2 * check_gate.divergence(z) - z.shape[1]).mean().item() / z.shape[1] actual = risk(check_gate(z), theta) # Repeated draws check the expected SURE identity, not just one batch. sure_rep, risk_rep = [], [] for _ in range(20): zz = theta + torch.randn_like(theta) sure_rep.append((((check_gate(zz) - zz).square().sum(1) + 2 * check_gate.divergence(zz) - zz.shape[1]).mean() / zz.shape[1]).item()) risk_rep.append(risk(check_gate(zz), theta)) # Adaptive normal-mean experiment: sparse and dense regimes, held-out risk. results = {} for name, scale in [("sparse", 0.35), ("dense", 2.0)]: train_theta = scale * torch.randn(64, 128, device=device) train_z = train_theta + torch.randn_like(train_theta) gate = fit_gate(train_z) test_theta = scale * torch.randn(256, 128, device=device) test_z = test_theta + torch.randn_like(test_theta) with torch.no_grad(): r_gate = risk(gate(test_z), test_theta) r_id = risk(test_z, test_theta) r_soft = risk(fixed_soft(test_z), test_theta) train_sure = gate.sure(train_z).item() / train_z.shape[-1] vals = gate.knot_values().detach().cpu().numpy().round(3).tolist() results[name] = {"identity": r_id, "fixed_soft": r_soft, "risk_fitted_gate": r_gate, "train_sure_per_dim": train_sure, "knot_values": vals} return {"device": device, "divergence_formula": div_formula, "divergence_finite_difference": div_fd, "fixed_gate_sure": sure, "fixed_gate_actual_risk": actual, "fixed_gate_sure_repeated_mean": float(np.mean(sure_rep)), "fixed_gate_actual_risk_repeated_mean": float(np.mean(risk_rep)), "fixed_gate_sure_minus_risk_repeated": float(np.mean(sure_rep) - np.mean(risk_rep)), "regimes": results} if __name__ == "__main__": out = run(device_or_cpu()) print(json.dumps(out, indent=2))