Transformerの中身をのぞく:attentionを手で計算する

// Self-attention を1ヘッド分、手で計算する。重みは学習前の乱数なので、パターン自体に意味はない。
const tokens = ['the', 'model', 'predicts', 'the', 'next'];
const d = 4;                                   // ベクトルの次元(実際のモデルは、はるかに大きい)

function rng(seed) {
  return () => {
    seed = (seed + 0x6D2B79F5) | 0;
    let t = Math.imul(seed ^ (seed >>> 15), 1 | seed);
    t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t;
    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
  };
}
const random = rng(2);
const matrix = (rows, cols) =>
  Array.from({ length: rows }, () => Array.from({ length: cols }, () => random() * 2 - 1));
const mul = (A, B) =>
  A.map(row => B[0].map((_, j) => row.reduce((sum, v, k) => sum + v * B[k][j], 0)));
const transpose = A => A[0].map((_, j) => A.map(row => row[j]));
const softmax = row => {
  const max = Math.max(...row);
  const exps = row.map(v => Math.exp(v - max));
  const sum = exps.reduce((a, b) => a + b, 0);
  return exps.map(v => v / sum);
};

const X = matrix(tokens.length, d);            // 各トークンのベクトル(埋め込み)
const Wq = matrix(d, d), Wk = matrix(d, d), Wv = matrix(d, d);
const Q = mul(X, Wq), K = mul(X, Wk), V = mul(X, Wv);

// 1. 各トークンの Q と、すべてのトークンの K の内積を取り、√d で割る
const scores = mul(Q, transpose(K)).map(row => row.map(v => v / Math.sqrt(d)));
// 2. 未来のトークンは見えないようにする(右上を -∞ にする)
const masked = scores.map((row, i) => row.map((v, j) => (j <= i ? v : -Infinity)));
// 3. 行ごとに softmax をかけて、合計が 1 の重みにする
const weights = masked.map(softmax);
// 4. 重みで V を足し合わせて、各トークンの新しいベクトルにする
const output = mul(weights, V);

console.log('weights:');
weights.forEach((row, i) => console.log(tokens[i].padEnd(9), row.map(w => w.toFixed(2)).join(' ')));
console.log('row sums:', weights.map(r => r.reduce((a, b) => a + b, 0).toFixed(2)).join(' '));
console.log('output of "next":', output[4].map(v => v.toFixed(2)).join(' '));

# Self-attention を1ヘッド分、手で計算する(attention.mjs と同じ計算を Python で書いたもの)。
import math

tokens = ['the', 'model', 'predicts', 'the', 'next']
d = 4                                    # ベクトルの次元(実際のモデルは、はるかに大きい)


def rng(seed):
    """JavaScript版と同じ式の乱数。同じ種なら、同じ数列になる。"""
    state = seed & 0xFFFFFFFF

    def random():
        nonlocal state
        state = (state + 0x6D2B79F5) & 0xFFFFFFFF
        t = ((state ^ (state >> 15)) * (1 | state)) & 0xFFFFFFFF
        t = ((t + (((t ^ (t >> 7)) * (61 | t)) & 0xFFFFFFFF)) & 0xFFFFFFFF) ^ t
        return ((t ^ (t >> 14)) & 0xFFFFFFFF) / 4294967296

    return random


random = rng(2)


def matrix(rows, cols):
    return [[random() * 2 - 1 for _ in range(cols)] for _ in range(rows)]


def mul(A, B):
    return [[sum(a * b for a, b in zip(row, col)) for col in zip(*B)] for row in A]


def transpose(A):
    return [list(col) for col in zip(*A)]


def softmax(row):
    top = max(row)
    exps = [math.exp(v - top) for v in row]
    total = sum(exps)
    return [v / total for v in exps]


X = matrix(len(tokens), d)               # 各トークンのベクトル(埋め込み)
Wq, Wk, Wv = matrix(d, d), matrix(d, d), matrix(d, d)
Q, K, V = mul(X, Wq), mul(X, Wk), mul(X, Wv)

# 1. 各トークンの Q と、すべてのトークンの K の内積を取り、√d で割る
scores = [[v / math.sqrt(d) for v in row] for row in mul(Q, transpose(K))]
# 2. 未来のトークンは見えないようにする(右上を -∞ にする)
masked = [[v if j <= i else -math.inf for j, v in enumerate(row)] for i, row in enumerate(scores)]
# 3. 行ごとに softmax をかけて、合計が 1 の重みにする
weights = [softmax(row) for row in masked]
# 4. 重みで V を足し合わせて、各トークンの新しいベクトルにする
output = mul(weights, V)

print('weights:')
for token, row in zip(tokens, weights):
    print(token.ljust(9), ' '.join(f'{w:.2f}' for w in row))
print('row sums:', ' '.join(f'{sum(row):.2f}' for row in weights))
print('output of "next":', ' '.join(f'{v:.2f}' for v in output[4]))