1import math, random, time
  2import numpy as np
  3import torch
  4from torch import nn
  5import torch.nn.functional as F
  6
  7SEED = 153
  8
  9def seed_all(seed=SEED):
 10    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 11
 12
 13def math_check():
 14    seed_all(7)
 15    beta, eta, steps, n, dim = 0.07, 0.013, 80, 12000, 3
 16    z = torch.zeros(n, dim)
 17    amp = math.sqrt(2 * beta * eta)
 18    for _ in range(steps):
 19        z += amp * torch.randn_like(z)
 20    empirical = z.var(0, unbiased=True).mean().item()
 21    predicted = 2 * beta * eta * steps
 22    x = torch.ones(n)
 23    for _ in range(steps):
 24        x *= (1 - eta)
 25    expected = (1 - eta) ** steps
 26    return {
 27        'diffusion_empirical_variance': empirical,
 28        'diffusion_predicted_variance': predicted,
 29        'diffusion_relative_error': abs(empirical - predicted) / predicted,
 30        'quadratic_contraction_empirical': x.mean().item(),
 31        'quadratic_contraction_expected': expected,
 32    }
 33
 34
 35class ParticleAttention(nn.Module):
 36    def __init__(self, d_model=24, heads=8, particle=False, beta=0.0, noise_eta=0.03):
 37        super().__init__()
 38        assert d_model % heads == 0
 39        self.d, self.h, self.dk = d_model, heads, d_model // heads
 40        self.particle, self.beta, self.noise_eta = particle, beta, noise_eta
 41        self.q = nn.Parameter(torch.randn(heads, d_model, self.dk) * .12)
 42        self.k = nn.Parameter(torch.randn(heads, d_model, self.dk) * .12)
 43        self.v = nn.Parameter(torch.randn(heads, d_model, self.dk) * .12)
 44        self.o = nn.Parameter(torch.randn(heads, self.dk, d_model) * .12)
 45        self.norm = nn.LayerNorm(d_model)
 46
 47    def forward(self, x, noise=True):
 48        # x: batch, sequence, model dimension. Each particle is one attention head.
 49        q = torch.einsum('bnd,hdk->bhnk', x, self.q)
 50        k = torch.einsum('bnd,hdk->bhnk', x, self.k)
 51        v = torch.einsum('bnd,hdk->bhnk', x, self.v)
 52        a = torch.softmax(torch.einsum('bhnk,bhmk->bhnm', q, k) / math.sqrt(self.dk), dim=-1)
 53        hv = torch.einsum('bhnm,bhmk->bhnk', a, v)
 54        outs = torch.einsum('bhnk,hkd->bhnd', hv, self.o)
 55        y = outs.mean(1)
 56        return self.norm(x + y), a, outs
 57
 58    @torch.no_grad()
 59    def langevin_noise(self):
 60        if self.particle and self.beta > 0:
 61            amp = math.sqrt(2 * self.beta * self.noise_eta)
 62            for p in (self.q, self.k, self.v, self.o):
 63                p.add_(amp * torch.randn_like(p))
 64
 65    @torch.no_grad()
 66    def diversity(self):
 67        # Mean pairwise cosine similarity of flattened particle parameters.
 68        w = torch.cat([z.flatten(1) for z in (self.q, self.k, self.v, self.o)], dim=1)
 69        c = F.normalize(w, dim=1) @ F.normalize(w, dim=1).T
 70        h = self.h
 71        return ((c.sum() - h) / (h * (h - 1))) if h > 1 else 1.0
 72
 73    def dropped_forward(self, x, keep):
 74        q = torch.einsum('bnd,hdk->bhnk', x, self.q[:keep])
 75        k = torch.einsum('bnd,hdk->bhnk', x, self.k[:keep])
 76        v = torch.einsum('bnd,hdk->bhnk', x, self.v[:keep])
 77        a = torch.softmax(torch.einsum('bhnk,bhmk->bhnm', q, k) / math.sqrt(self.dk), dim=-1)
 78        hv = torch.einsum('bhnm,bhmk->bhnk', a, v)
 79        y = torch.einsum('bhnk,hkd->bhnd', hv, self.o[:keep]).mean(1)
 80        return self.norm(x + y)
 81
 82
 83class TinyClassifier(nn.Module):
 84    def __init__(self, heads, beta):
 85        super().__init__()
 86        self.attn = ParticleAttention(heads=heads, particle=(beta > 0), beta=beta)
 87        self.fc = nn.Sequential(nn.Linear(24, 32), nn.Tanh(), nn.Linear(32, 2))
 88    def forward(self, x):
 89        y, a, outs = self.attn(x)
 90        return self.fc(y[:, 0]), a, outs
 91
 92
 93def data(seed=SEED):
 94    g = torch.Generator().manual_seed(seed)
 95    ntr, nte, seq, d = 768, 256, 8, 24
 96    # Label is encoded by a shared direction at token zero; other tokens are distractors.
 97    direction = F.normalize(torch.randn(d, generator=g), dim=0)
 98    def make(n):
 99        y = torch.randint(0, 2, (n,), generator=g)
100        x = .8 * torch.randn(n, seq, d, generator=g)
101        x[:, 0] += (2*y.float()-1)[:, None] * direction[None, :] * 1.4
102        return x, y
103    return make(ntr), make(nte)
104
105
106def run_one(heads, beta, epochs=18):
107    seed_all(SEED + heads + int(beta * 1e6))
108    (xt, yt), (xv, yv) = data()
109    model = TinyClassifier(heads, beta)
110    opt = torch.optim.Adam(model.parameters(), lr=0.012)
111    bs = 64; losses=[]; t0=time.time()
112    for ep in range(epochs):
113        perm = torch.randperm(len(xt))
114        model.train()
115        for ix in perm.split(bs):
116            logits, _, _ = model(xt[ix])
117            loss = F.cross_entropy(logits, yt[ix])
118            opt.zero_grad(); loss.backward(); opt.step()
119            model.attn.langevin_noise()
120            losses.append(loss.item())
121    model.eval()
122    with torch.no_grad():
123        logits, maps, outs = model(xv)
124        val_loss = F.cross_entropy(logits, yv).item()
125        acc = (logits.argmax(1) == yv).float().mean().item()
126        keep = max(1, heads // 2)
127        drop_acc = (model.attn.dropped_forward(xv, keep)[:,0] if False else
128                    model.fc(model.attn.dropped_forward(xv, keep)[:,0]).argmax(1) == yv).float().mean().item()
129        # Attention map similarity across particles, averaged over examples/queries.
130        flat = maps.permute(1,0,2,3).reshape(heads, -1)
131        sim = (F.normalize(flat, dim=1) @ F.normalize(flat, dim=1).T)
132        map_sim = ((sim.sum()-heads)/(heads*(heads-1))).item() if heads > 1 else 1.0
133    return {'heads':heads, 'beta':beta, 'val_loss':val_loss, 'val_accuracy':acc,
134            'half_particle_accuracy':drop_acc, 'parameter_cosine':float(model.attn.diversity()),
135            'attention_map_cosine':map_sim, 'last_train_loss':float(np.mean(losses[-20:])),
136            'seconds':time.time()-t0}
137
138
139def main():
140    print({'math_check': math_check()})
141    results=[]
142    # Standard fixed 8-head attention versus the same particle population with Langevin noise.
143    for beta in (0.0, 0.001, 0.005):
144        results.append(run_one(8, beta))
145    print({'results':results})
146
147if __name__ == '__main__':
148    main()