1"""Third-order Langevin optimizer MVP.
 2
 3State: parameter x, velocity v, acceleration a. Noise is injected only into a:
 4 a <- a + dt*(-grad - gamma*a) + sqrt(2*gamma*temperature*dt)*N(0,I)
 5 v <- v + dt*a
 6 x <- x + dt*v
 7"""
 8import math
 9import torch
10from torch.optim import Optimizer
11
12
13def cubic_unstable_rate(gamma: float, kappa: float) -> float:
14    """Positive root r of r^3 + gamma*r^2 - kappa = 0."""
15    if gamma <= 0:
16        raise ValueError("gamma must be positive")
17    if kappa <= 0:
18        return 0.0
19    lo, hi = 0.0, kappa ** (1.0 / 3.0) + math.sqrt(kappa / gamma) + 1.0
20    for _ in range(80):
21        mid = (lo + hi) * 0.5
22        if mid**3 + gamma * mid**2 < kappa:
23            lo = mid
24        else:
25            hi = mid
26    return (lo + hi) * 0.5
27
28
29class ThirdOrderLangevin(Optimizer):
30    """Noise-on-acceleration third-order Langevin optimizer.
31
32    ``adapt_dt=True`` optionally enforces dt*r <= rate_limit using a supplied
33    negative-curvature estimate via ``set_negative_curvature``. Curvature
34    estimation is deliberately kept outside the optimizer so callers can use
35    Hessian-vector products at a chosen cadence.
36    """
37    def __init__(self, params, dt=1e-2, gamma=1.0, temperature=0.0,
38                 adapt_dt=False, rate_limit=0.2):
39        if dt <= 0 or gamma <= 0 or temperature < 0 or rate_limit <= 0:
40            raise ValueError("invalid dt, gamma, temperature, or rate_limit")
41        defaults = dict(dt=float(dt), gamma=float(gamma),
42                        temperature=float(temperature), adapt_dt=adapt_dt,
43                        rate_limit=float(rate_limit))
44        super().__init__(params, defaults)
45        self._kappa = 0.0
46
47    @torch.no_grad()
48    def set_negative_curvature(self, kappa: float):
49        self._kappa = max(0.0, float(kappa))
50
51    @torch.no_grad()
52    def step(self, closure=None):
53        loss = None
54        if closure is not None:
55            with torch.enable_grad():
56                loss = closure()
57        # CUDA errors are allowed to propagate to the caller, which can rerun
58        # the experiment on CPU as required by the experiment harness.
59        for group in self.param_groups:
60            dt = group['dt']
61            gamma = group['gamma']
62            if group['adapt_dt'] and self._kappa > 0:
63                rate = cubic_unstable_rate(gamma, self._kappa)
64                dt = min(dt, group['rate_limit'] / rate)
65            noise_scale = math.sqrt(2.0 * gamma * group['temperature'] * dt)
66            for p in group['params']:
67                if p.grad is None:
68                    continue
69                if p.grad.is_sparse:
70                    raise RuntimeError("ThirdOrderLangevin does not support sparse gradients")
71                state = self.state[p]
72                if not state:
73                    state['v'] = torch.zeros_like(p, memory_format=torch.preserve_format)
74                    state['a'] = torch.zeros_like(p, memory_format=torch.preserve_format)
75                # Keep optimizer state fp32 where parameters are lower precision.
76                v, a = state['v'], state['a']
77                g = p.grad
78                if not torch.is_floating_point(g):
79                    g = g.float()
80                a.add_(-dt * g).add_(-dt * gamma * a)
81                if noise_scale:
82                    a.add_(torch.randn_like(a) * noise_scale)
83                v.add_(a, alpha=dt)
84                p.add_(v, alpha=dt)
85        return loss