LAB 09 · 视觉 Attention 完整层

from pathlib import Path
import json, sys, platform, hashlib
import numpy as np
import matplotlib.pyplot as plt
np.set_printoptions(precision=4, suppress=True)
OUTPUT = Path('outputs/LAB-B')
OUTPUT.mkdir(parents=True, exist_ok=True)
SEED = 20260918
DTYPE = np.float64
EPS = 1e-5
plt.rcParams.update({'font.size': 11, 'figure.dpi': 120, 'axes.spines.top': False, 'axes.spines.right': False})
image = np.array([[0., 0.2, 0.8, 1.0], [0.1, 0.3, 0.7, 0.9],
                  [0.9, 0.7, 0.2, 0.1], [1., 0.8, 0.3, 0.]], dtype=DTYPE)
def patchify(im):
    # Row-major order: top-left, top-right, bottom-left, bottom-right.
    return im.reshape(2, 2, 2, 2).transpose(0, 2, 1, 3).reshape(4, 4)
patches = patchify(image)
W_patch = np.array([[1., .1, 0., -.2], [.2, .8, -.1, 0.],
                    [0., -.2, 1., .2], [-.1, 0., .1, .9]])
position = np.array([[.2, 0., 0., 0.], [0., .2, 0., 0.],
                     [0., 0., .2, 0.], [0., 0., 0., .2]])
embedded = patches @ W_patch
X = embedded + position
print('image', image.shape, 'patches', patches.shape, 'embedding', embedded.shape, 'tokens X', X.shape)
print('patch vectors:\n', patches)
print('X = patches @ W_patch + position:\n', X)
fig, ax = plt.subplots(1, 2, figsize=(9, 3.8))
a = ax[0].imshow(image, cmap='gray', vmin=0, vmax=1)
ax[0].axvline(1.5, color='#e6b450', lw=2); ax[0].axhline(1.5, color='#e6b450', lw=2)
for i, (x, y) in enumerate([(0.5, .5), (2.5, .5), (.5, 2.5), (2.5, 2.5)]):
    ax[0].text(x,y,f'P{i}',ha='center',color='#f2ab32',fontweight='bold')
ax[0].set_title('Constructed image | 4 patches')
im = ax[1].imshow(X, cmap='viridis'); fig.colorbar(im, ax=ax[1], shrink=.8)
ax[1].set(title='Projected pixels + position', xlabel='Channel', ylabel='Patch token')
fig.tight_layout(); fig.savefig(OUTPUT/'01-patches.png'); plt.show()

D, HEADS, DH, DFF = 4, 2, 2, 6
rng = np.random.default_rng(SEED)
weights = {name: rng.normal(0, .35, shape).astype(DTYPE) for name, shape in [
    ('W_Q', (D,D)), ('W_K', (D,D)), ('W_V', (D,D)), ('W_O', (D,D)),
    ('W_1', (D,DFF)), ('W_2', (DFF,D))]}
def ln(x):
    return (x-x.mean(axis=-1,keepdims=True))/np.sqrt(x.var(axis=-1,keepdims=True)+EPS)
def softmax(x):
    e = np.exp(x-x.max(axis=-1,keepdims=True))
    return e/e.sum(axis=-1,keepdims=True)
def gelu(x):
    # The tanh approximation is declared explicitly; no library default ambiguity.
    return .5*x*(1+np.tanh(np.sqrt(2/np.pi)*(x+.044715*x**3)))
def encoder(x):
    t = {'input': x.copy(), 'ln1': ln(x)}
    for name in ('Q','K','V'):
        t[name] = t['ln1'] @ weights['W_'+name]
        t[name+'_heads'] = t[name].reshape(len(x),HEADS,DH).transpose(1,0,2)
    t['scores'] = t['Q_heads'] @ t['K_heads'].transpose(0,2,1) / np.sqrt(DH)
    t['attention'] = softmax(t['scores'])
    t['read_values'] = t['attention'] @ t['V_heads']
    t['concat'] = t['read_values'].transpose(1,0,2).reshape(len(x),D)
    t['projected_attention'] = t['concat'] @ weights['W_O']
    t['residual1'] = x + t['projected_attention']
    t['ln2'] = ln(t['residual1'])
    t['mlp_hidden_pre'] = t['ln2'] @ weights['W_1']
    t['mlp_hidden'] = gelu(t['mlp_hidden_pre'])
    t['mlp_output'] = t['mlp_hidden'] @ weights['W_2']
    t['output'] = t['residual1'] + t['mlp_output']
    return t
trace = encoder(X)
for name, value in trace.items(): print(f'{name:22s} {str(value.shape):14s}')
print('Head 0 attention weights (row=query, column=key):\n', trace['attention'][0])
print('Head 0 read-back values:\n', trace['read_values'][0])
print('Final output:\n', trace['output'])

for name in ['ln1','Q','K','V','scores','concat','projected_attention','residual1','ln2','mlp_hidden_pre','mlp_hidden','mlp_output','output']:
    print(name, trace[name].shape)
    print(trace[name])

h, q = 0, 0
manual_read = sum(trace['attention'][h,q,k]*trace['V_heads'][h,k] for k in range(4))
assert np.allclose(manual_read, trace['read_values'][h,q], atol=1e-12)
assert np.allclose(trace['attention'].sum(axis=-1), 1)
assert np.allclose(trace['residual1']-X, trace['projected_attention'])
assert np.allclose(trace['output']-trace['residual1'], trace['mlp_output'])
assert all(np.isfinite(v).all() for v in trace.values())
print('Manual readback:', manual_read)
print('Vectorized readback:', trace['read_values'][h,q])
print('Softmax row sums:', trace['attention'].sum(axis=-1))
fig, axes = plt.subplots(1, 3, figsize=(12, 3.4))
for h in range(HEADS):
    im=axes[h].imshow(trace['attention'][h], vmin=0, vmax=1, cmap='Blues')
    axes[h].set(title=f'Head {h}: all-to-all read', xlabel='Read from key', ylabel='Query')
    for i in range(4):
        for j in range(4): axes[h].text(j,i,f'{trace["attention"][h,i,j]:.2f}',ha='center',va='center')
fig.colorbar(im, ax=list(axes[:2]), fraction=.025)
axes[2].plot(np.linalg.norm(trace['projected_attention'],axis=1),'o-',label='Attention branch')
axes[2].plot(np.linalg.norm(trace['mlp_output'],axis=1),'s--',label='MLP branch')
axes[2].set(xlabel='Token',ylabel='Branch output L2 norm',title='Two updates to a residual path')
axes[2].legend(fontsize=9); fig.subplots_adjust(wspace=.5)
fig.savefig(OUTPUT/'02-readback.png',bbox_inches='tight'); plt.show()

changed_image = image.copy(); changed_image[3,3] += .4
X_pixel = patchify(changed_image) @ W_patch + position
pixel_trace = encoder(X_pixel)
no_position = encoder(embedded)
identical_input = np.repeat(embedded[:1], 4, axis=0)
identical = encoder(identical_input)
permutation = np.array([2,0,3,1])
permuted = encoder(embedded[permutation])
assert np.allclose(identical['attention'], .25)
assert np.allclose(identical['output'], identical['output'][0])
assert np.allclose(permuted['output'], no_position['output'][permutation])
changes = {
    'one_pixel': {k: np.linalg.norm(pixel_trace[k]-trace[k],axis=-1).tolist() for k in ['input','projected_attention','residual1','mlp_output','output']},
    'remove_position': np.linalg.norm(no_position['output']-trace['output'],axis=-1).tolist(),
    'identical_attention_max_deviation_from_uniform': float(np.abs(identical['attention']-.25).max()),
    'permutation_max_error': float(np.abs(permuted['output']-no_position['output'][permutation]).max())}
print(json.dumps(changes,indent=2))
fig, ax=plt.subplots(figsize=(8,3.8))
for name, style in [('input','o--'),('projected_attention','s-'),('output','^-')]:
    ax.plot(range(4), changes['one_pixel'][name], style, label=name)
ax.set(xlabel='Token index',ylabel='L2 change',title='Change one pixel in patch 3; all weights fixed',xticks=range(4))
ax.legend(); fig.tight_layout(); fig.savefig(OUTPUT/'03-change-one-patch.png'); plt.show()

import csv
artifact = {'kind':'exact_constructed_calculation_not_trained_model', 'seed':SEED,
            'environment':{'python':sys.version,'numpy':np.__version__,'platform':platform.platform()},
            'structure':{'norm':'pre-norm','heads':HEADS,'d_model':D,'d_head':DH,'d_ff':DFF,'epsilon':EPS,
                         'mask':'none','dropout':0,'gelu':'tanh approximation','class_token':False},
            'image':image.tolist(),'patches':patches.tolist(),'position':position.tolist(),
            'patch_projection':W_patch.tolist(),'weights':{k:v.tolist() for k,v in weights.items()},
            'trace':{k:v.tolist() for k,v in trace.items()},'changes':changes,
            'checks':{'manual_readback':True,'softmax_rows_sum_one':True,'residuals':True,
                      'identical_tokens_uniform':True,'no_position_permutation_equivariance':True}}
(OUTPUT/'results.json').write_text(json.dumps(artifact,indent=2),encoding='utf-8')
np.savez_compressed(OUTPUT/'tensors.npz', **trace, **weights, image=image, position=position, W_patch=W_patch)
with (OUTPUT/'token-changes.csv').open('w',newline='') as f:
    w=csv.writer(f); w.writerow(['token','input_l2_change','attention_branch_l2_change','output_l2_change'])
    for i in range(4): w.writerow([i]+[changes['one_pixel'][k][i] for k in ['input','projected_attention','output']])
print('Saved:', sorted(p.name for p in OUTPUT.iterdir()))