← Lesson
Multi-Head Attention
Joeven
Run
Reset
Python loads on first run
import math def softmax(xs): m = max(xs) exps = [math.exp(x - m) for x in xs] s = sum(exps) return [e / s for e in exps] def dot(a, b): return sum(x * y for x, y in zip(a, b)) X = [ [1.0, 0.0, 0.0, 1.0], [0.9, 0.1, 1.0, 0.0], ] def head(X, lo, hi): n = len(X) d = hi - lo rows = [] for i in range(n): qi = X[i][lo:hi] scores = [dot(qi, X[j][lo:hi]) / math.sqrt(d) for j in range(n)] rows.append(softmax(scores)) return rows print("head0 id subspace") for row in head(X, 0, 2): print([round(x, 3) for x in row]) print("head1 time subspace") for row in head(X, 2, 4): print([round(x, 3) for x in row]) print("token0", X[0]) print("token1", X[1])
Run to execute this in your browser. Nothing is sent to a server.