Self-Attention

জিনিসটা কী, কেন দরকার

Library তে বই খোঁজার কথা ভাবো। তুমি ঢোকো একটা query নিয়ে ("দেরিতে delivery নিয়ে বই")। প্রতিটা বইয়ের পিঠে একটা label, একটা key। তোমার query কে প্রতিটা key এর সাথে মেলাও, যাদের key সবচেয়ে বেশি মেলে সেই বইগুলো নামাও, আর তাদের ভেতরের জিনিস, মানে value, পড়ো, যেটা বেশি মেলে সেখান থেকে বেশি নিয়ে।

Self-attention হলো এই খোঁজা, একটা বাক্যের প্রতিটা token একসাথে করছে। প্রতিটা token বাকি সবাইকে একটা প্রশ্ন করে, যারা মেলে তাদের ভেতরের জিনিস জড়ো করে, আর এমন একটা নতুন vector হয়ে যায় যেটা তার context জানে। Part 3 ধারণাটা দেখিয়েছিল হাতে বানানো vector দিয়ে; এখানে আসল mechanism, শেখা query, key আর value matrix সহ, causal mask যেটা একটা শব্দকে ভবিষ্যতের শব্দে উঁকি দিতে দেয় না, আর একাধিক head। প্রতিটা transformer এর মূল operation এটাই, প্রতিটা layer এ বারবার।

ভেতরে কীভাবে চলে

n টা token, প্রতিটা d সাইজের vector (embedding আর position), একটা (n, d) shape এর matrix X এ সাজানো:

1. Projection। তিনটা শেখা weight matrix প্রতিটা token এর তিনটা রূপ বানায়: Q = X·W_q (আমি কী খুঁজছি?), K = X·W_k (আমি কী দিতে পারি?), V = X·W_v (আমাকে বাছলে কী পাঠাব?)।
2. Score। S = Q·Kᵀ / √d: প্রতিটা query র সাথে প্রতিটা key এর dot product, একটা (n, n) এর table। √d দিয়ে ভাগ করলে সংখ্যাগুলো এমন range এ থাকে যেখানে softmax একটা মাত্র 1 না দিয়ে weight ছড়িয়ে দেয়।
3. Causal mask। Text generation এ token i তার পরের token দেখতে পারবে না, নইলে training এ চুরি হতো: answer (পরের শব্দ) input এই থাকত। Diagonal এর উপরের প্রতিটা score −∞ করে দাও; softmax এর পরে সেই weight গুলো ঠিক 0।
4. প্রতিটা row কে softmax করে weight বানাও, যোগফল 1।
5. মেশাও। Output = weights · V: প্রতিটা token এর নতুন vector হলো সে যাদের দিকে তাকিয়েছে তাদের value vector এর weight দেওয়া গড়।

Attention(Q, K, V) = softmax(Q·Kᵀ / √d + mask) · V

Multi-head attention। প্রতি layer এ একটাই attention pattern সীমিত: একটা token কে হয়তো একসাথে তার subject, তার verb আর আগের বাক্য ট্র্যাক করতে হয়। তাই d টা সংখ্যাকে h টা head এ ভাগ করা হয় (ধরো d/8 করে 8 টা head)। প্রতিটা head এর নিজের Q, K, V আর নিজের pattern; output গুলো জোড়া লাগিয়ে আরেকটা matrix দিয়ে মেশানো হয়। আলাদা head আলাদা কাজ শেখে।

খরচ। Score এর table n × n। Context দ্বিগুণ করলে attention এর কাজ আর memory চার গুণ। এই একটা সত্য context এর সীমা, দাম আর Part 9 এর অনেকটা ঠিক করে দেয়।

আগে অনুমান করো

Causal mask সহ, একটা sequence এ 3 নম্বর token থেকে 5 নম্বর token এর attention weight কত?

ঠিক 0খুব ছোট কিন্তু 0 এর বেশি1, কারণ পরের token ই সবচেয়ে জরুরিদুইটা token কতটা একরকম তার উপর নির্ভর করে

ছবিতে দেখো

Play

Code এ ধাপে ধাপে attention

4 টা token এ single-head causal self-attention, NumPy তে ধাপে ধাপে, PyTorch এর scaled_dot_product_attention এর সাথে মেলানো। তারপর একই input দুইটা head এ ভাগ, আর √d দিয়ে scale কেন করা হয় তার একটা মাপ।

import numpy as np
import torch

np.set_printoptions(precision=2, suppress=True)
rng = np.random.default_rng(0)

tokens = ["the", "order", "was", "late"]
X = rng.normal(size=(4, 8)).round(1)            # 4 tokens, each an 8-number embedding (+ position)
Wq, Wk, Wv = (rng.normal(size=(8, 8)).round(1) / np.sqrt(8) for _ in range(3))

# 1) Every token makes a query, a key and a value with three learned matrices.
Q, K, V = X @ Wq, X @ Wk, X @ Wv

# 2) Scores: each query against every key, scaled by sqrt(d) so softmax doesn't saturate.
scores = Q @ K.T / np.sqrt(8)

# 3) Causal mask: a token may not look at tokens after it (that would be reading the answer).
mask = np.triu(np.ones((4, 4), dtype=bool), k=1)
scores = np.where(mask, -np.inf, scores)

# 4) Softmax per row -> attention weights; 5) weighted mix of values.
weights = np.exp(scores - scores.max(axis=1, keepdims=True))
weights /= weights.sum(axis=1, keepdims=True)
out = weights @ V

print("attention weights (row = who is looking, column = who is looked at):")
print("         " + "".join(f"{t:>7}" for t in tokens))
for t, row in zip(tokens, weights):
    print(f"{t:>8} " + "".join(f"{w:7.2f}" for w in row))
print("\noutput shape:", out.shape, " (one new 8-number vector per token)")

# The same thing with PyTorch's fused kernel.
ref = torch.nn.functional.scaled_dot_product_attention(
    torch.tensor(Q)[None], torch.tensor(K)[None], torch.tensor(V)[None], is_causal=True)[0].numpy()
print("matches torch scaled_dot_product_attention:", np.allclose(out, ref))

# Multi-head: split the 8 numbers into 2 heads of 4; each head attends on its own.
print("\n2 heads, row for 'late':")
for h in range(2):
    s = slice(h * 4, (h + 1) * 4)
    sc = Q[:, s] @ K[:, s].T / np.sqrt(4)
    sc = np.where(mask, -np.inf, sc)
    w = np.exp(sc - sc.max(axis=1, keepdims=True)); w /= w.sum(axis=1, keepdims=True)
    print(f"  head {h + 1}: " + "  ".join(f"{t} {x:.2f}" for t, x in zip(tokens, w[3])))

# Why scale by sqrt(d)? Dot products of random vectors grow with d.
for d in (8, 512):
    a, b = rng.normal(size=(2000, d)), rng.normal(size=(2000, d))
    print(f"d={d:<3}: std of raw dot products {np.std((a * b).sum(1)):5.1f}, after /sqrt(d) {np.std((a * b).sum(1) / np.sqrt(d)):.2f}")

Output:

attention weights (row = who is looking, column = who is looked at):
             the  order    was   late
     the    1.00   0.00   0.00   0.00
   order    0.99   0.01   0.00   0.00
     was    0.29   0.36   0.35   0.00
    late    0.22   0.27   0.21   0.30

output shape: (4, 8)  (one new 8-number vector per token)
matches torch scaled_dot_product_attention: True

2 heads, row for 'late':
  head 1: the 0.16  order 0.37  was 0.17  late 0.30
  head 2: the 0.31  order 0.17  was 0.27  late 0.25
d=8  : std of raw dot products   2.7, after /sqrt(d) 0.97
d=512: std of raw dot products  22.6, after /sqrt(d) 1.00

কী দেখলাম:

• উপরের ত্রিভুজ পুরোটাই 0.00: mask কাজ করছে। "the" এর দেখার মতো শুধু নিজেই আছে, তাই তার weight 1.00।
• Row গুলো আলাদা। "order" 0.99 দেয় "the" কে; "was" তিনটা token এ প্রায় সমান ছড়ায়; "late" চারটায়। এখানে weight গুলো random projection থেকে এসেছে; train করা model এ এগুলো ব্যাকরণ আর অর্থ ধরে।
• PyTorch এর fused kernel একই output দেয়। আসল code এ এটাই ডাকবে (অনেক দ্রুত আর কম memory, FlashAttention ধরনের kernel দিয়ে), কিন্তু হিসাব করে হুবহু এই পাঁচটা ধাপ।
• দুইটা head, দুইটা pattern: "late" এর জন্য head 1 সবচেয়ে বেশি তাকায় "order" এর দিকে (0.37), head 2 "the" এর দিকে (0.31)। একই token, আলাদা প্রশ্ন।
• √d কেন: d = 512 এ কাঁচা dot product এর ছড়ানো 22.6, যেটা softmax কে প্রায় সব weight একটা token এ দিতে ঠেলে দিত আর gradient মিলিয়ে দিত। √512 দিয়ে ভাগ করলে আবার 1.00, d = 8 এর মতোই।

Production এ যা মাথায় রাখবে

• Fused kernel ব্যবহার করো। scaled_dot_product_attention (বা FlashAttention) পুরো n × n matrix কখনো জমায় না, আর কয়েক গুণ দ্রুত। Production code এ হাতে attention লেখা একটা performance bug।
• Context এর সাথে memory বাড়ে। Inference এ আগের প্রতিটা token এর K আর V রেখে দেওয়া হয় (KV cache, এই part এর lesson 6)। একটা GPU কতজন user সামলাবে, সেটা প্রায়ই weight না, এই cache ঠিক করে।
• Mask হলো সঠিকতার সীমানা। Training এ ভুল mask ভবিষ্যতের token leak করে, ফলে এমন model হয় যেটা training এ দারুণ দেখায় আর generation এ ফেল করে। Batch এর padding token দেরও mask লাগে, নইলে তারা attention weight পেয়ে যায়।
• Grouped-query attention (GQA), এখনকার বেশিরভাগ model এ যেটা আছে, কয়েকটা query head এর মধ্যে K আর V ভাগ করে নেয়, যাতে অল্প মান হারিয়েই KV cache ছোট হয়।

এক কথায়

Self-attention প্রতিটা token কে query, key আর value তে project করে, প্রতিটা query কে প্রতিটা key এর সাথে score করে (√d দিয়ে scale), ভবিষ্যৎ mask করে, score কে softmax করে weight বানায় আর value মেশায়; একাধিক head একসাথে আলাদা প্রশ্ন করে, আর n × n এর score table এর কারণেই লম্বা context এর খরচ বর্গ হারে বাড়ে।

নিজে করে দেখো

Example থেকে causal mask সরাও (np.where লাইনটা comment করো আর is_causal=False দাও) আর "the" এর row দেখো। কী বদলায়, আর mask ছাড়া একটা পরের-শব্দ predict করা model train করা অকাজের কেন?

Answer

"the" এখন চারটা token এই তাকায়, তার পরের "order", "was" আর "late" সহ। Mask ছাড়া train করা পরের-শব্দের model যে শব্দটা predict করার কথা সেটাই সরাসরি দেখতে পেত, তাই কপি করেই training loss প্রায় শূন্যে নামত, আর generation এর সময়, যখন ভবিষ্যতের token এখনো নেই, তখন কাজের কিছুই শেখা থাকত না। (BERT এর মতো যেসব model ফাঁকা জায়গা ভরে, তারা ইচ্ছা করেই causal mask রাখে না, কিন্তু তাদের বাম থেকে ডানে generate করতে train করা হয় না।)