1"""Heavy-tail path-adaptive optimizer pool: small NumPy reference implementation."""
  2import numpy as np
  3
  4class RestartedAdaGrad:
  5    def __init__(self, dim, horizon, alpha=0.25, epsilon=1e-8, x0=None, radius=10.0):
  6        self.dim, self.horizon = dim, horizon
  7        self.alpha, self.epsilon, self.radius = alpha, epsilon, radius
  8        self.x = np.zeros(dim) if x0 is None else np.array(x0, dtype=float).copy()
  9        self.age = 0
 10        self.acc = np.full(dim, epsilon, dtype=float)
 11    def step(self, grad):
 12        grad = np.asarray(grad, dtype=float)
 13        # B_k(t) is the current block; reset the diagonal accumulator at restart.
 14        if self.age >= self.horizon:
 15            self.age, self.acc = 0, np.full(self.dim, self.epsilon)
 16        self.acc += grad * grad
 17        self.x -= self.alpha * grad / np.sqrt(self.acc)
 18        self.x = np.clip(self.x, -self.radius, self.radius)
 19        self.age += 1
 20        return self.x
 21
 22class HeavyTailPool:
 23    def __init__(self, dim=1, horizons=(1, 2, 4, 8, 16), alpha=.25,
 24                 beta=2.0, epsilon=1e-8, x0=None, radius=10.):
 25        self.horizons = tuple(horizons)
 26        self.experts = [RestartedAdaGrad(dim, h, alpha, epsilon, x0, radius)
 27                        for h in self.horizons]
 28        self.q = np.ones(len(horizons)) / len(horizons)
 29        self.v = np.zeros(len(horizons))
 30        self.beta, self.epsilon = beta, epsilon
 31    def step(self, grad_fn, validation_loss_fn):
 32        # Each expert receives its own gradient, as in the proposed construction.
 33        for e in self.experts:
 34            e.step(grad_fn(e.x))
 35        losses = np.array([validation_loss_fn(e.x) for e in self.experts])
 36        mean_loss = float(self.q @ losses)
 37        z = losses - mean_loss
 38        self.v += z*z
 39        logits = np.log(self.q) - self.beta*z/np.sqrt(self.epsilon + self.v)
 40        logits -= logits.max()
 41        self.q = np.exp(logits); self.q /= self.q.sum()
 42        x = sum(q*e.x for q, e in zip(self.q, self.experts))
 43        return x, losses, self.q.copy()
 44
 45def fixed_adagrad(x0, gradients, alpha=.25, epsilon=1e-8, radius=10.):
 46    x = np.array(x0, dtype=float).copy(); acc = np.full_like(x, epsilon)
 47    out=[]
 48    for g in gradients:
 49        acc += g*g; x -= alpha*g/np.sqrt(acc); x=np.clip(x,-radius,radius); out.append(x.copy())
 50    return np.array(out)
 51
 52def run_shift(T=240, shift=80, seed=0, horizons=(1,2,4,8,16,32), alpha=.32, beta=3.):
 53    rng=np.random.default_rng(seed); target=np.zeros(T); target[shift:]=2.0
 54    pool=HeavyTailPool(1,horizons,alpha,beta,x0=[0.],radius=4.)
 55    xs=[]; qs=[]; losses=[]
 56    for t in range(T):
 57        # convex quadratic observation with moderate heavy-tailed gradient noise
 58        def gf(x, t=t): return (x-target[t:t+1]) + .06*rng.standard_t(1.7, size=1)
 59        def lf(x, t=t): return .5*float((x[0]-target[t])**2)
 60        x,l,q=pool.step(gf,lf); xs.append(float(x[0])); qs.append(q); losses.append(float(.5*(x[0]-target[t])**2))
 61    return np.array(xs),np.array(losses),np.array(qs),target
 62
 63def main():
 64    # 1) Exact constant-gradient scaling: displacement is alpha sum 1/sqrt(eps+t g^2).
 65    # Prediction: for fixed g and large H, displacement ~ 2 alpha sqrt(H)/|g|.
 66    alpha=.7; g=2.; eps=1e-8
 67    hs=np.array([16,32,64,128,256,512,1024])
 68    observed=[]; predicted=[]
 69    for h in hs:
 70        a=eps+np.arange(1,h+1)*g*g
 71        observed.append(alpha*np.sum(g/np.sqrt(a)))
 72        predicted.append(2*alpha*np.sqrt(h)) # / sign(g), because |g| cancels in 1-D
 73    slope=np.polyfit(np.log(hs),np.log(observed),1)[0]
 74    print('PREDICTION displacement ~ sqrt(H):')
 75    print('H observed predicted ratio')
 76    for h,o,p in zip(hs,observed,predicted): print(h, f'{o:.6f}', f'{p:.6f}', f'{o/p:.4f}')
 77    print('observed_loglog_slope',slope,'predicted',.5)
 78
 79    # 2) Meta mechanism: on fixed losses, log q_i/q_j follows sqrt(T) for a constant gap.
 80    # Construct two experts with losses [0, delta], initially equal. z=[-delta/2,+delta/2].
 81    beta=2.; delta=.2; Ts=np.array([25,100,400,1600]); logodds=[]; pred=[]
 82    for T in Ts:
 83        v=np.zeros(2); lo=0.
 84        for _ in range(T):
 85            z=np.array([-delta/2,delta/2]); v += z*z
 86            lo += float((-beta*z[0]/np.sqrt(eps+v[0])) - (-beta*z[1]/np.sqrt(eps+v[1])))
 87        logodds.append(lo)
 88        # Each expert has v_k=t*delta^2/4, so the two log-weight
 89        # increments sum asymptotically to 2*beta/sqrt(t).
 90        pred.append(4*beta*np.sqrt(T))
 91    meta_slope=np.polyfit(np.log(Ts),np.log(logodds),1)[0]
 92    print('\nPREDICTION meta log-odds ~ gap*sqrt(T):')
 93    print('T observed predicted ratio')
 94    for t,o,p in zip(Ts,logodds,pred): print(t, f'{o:.6f}', f'{p:.6f}', f'{o/p:.4f}')
 95    print('observed_loglog_slope',meta_slope,'predicted',.5)
 96
 97    # 3) Dynamic prediction: after a jump, best expert's recovery lag scales with H.
 98    # Recovery lag is the first time its squared error is below 0.25 after shift.
 99    print('\nPREDICTION dynamic recovery lag is O(H):')
100    for h in [1,2,4,8,16,32]:
101        x,loss,q,target=run_shift(T=180,shift=60,seed=11,horizons=(h,),alpha=.32,beta=3.)
102        post=np.where(np.abs(x[60:]-2.)<.5)[0]
103        lag=int(post[0]) if len(post) else 180-60
104        print(h,lag, 'lag_over_H',f'{lag/h:.3f}')
105
106    # Mini comparison: same shifted quadratic, pool versus fixed AdaGrad, equal one gradient/expert
107    seeds=range(10); pool_scores=[]; base_scores=[]; pool_recovery=[]; base_recovery=[]
108    for seed in seeds:
109        x,loss,q,target=run_shift(seed=seed)
110        rng=np.random.default_rng(seed); gs=[]
111        for t in range(240): gs.append(np.array([0.0])) # use deterministic baseline trajectory below
112        # baseline gets the same target sequence and no noisy gradient, isolating adaptation behavior
113        tg=np.zeros(240); tg[80:]=2.; xx=fixed_adagrad([0.], [np.array([0.])-tg[t:t+1] for t in range(240)], alpha=.32)
114        bl=.5*(xx[:,0]-tg)**2
115        pool_scores.append(loss[80:].mean()); base_scores.append(bl[80:].mean())
116        pool_recovery.append(np.where(np.abs(x[80:]-2.)<.5)[0][0] if np.any(np.abs(x[80:]-2.)<.5) else 160)
117        base_recovery.append(np.where(np.abs(xx[80:,0]-2.)<.5)[0][0] if np.any(np.abs(xx[80:,0]-2.)<.5) else 160)
118    print('\nCOMPARISON (pool has noisy heavy-tail gradients; baseline is fixed AdaGrad):')
119    print('pool post-shift loss',np.mean(pool_scores),'+/-',np.std(pool_scores)/np.sqrt(10))
120    print('baseline post-shift loss',np.mean(base_scores),'+/-',np.std(base_scores)/np.sqrt(10))
121    print('pool recovery',np.mean(pool_recovery),'baseline recovery',np.mean(base_recovery))
122
123if __name__=='__main__': main()