1"""Late-time fractional-order optimizer MVP.
 2
 3The fractional direction is a normalized piecewise-constant Caputo-memory
 4average of recent gradients. Normalization keeps the learning-rate meaning
 5comparable to Adam/SGD while retaining the intended power-law kernel shape.
 6"""
 7from collections import deque
 8import math
 9import torch
10
11
12class LateTimeOrderEstimator:
13    def __init__(self, alpha0=0.5, amin=0.2, amax=0.95, beta=0.8,
14                 warmup=40, lag=20, check_every=5, monotone_checks=2,
15                 eps=1e-12):
16        self.alpha = float(alpha0)
17        self.amin, self.amax = amin, amax
18        self.beta, self.warmup, self.lag = beta, warmup, lag
19        self.check_every, self.monotone_checks, self.eps = check_every, monotone_checks, eps
20        self.signals = []
21        self.ratios = deque(maxlen=monotone_checks)
22
23    def update(self, step, loss):
24        # A positive loss is the scalar observation M. EMA suppresses minibatch noise.
25        loss = max(float(loss), self.eps)
26        if self.signals:
27            self.signals.append(self.beta * self.signals[-1] + (1-self.beta) * loss)
28        else:
29            self.signals.append(loss)
30        if step < self.warmup or step % self.check_every:
31            return self.alpha
32        old_i = max(0, len(self.signals)-1-self.lag)
33        if old_i == len(self.signals)-1:
34            return self.alpha
35        ratio = self.signals[-1] / max(self.signals[old_i], self.eps)
36        self.ratios.append(ratio)
37        # The late-time hypothesis requires decreasing positive observations.
38        if not (0.0 < ratio < 1.0) or len(self.ratios) >= 2 and any(
39                self.ratios[i] >= self.ratios[i-1] for i in range(1, len(self.ratios))):
40            return self.alpha
41        rho = (step + 1) / float(step + 1 - self.lag)
42        raw = -math.log(ratio) / math.log(rho)
43        raw = min(self.amax, max(self.amin, raw))
44        self.alpha = self.beta * self.alpha + (1-self.beta) * raw
45        return self.alpha
46
47
48class FractionalMemory(torch.optim.Optimizer):
49    def __init__(self, params, lr=0.05, alpha=0.5, history=32, weight_decay=0.0,
50                 adaptive=False, estimator=None):
51        defaults = dict(lr=lr, alpha=alpha, history=history, weight_decay=weight_decay)
52        super().__init__(params, defaults)
53        self.adaptive = adaptive
54        self.estimator = estimator or LateTimeOrderEstimator(alpha0=alpha)
55        self.step_count = 0
56
57    @torch.no_grad()
58    def step(self, closure=None, loss_value=None):
59        if closure is not None:
60            with torch.enable_grad():
61                loss = closure()
62            loss_value = float(loss)
63        self.step_count += 1
64        alpha = self.estimator.update(self.step_count, loss_value) if self.adaptive and loss_value is not None else None
65        for group in self.param_groups:
66            a = alpha if alpha is not None else group['alpha']
67            a = min(0.999, max(0.05, float(a)))
68            H = group['history']
69            # normalized interval weights: w_j=(j+1)^a-j^a, recent gradient j=0 largest
70            weights = torch.tensor([(j+1)**a-j**a for j in range(H)], device=group['params'][0].device)
71            weights = weights / weights.sum()
72            for p in group['params']:
73                if p.grad is None:
74                    continue
75                g = p.grad.detach().clone()
76                state = self.state[p]
77                hist = state.setdefault('grad_history', deque(maxlen=H))
78                hist.appendleft(g)
79                n = len(hist)
80                v = sum(weights[j] * hist[j] for j in range(n))
81                if group['weight_decay']:
82                    p.mul_(1 - group['lr'] * group['weight_decay'])
83                p.add_(v, alpha=-group['lr'])
84        return None
85
86
87class AdamW(torch.optim.AdamW):
88    """Named wrapper so the experiment has an explicit standard baseline."""
89    pass