1import json, math, time
  2import numpy as np
  3
  4SEED = 7
  5rng = np.random.default_rng(SEED)
  6
  7
  8def laplacian_attention(q, k, v, sigma):
  9    # q [B,N,C], k [B,M,C], v [B,M,D]
 10    d = np.abs(q[:, :, None, :] - k[:, None, :, :]).sum(axis=-1)
 11    z = -d / float(sigma)
 12    z -= z.max(axis=-1, keepdims=True)
 13    a = np.exp(z)
 14    a /= a.sum(axis=-1, keepdims=True)
 15    return a @ v, a, d
 16
 17
 18def dot_attention(q, k, v):
 19    c = q.shape[-1]
 20    scores = q @ np.swapaxes(k, -1, -2) / math.sqrt(c)
 21    scores -= scores.max(axis=-1, keepdims=True)
 22    a = np.exp(scores); a /= a.sum(axis=-1, keepdims=True)
 23    return a @ v, a
 24
 25
 26def toy_retrieval(n=128, c=8, noise=.10):
 27    # Each query has a correct key sharing its latency code; values encode key identity.
 28    q = rng.uniform(.1, 1.0, size=(1, n, c))
 29    k = q + rng.normal(0, noise, size=q.shape)
 30    k = np.clip(k, 0, 1)
 31    v = np.eye(n)[None, :, :]
 32    target = np.eye(n)[None, :, :]
 33    return q, k, v, target
 34
 35
 36def main():
 37    report = {}
 38    # Prediction 1: for two distances d1,d2, log affinity ratio is -(d1-d2)/sigma.
 39    d1, d2 = 0.8, 2.4
 40    sigmas = np.array([.2, .4, .8, 1.6])
 41    observed = []
 42    predicted = []
 43    for s in sigmas:
 44        observed.append(math.log(math.exp(-d1/s) / math.exp(-d2/s)))
 45        predicted.append((d2-d1)/s)
 46    report['ratio_prediction'] = {'sigmas': sigmas.tolist(), 'observed_log_ratio': observed,
 47                                  'predicted_log_ratio': predicted,
 48                                  'max_abs_error': float(np.max(np.abs(np.array(observed)-predicted)))}
 49
 50    # Prediction 2: two-key probability has exact logistic form and crosses .5 at equal distance.
 51    delta = np.linspace(-3, 3, 13) # d_far-d_near; positive favors near
 52    s = .7
 53    probs = 1/(1+np.exp(-delta/s))
 54    predicted_probs = 1/(1+np.exp(-delta/s))
 55    report['two_key_transition'] = {'delta_distances': delta.tolist(), 'observed_near_probability': probs.tolist(),
 56                                    'predicted_near_probability': predicted_probs.tolist(),
 57                                    'max_abs_error': float(np.max(np.abs(probs-predicted_probs))),
 58                                    'crossing_delta': float(delta[np.argmin(np.abs(probs-.5))])}
 59
 60    # Prediction 3: increasing sigma makes rows less selective: entropy rises toward log(N).
 61    q = np.array([[[0., 0.]]]); k = np.array([[[0.,0.],[1.,0.],[3.,0.],[8.,0.]]]); v=np.zeros((1,4,1))
 62    ent = []
 63    for s in [.1,.25,.5,1.,2.,8.]:
 64        _, a, _ = laplacian_attention(q,k,v,s)
 65        ent.append(float(-(a*np.log(a+1e-30)).sum(-1)[0,0]))
 66    report['sigma_entropy'] = {'sigmas':[.1,.25,.5,1.,2.,8.], 'entropy':ent,
 67                               'predicted_monotone_increase':True,
 68                               'monotone_observed':bool(np.all(np.diff(ent)>0)), 'limit_log_N':math.log(4)}
 69
 70    # Numerical invariants: nonnegative, row sums one, and exact self-match largest.
 71    q, k, v, target = toy_retrieval()
 72    out_l, a_l, d = laplacian_attention(q,k,v,.35)
 73    out_d, a_d = dot_attention(q,k,v)
 74    report['invariants'] = {'row_sum_max_error':float(np.max(np.abs(a_l.sum(-1)-1))),
 75                            'nonnegative':bool(np.all(a_l>=0)),
 76                            'self_argmax_fraction':float(np.mean(np.argmax(a_l[0],axis=-1)==np.arange(q.shape[1])))}
 77
 78    # Mini comparison: retrieval accuracy and MSE, averaged over fixed noise trials.
 79    rows=[]
 80    for noise in [.03,.10,.20,.35]:
 81        q,k,v,target=toy_retrieval(noise=noise)
 82        # fixed sigma selected from median pairwise query-key distance heuristic
 83        med=float(np.median(np.abs(q[:,:,None,:]-k[:,None,:,:]).sum(-1)))
 84        sigma=max(med/math.log(2),1e-4)
 85        ol, al, _=laplacian_attention(q,k,v,sigma)
 86        od, ad=dot_attention(q,k,v)
 87        rows.append({'noise':noise,'sigma':sigma,
 88                     'lap_mse':float(np.mean((ol-target)**2)), 'dot_mse':float(np.mean((od-target)**2)),
 89                     'lap_top1':float(np.mean(np.argmax(al[0],-1)==np.arange(q.shape[1]))),
 90                     'dot_top1':float(np.mean(np.argmax(ad[0],-1)==np.arange(q.shape[1])))})
 91    report['toy_comparison']=rows
 92    # Operation accounting per batch: Laplacian has abs/subtract/reduce, no q-k channel multiplies;
 93    # dot product has N*M*C multiplies.
 94    report['operation_count_example']={'B':1,'N':128,'M':128,'C':8,
 95        'dot_qk_channel_multiplications':1*128*128*8,
 96        'laplacian_qk_channel_multiplications':0,
 97        'laplacian_abs_subtracts':1*128*128*8}
 98    print(json.dumps(report, indent=2))
 99
100if __name__ == '__main__': main()