Self-attention আর multi-head attention — matrix দিয়ে হাতে-কলমে

পর্ব ৩২-এ (Attention — query, key, value আর alignment) decoder-এর একটি query encoder-এর সব token-এর দিকে তাকিয়েছিল। এবার প্রশ্নটি একটু বদলাই: একই sequence-এর token যদি একে অন্যের দিকে তাকায়? এই একটি বদল থেকেই Transformer-এর মূল অংশটি আসে। Dot product মনে না থাকলে AI-এর ভিত্তি কোর্সের পর্ব ৩৩ দেখে নিন; একবার হাতে করা attention আছে AI-এর ভিত্তি কোর্সের পর্ব ৭০-এ। Batched matmul আর einsum এসেছে পর্ব ৭-এ। এখানে আমরা গভীরে যাব।

১ · বাক্য নিজেকেই পড়ে — encoder–decoder attention থেকে self-attention

বাজারঘরের একটি review নিন: “ফোনটা ভালো, কিন্তু সেটি দেরিতে এসেছে”. “সেটি” বুঝতে দূরের “ফোনটা” দরকার। পর্ব ২৮-এর RNN-এ এই তথ্যকে মাঝের প্রতিটি ধাপ পার হতে হয়। পর্ব ৩০-এ দেখেছেন, পথ লম্বা হলে তথ্য হারায়।

তিন রকম attention পাশাপাশি রাখি। Attention-এ তিনটি জিনিস লাগে: query (কী খুঁজছি), key (কার কাছে কী আছে তার লেবেল) আর value (আসল যে তথ্য মেশানো হবে)।

২ · Q, K, V — একই input থেকে তিনটি শেখা projection

Input X-এর shape (B, T, d). B হলো batch-এ কয়টি উদাহরণ, T কয়টি token, d প্রতি token-এ কয়টি সংখ্যা। একজন বক্তার তিনটি ভূমিকা: কী জানতে চান, নিজের সম্পর্কে কী লেবেল দেখান, আর ডাক পেলে কী দেন। এই তিনের জন্য তিনটি আলাদা শেখা matrix থাকে। একে বলে learned projection.

Q = XW_Q,   K = XW_K,   V = XW_V

Z = A V,   A = softmax(Q Kᵀ / √d_k)

এখানে d_k হলো key-র দৈর্ঘ্য, Kᵀ মানে শেষ দুই অক্ষ অদলবদল করা K, আর softmax চলে প্রতিটি সারিতে আলাদা। এটিই scaled dot-product attention. A-কে বলে attention weights; A[i, j] মানে token i, token j-এর value-র কত ভাগ নেবে।

নিচের কোডে বাজারঘরের ২টি review, প্রতিটিতে ৪টি token. Embedding আর W এলোমেলো, কারণ এখানে শুধু হিসাবের ছাঁচ দেখছি, এখনো শেখানো হয়নি। আমরা row-vector নিয়মে লিখছি, তাই XW; nn.Linear একই কাজ করে।

import numpy as np
np.set_printoptions(precision=3, suppress=True)
rng = np.random.default_rng(42)
B, T, d = 2, 4, 8                 # ২টি review, প্রতিটিতে ৪ token, token-প্রতি ৮টি সংখ্যা
X = rng.normal(size=(B, T, d))
Wq, Wk, Wv = rng.normal(size=(3, d, d)) / np.sqrt(d)   # শেখার মতো ৩টি matrix

def softmax(z):                   # বড় মান বাদ দিয়ে নিরাপদ softmax (পর্ব ০৩)
    e = np.exp(z - z.max(-1, keepdims=True))
    return e / e.sum(-1, keepdims=True)

def self_attention(Q, K, V):
    S = Q @ K.transpose(0, 2, 1) / np.sqrt(K.shape[-1])
    A = softmax(S)
    return A @ V, A, S

Q, K, V = X @ Wq, X @ Wk, X @ Wv
Z, A, S = self_attention(Q, K, V)
for name, arr in zip("XQKVSAZ", (X, Q, K, V, S, A, Z)):
    print(name, arr.shape)
print("প্রতি সারির যোগফল:", A[0].sum(-1))
print("A[0]:\n", A[0])
print("A প্রতিসম?", np.allclose(A[0], A[0].T))

by_hand = sum(A[0, 1, j] * V[0, j] for j in range(T))  # token ১-এর output, লুপে
print("লুপ বনাম matrix, max|diff|:", np.abs(by_hand - Z[0, 1]).max())
same = np.einsum("bts,bsd->btd", A, V)                 # একই কাজ einsum-এ
print("einsum বনাম @, max|diff|:", np.abs(same - Z).max())

_, A0, _ = self_attention(X, X, X)                     # projection ছাড়া
d0, d1 = np.diagonal(A0, axis1=1, axis2=2), np.diagonal(A, axis1=1, axis2=2)
print(f"নিজের দিকে গড় weight: projection ছাড়া {d0.mean():.3f}, সহ {d1.mean():.3f}")

X (2, 4, 8)
Q (2, 4, 8)
K (2, 4, 8)
V (2, 4, 8)
S (2, 4, 4)
A (2, 4, 4)
Z (2, 4, 8)
প্রতি সারির যোগফল: [1. 1. 1. 1.]
A[0]:
 [[0.272 0.175 0.227 0.327]
 [0.187 0.313 0.28  0.22 ]
 [0.289 0.195 0.318 0.199]
 [0.235 0.27  0.34  0.155]]
A প্রতিসম? False
লুপ বনাম matrix, max|diff|: 0.0
einsum বনাম @, max|diff|: 0.0
নিজের দিকে গড় weight: projection ছাড়া 0.537, সহ 0.262

Shape-গুলো মুখে বলতে পারাই এই অংশের আসল লক্ষ্য। Q আর K-এর শেষ অক্ষ অবশ্যই সমান (d_k), নইলে dot product হয় না। V-এর শেষ অক্ষ (d_v) আলাদা হতে পারে, কারণ V-কে শুধু A দিয়ে মেশানো হয়; বাস্তবে সাধারণত তিনটিই সমান রাখা হয়। আর Q-এর T আর K-এর T সমান হতে হবে, এমন শর্ত নেই। Decoder-এর ১০ token encoder-এর ৩০ token-এর দিকে তাকালে A হয় (B, 10, 30); এই cross-attention পর্ব ৩৫-এ (Transformer block) আসবে। Self-attention হলো সেই বিশেষ ক্ষেত্র, যেখানে দুটি T একই।

৩ · ৩টি token হাতে হিসাব

কাগজ-কলম নিন। ৩টি token, d = ২, ধরুন “ফোন”, “ভালো”, “ছিল”. সংখ্যাগুলো হাতে হিসাবের সুবিধার জন্য বানানো। W_Q = [[1,0],[1,1]], W_K = [[1,1],[0,1]], W_V = [[1,0],[0,2]].

একটি জিনিস লক্ষ করুন। ওজন ঠিক করে Q আর K, কিন্তু মিশ্রণ হয় V দিয়ে। পরীক্ষায় হাতে করার সময় এই তিনটি গুলিয়ে ফেলাই সবচেয়ে সাধারণ ভুল।

A-র একটি সারি পড়তে শিখুন। সারি ৩ বলছে, “ছিল”-এর নতুন vector তৈরি হবে “ফোন”-এর value-র ৩০.৫৭%, “ভালো”-র ৭.৪৩% আর নিজের ৬২% নিয়ে। এখানে W আমরা বানিয়েছি, তাই এই সংখ্যার কোনো ভাষাগত অর্থ নেই। Training-এ W বদলায়, আর সেই সঙ্গে A-ও। হিসাব মিলিয়ে নেওয়ার সহজ উপায়: প্রতিটি সারির যোগফল ১ কি না, আর Z-এর প্রতিটি মান V-এর সেই কলামের সর্বনিম্ন ও সর্বোচ্চ মানের ভেতরে কি না। মিশ্রণ বলে এর বাইরে যেতে পারে না।

৪ · কেন √d_k দিয়ে ভাগ

Softmax score-এর পার্থক্য দেখে। পার্থক্য খুব বড় হলে একটি token প্রায় ১ পায়, বাকিরা প্রায় ০. একে বলে softmax saturation. তখন gradient প্রায় শূন্য, কারণ ∂p_i/∂z_i = p_i(1 − p_i). Gradient-এর এই হিসাবের ভিত্তি আছে মেশিন লার্নিং কোর্সের পর্ব ১১-এ। ধাপে ধাপে দেখুন।

এটি প্রমাণের ধারণাটুকু। এবার মাপি: প্রতিটি d_k-এ ৩০০টি query আর ৮টি করে key, উপাদানগুলো standard normal.

import numpy as np, matplotlib.pyplot as plt
rng = np.random.default_rng(42)
ds = [4, 16, 64, 256, 1024]
def softmax(z):
    e = np.exp(z - z.max(-1, keepdims=True))
    return e / e.sum(-1, keepdims=True)
rows = []
for dk in ds:
    q = rng.normal(size=(300, dk))          # ৩০০টি query
    k = rng.normal(size=(300, 8, dk))       # প্রতিটির জন্য ৮টি key
    s = np.einsum("nd,ntd->nt", q, k)       # dot product, (300, 8)
    res = [dk, s.var()]
    for z in (s, s / np.sqrt(dk)):          # scaling ছাড়া, তারপর সহ
        p = softmax(z)
        g = (p * (1 - p)).max(-1)           # সবচেয়ে বড় ∂p_i/∂z_i
        res += [z.var(), p.max(-1).mean(), g.mean()]
    rows.append(res)
print("  d_k   var(raw) var(scaled) maxp(raw) maxp(scaled) grad(raw) grad(scaled)")
for r in rows:
    print(f"{r[0]:5d} {r[2]:9.1f} {r[5]:10.2f} {r[3]:9.3f} {r[6]:12.3f} {r[4]:9.4f} {r[7]:11.3f}")
r = np.array(rows)
fig, ax = plt.subplots(1, 2, figsize=(8, 3.2))
ax[0].loglog(r[:, 0], r[:, 2], "o-", label="raw q·k")
ax[0].loglog(r[:, 0], r[:, 5], "s-", label="q·k / sqrt(d_k)")
ax[0].set_xlabel("d_k"); ax[0].set_ylabel("variance of score"); ax[0].legend()
ax[1].semilogx(r[:, 0], r[:, 3], "o-", label="raw q·k")
ax[1].semilogx(r[:, 0], r[:, 6], "s-", label="q·k / sqrt(d_k)")
ax[1].axhline(1 / 8, ls=":", c="gray"); ax[1].set_ylim(0, 1.05)
ax[1].set_xlabel("d_k"); ax[1].set_ylabel("mean largest weight (8 keys)"); ax[1].legend()
fig.subplots_adjust(wspace=0.32, bottom=0.18); plt.show()

  d_k   var(raw) var(scaled) maxp(raw) maxp(scaled) grad(raw) grad(scaled)
    4       3.8       0.94     0.520        0.336    0.2095       0.206
   16      16.1       1.01     0.742        0.350    0.1532       0.215
   64      67.7       1.06     0.875        0.367    0.0838       0.219
  256     261.6       1.02     0.938        0.357    0.0427       0.218
 1024     991.9       0.97     0.971        0.357    0.0214       0.215

শেখানোর সময় কি সত্যিই ফারাক পড়ে? — সত্যিকারের run

Variance মাপা এক কথা, training আরেক কথা। তাই PyTorch-এ (torch 2.15.0, CPU-তে) একটি ছোট খোঁজার কাজ চালিয়েছি। বাজারঘরের ৮টি পণ্য-কোড আর প্রতিটির একটি দামের স্তর দেওয়া আছে, ডেটা seed দিয়ে বানানো। একটি কোড দিলে model-কে তার দামের স্তর বলতে হবে (১০টির মধ্যে একটি)। এখানে query আসে আলাদা কোড থেকে, তাই গঠনটি cross-attention-এর মতো, কিন্তু scaling-এর হিসাব একই। মূল অংশ, যেমন চলেছে:

def forward(self, k, v, q):
    S = self.wq(self.ek(q))[:, None] @ self.wk(self.ek(k)).transpose(1, 2) / self.scale
    A = S.softmax(-1)                 # A is a middle step, not the output
    return self.out((A @ self.ev(v)).squeeze(1)), A

self.scale হয় √d, নয় ১. AdamW, learning rate 0.001, batch ১২৮, ৬০০ ধাপ, ১,০০০টি আলাদা পরীক্ষার উদাহরণ; accuracy দেখা হয়েছে প্রতি ২০ ধাপে। Embedding-এর শুরুর sd ১. সময়ের সীমায় থেমে যাওয়ায় d = ১০২৪-এর শেষ চারটি সারি একই script-এর আলাদা run থেকে নেওয়া; প্রতিটি সারি নিজের seed থেকে নতুন শুরু করে।

   d  scale seed  maxA0  gradW0  acc600  step90
  16   True    0  0.180   0.056   1.000     180
  16   True    1  0.195   0.069   1.000     160
  16   True    2  0.174   0.055   1.000     200
  16  False    0  0.379   0.278   1.000     180
  16  False    1  0.435   0.315   1.000     140
  16  False    2  0.356   0.202   1.000     200
 256   True    0  0.195   0.199   1.000      20
 256   True    1  0.193   0.251   1.000      20
 256   True    2  0.194   0.227   1.000      20
 256  False    0  0.811   2.456   0.987      80
 256  False    1  0.831   2.996   0.942     260
 256  False    2  0.807   3.173   0.952     220
1024   True    0  0.191   0.482   0.993     100
1024   True    1  0.197   0.363   0.940      20
1024   True    2  0.188   0.397   1.000      20
1024  False    0  0.922   7.469   0.524    None
1024  False    1  0.926   6.214   0.564    None
1024  False    2  0.892   7.519   0.608    None

maxA0 শুরুতে সবচেয়ে বড় weight-এর গড়, step90 কোন ধাপে accuracy প্রথম ০.৯ ছুঁল (None মানে ৬০০ ধাপে ছোঁয়নি)। d = ১৬-এ প্রায় কোনো পার্থক্য নেই। d = ২৫৬-এ scale করা model ২০ ধাপেই ০.৯, ভাগ ছাড়া model-এর লাগল ৮০ থেকে ২৬০. আর d = ১০২৪-এ ভাগ ছাড়া তিনটি run ৬০০ ধাপ পরেও ০.৫২ থেকে ০.৬১. শুরুতেই weight প্রায় ০.৯ একটি token-এ, যা প্রায় সবসময় ভুল token.

দুটি সতর্ক কথা। এক, প্রথম ধাপের gradient (gradW0) ভাগ ছাড়া model-এই বড়, কারণ score-ই বড়। অর্থাৎ এই run “gradient ছোট বলেই ধীর” কথাটি সরাসরি প্রমাণ করে না; ধীর হওয়ার কারণ আমরা আলাদা করে মাপিনি। দুই, একই ভাগ-না-করা ধারণা আরেকটি কাজে কোনো ক্ষতিই করেনি। সেখানে weight নিজেই output, আর loss সরাসরি log p-এর ওপর: “মিলে যাওয়া কোডটি কোথায়” খোঁজা। d = ১০২৪-এ ভাগ ছাড়া শুরুর loss ১৫–১৭ (ln 9 ≈ ২.২-এর জায়গায়), তবু তিনটি seed-এই প্রথম মাপে, ১০ ধাপে, accuracy ০.৯৫ পার। ফলে √d_k কোনো সার্বজনীন উদ্ধারকর্তা নয়: যেখানে A মাঝের ধাপ, সেখানে এই একটি ছোট পরীক্ষায় এর ফল স্পষ্ট দেখা গেল। প্রতিটি সংখ্যা এই ছোট কাজের, তিনটি seed-এর; সাধারণ নিয়ম নয়।

৫ · Causal আর padding mask — softmax-এর আগে −∞

সবাইকে সবসময় দেখতে দেওয়া যায় না। দুটি কারণ আলাদা। এক, পরের token অনুমান করার সময় ভবিষ্যৎ দেখা নকল করার সমান। পরীক্ষার হলে উত্তরপত্র দেখে উত্তর লেখার মতো। এটি আটকায় causal mask. দুই, batch-এ ছোট review-এর শেষে pad আছে (পর্ব ১১, পর্ব ২৭), সেগুলোর কোনো অর্থ নেই। এটি আটকায় padding mask.

কেন softmax-এর আগে? −∞-এর e-ঘাত ঠিক ০, তাই নিষিদ্ধ ঘর weight পায় না, আর বাকিদের যোগফল আপনাআপনি ১ থাকে। Softmax-এর পরে শূন্য করলে সারির যোগফল ১-এর নিচে নেমে যেত। Causal mask বানাতে উপরের ত্রিভুজ, মানে j > i: np.triu(..., k=1). Shape মেলাতে broadcasting: causal (T, T) আর pad (B, 1, T) মিলে (B, T, T).

দুটি mask-এর আকারের পার্থক্যটি খেয়াল করুন। Causal mask সব উদাহরণের জন্য এক, কারণ এটি শুধু অবস্থানের ওপর নির্ভর করে। Padding mask প্রতিটি উদাহরণের জন্য আলাদা, কারণ দৈর্ঘ্য আলাদা। আর pad আটকানো হয় কলামে, অর্থাৎ key হিসেবে: pad-কে কেউ দেখবে না। Encoder-এ শুধু padding mask লাগে। Decoder-এর self-attention-এ দুটিই। Cross-attention-এ encoder-এর padding mask, কারণ key তখন encoder-এর token.

হাতে একবার করুন। অংশ ৩-এর S-এর দ্বিতীয় সারি ছিল [১.৪১৪২, ০.৭০৭১, ২.১২১৩]. Causal mask তৃতীয় ঘরে −∞ বসায়। e-ঘাত হয় ৪.১১৩৩, ২.০২৮১, ০; যোগফল ৬.১৪১৪. তাই সারিটি [০.৬৬৯৮, ০.৩৩০২, ০]. Mask ছাড়া ছিল [০.২৮৪০, ০.১৪০০, ০.৫৭৬০]. লক্ষ করুন, বাকি দুই ঘরের অনুপাত একই থাকে (প্রায় ২.০৩ : ১), শুধু যোগফল আবার ১ হয়। Softmax-এর আগে mask দেওয়ার সুবিধা ঠিক এটিই।

import numpy as np
np.set_printoptions(precision=2, suppress=True, linewidth=60)
rng = np.random.default_rng(7)
T, d = 5, 8
X = rng.normal(size=(2, T, d))                # batch: দৈর্ঘ্য ৫ আর ৩ (শেষে ২টি pad)
pad = np.array([[0, 0, 0, 0, 0], [0, 0, 0, 1, 1]], dtype=bool)
causal = np.triu(np.ones((T, T), dtype=bool), k=1)
block = causal[None] | pad[:, None, :]        # (1,T,T) | (2,1,T) → (2,T,T)
print("causal:\n", causal.astype(int))
print("pad-সহ, উদাহরণ ১:\n", block[1].astype(int))

def attend(X, block):
    S = X @ X.transpose(0, 2, 1) / np.sqrt(d)
    S = np.where(block, -np.inf, S)           # softmax-এর আগে −∞
    e = np.exp(S - S.max(-1, keepdims=True))
    A = e / e.sum(-1, keepdims=True)
    return A @ X, A

Z, A = attend(X, block)
print("A, উদাহরণ ১:\n", A[1])
print("সারির যোগফল:", A[1].sum(-1))
Y = X.copy(); Y[1, 3:] = 99.0                 # শুধু pad-এর ভেতরের মান বদল
Zp, _ = attend(Y, block)
print("pad বদলালে আসল token-এ পরিবর্তন:", np.abs(Zp[1, :3] - Z[1, :3]).max())
Y = X.copy(); Y[0, 4] = 99.0                  # উদাহরণ ০-এর শেষ token বদল
Zf, _ = attend(Y, causal[None] | pad[:, None, :])
print("ভবিষ্যৎ বদলালে আগের token-এ পরিবর্তন:", np.abs(Zf[0, :4] - Z[0, :4]).max())
left = np.array([[1, 1, 0, 0, 0], [0, 0, 0, 0, 0]], dtype=bool)  # বাঁয়ে pad
_, B = attend(X, causal[None] | left[:, None, :])
print("বাঁয়ে pad + causal, A-এর সারি ০:", B[0, 0], "| NaN সারি:", np.isnan(B[0]).any(-1).astype(int))

causal:
 [[0 1 1 1 1]
 [0 0 1 1 1]
 [0 0 0 1 1]
 [0 0 0 0 1]
 [0 0 0 0 0]]
pad-সহ, উদাহরণ ১:
 [[0 1 1 1 1]
 [0 0 1 1 1]
 [0 0 0 1 1]
 [0 0 0 1 1]
 [0 0 0 1 1]]
A, উদাহরণ ১:
 [[1.   0.   0.   0.   0.  ]
 [0.04 0.96 0.   0.   0.  ]
 [0.09 0.27 0.64 0.   0.  ]
 [0.16 0.1  0.74 0.   0.  ]
 [0.2  0.5  0.29 0.   0.  ]]
সারির যোগফল: [1. 1. 1. 1. 1.]
pad বদলালে আসল token-এ পরিবর্তন: 0.0
ভবিষ্যৎ বদলালে আগের token-এ পরিবর্তন: 0.0
বাঁয়ে pad + causal, A-এর সারি ০: [nan nan nan nan nan] | NaN সারি: [1 1 0 0 0]

৬ · Multi-head — কাটা, আলাদা attention, জোড়া

একটি attention একটি সারিতে একটিই ওজনের ছক দেয়। Multi-head attention-এ h-টি ছক। Vector d-কে h ভাগে কাটা হয়, d_k = d/h. প্রতি ভাগে আলাদা attention, তারপর জোড়া, শেষে আরেকটি projection W_O.

২০১৭ সালের প্রবন্ধের মূল মডেলের মান ছিল d = ৫১২, h = ৮. সেই মান ধরে প্রতিটি shape গুনে দেখি।

এবার একই ছাঁচ NumPy-তে, ছোট মানে (d = ৬৪, h = ৮), shape পরীক্ষা-সহ। কোডের শেষ তিন লাইন পরের অংশের প্রশ্নের জন্য।

import numpy as np
rng = np.random.default_rng(42)
B, T, d, h = 2, 6, 64, 8
dk = d // h
Wq, Wk, Wv, Wo = rng.normal(size=(4, d, d)) / np.sqrt(d)

def split(x):                                 # (B,T,d) → (B,h,T,dk)
    return x.reshape(x.shape[0], x.shape[1], h, dk).transpose(0, 2, 1, 3)

def mha(X, show=False):
    Q, K, V = split(X @ Wq), split(X @ Wk), split(X @ Wv)
    S = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(dk)
    e = np.exp(S - S.max(-1, keepdims=True))
    A = e / e.sum(-1, keepdims=True)
    H = A @ V                                 # প্রতিটি head আলাদা
    C = H.transpose(0, 2, 1, 3).reshape(X.shape)   # concat
    if show:
        for n, a in [("X", X), ("Q", Q), ("S", S), ("A", A), ("H", H), ("C", C)]:
            print(f"{n}: {a.shape}")
    return C @ Wo

X = rng.normal(size=(B, T, d))
Z = mha(X, show=True)
assert Z.shape == X.shape
Q = split(X @ Wq); manual = X[0, 3] @ Wq      # split কি শুধু কাটা? token ৩, head ২
print("Q[0,2,3] = projection-এর ১৬–২৩ ঘর?", np.allclose(Q[0, 2, 3], manual[16:24]))

print("params (bias ছাড়া) 4d² =", 4 * d * d, "| h = 8-এ head-প্রতি dk =", dk)
p = rng.permutation(T)                        # token-এর ক্রম এলোমেলো
print("perm:", p)
print("Z(PX) বনাম P·Z(X), max|diff|:", np.abs(mha(X[:, p]) - Z[:, p]).max())
pos = rng.normal(size=(T, d))                 # প্রতিটি ঘরের একটি নিজস্ব vector
base = mha(X + pos)
print("অবস্থান যোগ করলে, max|diff|:", np.abs(mha(X[:, p] + pos) - base[:, p]).max().round(3))

X: (2, 6, 64)
Q: (2, 8, 6, 8)
S: (2, 8, 6, 6)
A: (2, 8, 6, 6)
H: (2, 8, 6, 8)
C: (2, 6, 64)
Q[0,2,3] = projection-এর ১৬–২৩ ঘর? True
params (bias ছাড়া) 4d² = 16384 | h = 8-এ head-প্রতি dk = 8
perm: [0 2 5 4 3 1]
Z(PX) বনাম P·Z(X), max|diff|: 6.661338147750939e-16
অবস্থান যোগ করলে, max|diff|: 2.867

এই সূত্রের দুটি ব্যবহারিক ফল আছে। এক, h অবশ্যই d-কে নিঃশেষে ভাগ করবে। ৫১২ ঘরকে ৭ ভাগে কাটা যায় না। দুই, h বাড়লে d_k কমে: d = ৫১২-এ h = ১৬ দিলে head-প্রতি মাত্র ৩২ ঘর। W_O কেন লাগে? Concat-এর পর head-গুলোর ঘর পাশাপাশি বসে থাকে, মেশে না। W_O প্রতিটি token-এর ভেতরে সেগুলো মেশায়, যাতে এক head-এর ফল পরের স্তরের সব head-এর কাজে আসে। Einsum-এ score লেখা যায় np.einsum("bhtd,bhsd->bhts", Q, K): অক্ষের নাম দেখেই বোঝা যায়, d মিলিয়ে যোগ হচ্ছে, আর t, s আলাদা থাকছে।

৭ · ক্রম বদলালে কী হয় — permutation equivariance

আগের আউটপুটের শেষ তিন লাইন দেখুন। Token-এর ক্রম এলোমেলো করে চালালে output-ও ঠিক একইভাবে এলোমেলো হলো; পার্থক্য শুধু floating-point-এর ত্রুটি। একে বলে permutation equivariance. “ডেলিভারি দেরি হয়নি, দাম ভালো” আর “দাম দেরি হয়নি, ডেলিভারি ভালো”-তে একই শব্দ একই vector পেলে “দেরি”-র output দুই বাক্যে হুবহু এক। অর্থ বদলাল, কিন্তু model টের পেল না।

কারণ হাতে দেখানো যায়। ধরুন P একটি permutation matrix, আর PᵀP = I. তাহলে:

X → PX ⇒ Q, K, V → PQ, PK, PV

S → P S Pᵀ ⇒ A → P A Pᵀ

Z → (P A Pᵀ)(P V) = P A V = P Z

Softmax সারি ধরে চলে, আর সারির ভেতরের ক্রম বদলালে softmax-এর মান একইভাবে সরে যায়, তাই দ্বিতীয় লাইন ঠিক আছে। শেষ লাইনে PᵀP = I কাজ করে।

হাতের হিসাবেও দেখা যায়। অংশ ৩-এর উদাহরণে “ফোন” আর “ভালো” অদলবদল করলে X = [[0,1],[1,0],[1,1]] হয়। Q, K, V-এর প্রথম দুই সারিও অদলবদল হয়, A-র প্রথম দুই সারি আর প্রথম দুই কলামও। তাই নতুন output হয় Z₂, Z₁, Z₃: [০.৮৬০০, ১.৪৩২০], [০.৮০২২, ১.১৯৭৮], [০.৯২৫৭, ১.৩৮৮৬]. মান একই, শুধু জায়গা বদল।

কোনো ঘরে অবস্থানের একটি নিজস্ব vector যোগ করলে (শেষ লাইন) এই সমতা ভেঙে যায়: পার্থক্য আর ত্রুটির মাপে থাকে না। সে কারণে পরের পর্ব ৩৪ (Positional encoding — sinusoidal, learned আর RoPE) দরকার। আর causal mask-ও অবস্থান-নির্ভর, তাই সেটি থাকলেও এই সমতা থাকে না; অনুশীলনে মাপবেন।

৮ · সত্যিকারের run — নিজের module বনাম PyTorch

PyTorch পাতার ভেতরে চলে না, তাই এটি ডেস্কটপের সত্যিকারের run: Python 3.12.13, লেখার সময়ের torch 2.15.0, CPU build. নিচের class-টি যেমন চলেছে তেমনই।

import math, torch
import torch.nn as nn, torch.nn.functional as F

class MyMHA(nn.Module):
    def __init__(self, d, h, bias=True):
        super().__init__()
        assert d % h == 0, "d must divide evenly by h"
        self.d, self.h, self.dk = d, h, d // h
        self.wq, self.wk, self.wv, self.wo = (nn.Linear(d, d, bias=bias) for _ in range(4))

    def split(self, x):                      # (B,T,d) -> (B,h,T,dk)
        B, T, _ = x.shape
        return x.view(B, T, self.h, self.dk).transpose(1, 2)

    def forward(self, x, attn_mask=None, key_padding_mask=None, use_sdpa=False):
        B, T, _ = x.shape
        Q, K, V = self.split(self.wq(x)), self.split(self.wk(x)), self.split(self.wv(x))
        block = None                         # True = not allowed to look
        if attn_mask is not None:
            block = attn_mask[None, None]    # (1,1,T,T)
        if key_padding_mask is not None:
            pad = key_padding_mask[:, None, None, :]   # (B,1,1,T)
            block = pad if block is None else block | pad
        if use_sdpa:                         # note the ~ : True = may look
            m = None if block is None else ~block
            Z, A = F.scaled_dot_product_attention(Q, K, V, attn_mask=m), None
        else:
            S = Q @ K.transpose(-2, -1) / math.sqrt(self.dk)
            if block is not None:
                S = S.masked_fill(block, float("-inf"))
            A = S.softmax(-1)
            Z = A @ V                        # (B,h,T,dk)
        Z = Z.transpose(1, 2).reshape(B, T, self.d)
        return self.wo(Z), A

এবার nn.MultiheadAttention-এ একই weight বসিয়ে মেলাই। এই module তিনটি W একসঙ্গে একটি parameter-এ রাখে। আউটপুটের প্রথম লাইনে দেখবেন in_proj_weight-এর shape (96, 32): ৩ × ৩২ সারি, যা ভেতরে তিন ভাগে কাটা হয়। তাই আমাদের তিনটি weight torch.cat দিয়ে সারি বরাবর জোড়া হয়েছে, Q, K, V ক্রমে। average_attn_weights=False না দিলে head-গুলোর গড় ফেরত আসে, তখন shape (B, T, T). পরীক্ষার script-এর মূল অংশ সংক্ষেপ করে দেখানো হলো; পুরো script-এ আরও কয়েকটি একই ধরনের তুলনা আছে, আউটপুটে সেগুলোও এসেছে।

B, T, d, h = 2, 5, 32, 4                     # torch.manual_seed(42)
x = torch.randn(B, T, d)
pad = torch.tensor([[0,0,0,0,0],[0,0,0,1,1]], dtype=torch.bool)
causal = torch.triu(torch.ones(T, T, dtype=torch.bool), 1)
mine = MyMHA(d, h).eval()
ref = nn.MultiheadAttention(d, h, dropout=0.0, batch_first=True).eval()
with torch.no_grad():
    ref.in_proj_weight.copy_(torch.cat([mine.wq.weight, mine.wk.weight, mine.wv.weight], 0))
    ref.in_proj_bias.copy_(torch.cat([mine.wq.bias, mine.wk.bias, mine.wv.bias], 0))
    ref.out_proj.weight.copy_(mine.wo.weight)
    ref.out_proj.bias.copy_(mine.wo.bias)
    a, A = mine(x, causal, pad)
    b, _ = mine(x, causal, pad, use_sdpa=True)
    c, W = ref(x, x, x, attn_mask=causal, key_padding_mask=pad,
               need_weights=True, average_attn_weights=False)

~/bazarghar-dl$ uv run python check_match.py
[('in_proj_weight', (96, 32)), ('in_proj_bias', (96,)), ('out_proj.weight', (32, 32)), ('out_proj.bias', (32,))]
output shape (2, 5, 32) weights shape (2, 4, 5, 5)
mine vs my sdpa path      max|diff| = 8.94e-08
mine vs nn.MHA (weights)  max|diff| = 1.19e-07
mine vs nn.MHA (no wts)   max|diff| = 1.19e-07
attn weights mine vs nn   max|diff| = 5.96e-08
sdpa right vs wrong mask polarity max|diff| = 1.61, NaN in wrong: False
pad content changed: real tokens max|diff| = 0.00e+00
last token changed, causal: earlier max|diff| = 0.00e+00, last = 69.7
last token changed, no mask: earlier max|diff| = 105.1
left pad + causal, NaN per position (mine) : [[1, 1, 0, 0, 0], [0, 0, 0, 0, 0]]
left pad + causal, NaN per position (nn)   : [[1, 1, 0, 0, 0], [0, 0, 0, 0, 0]]
left pad + causal, NaN per position (sdpa) : [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]

পার্থক্য ১০^−৭-এর মাপে: float32-এর গোলকরণের ত্রুটি। তিনটি পথ একই গণিত করছে। কিন্তু দুটি লাইন আলাদা করে পড়ুন।

Mask-এর অর্থ উল্টো। একই boolean mask দুই API-তে দুই রকম মানে দেয়। ভুল দিকে দিলে কোনো error নেই, NaN-ও নেই; এখানে output ১.৬১ পর্যন্ত বদলে গেছে, নীরবে। টেবিলটি মুখস্থ না করে প্রতিবার docs দেখাই নিরাপদ।

NaN সব পথে এক নয়। বাঁয়ে pad আর causal একসঙ্গে দিলে আমাদের module আর nn.MultiheadAttention-এ NaN, কিন্তু F.scaled_dot_product_attention-এ NaN নেই (পুরো আটকানো সারির output সেখানে শূন্য এসেছে)। একই nn.MultiheadAttention-এর ভেতরেও পথ ভেদে ফল আলাদা:

eval need_weights True [[1, 1, 0, 0, 0], [0, 0, 0, 0, 0]]
eval need_weights False [[1, 1, 0, 0, 0], [0, 0, 0, 0, 0]]
train need_weights True [[1, 1, 0, 0, 0], [0, 0, 0, 0, 0]]
train need_weights False [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]
sdpa row 0 first head, all zero? 0.0 NaN False

মানে এই version-এ train-এ need_weights=False ঠিক চলল, অথচ eval-এ একই mask-এ NaN. ভেতরের কোডের পথ সম্ভবত আলাদা; অন্য version-এ এই ফল নাও মিলতে পারে। শিক্ষা: pad সারির ওপর ভরসা না করে নিজে পরীক্ষা করুন, আর সম্ভব হলে এমন mask দিন যাতে কোনো সারি পুরো আটকে না যায়।

তাহলে কোনটি ব্যবহার করবেন? শেখার জন্য নিজে লেখা ভালো, কারণ প্রতিটি shape দেখা যায়। নিজের model-এ সাধারণত F.scaled_dot_product_attention সুবিধার: হিসাবটি এক লাইনে, আর PyTorch নিজে দ্রুত পথ বেছে নিতে পারে, তবে এটি weight ফেরত দেয় না; weight দেখতে চাইলে নিজের পথ লাগে। nn.MultiheadAttention-এ সব একসঙ্গে পাওয়া যায়, কিন্তু default, mask আর ভেতরের পথগুলো মনে রাখতে হয়। যেটিই নিন, উপরের মতো ছোট একটি তুলনা আর leak test সঙ্গে রাখুন।

Parameter গোনা

~/bazarghar-dl$ uv run python count_params.py
d= 512 h= 1 dk= 512  mine=1,050,624  nn.MHA=1,050,624  formula 4d^2+4d=1,050,624
d= 512 h= 2 dk= 256  mine=1,050,624  nn.MHA=1,050,624  formula 4d^2+4d=1,050,624
d= 512 h= 8 dk=  64  mine=1,050,624  nn.MHA=1,050,624  formula 4d^2+4d=1,050,624
d= 512 h=16 dk=  32  mine=1,050,624  nn.MHA=1,050,624  formula 4d^2+4d=1,050,624
d=  64 h= 8 dk=   8  mine=   16,640  nn.MHA=   16,640  formula 4d^2+4d=   16,640
bias=False, d=512, h=8: 1,048,576 = 4d^2 = 1,048,576
  wq: 262,656
  wk: 262,656
  wv: 262,656
  wo: 262,656
nn.MHA(512, 7) -> embed_dim must be divisible by num_heads

h = ১ থেকে ১৬, parameter এক। Attention-এর parameter T-এর ওপরও নির্ভর করে না: ১০ শব্দের বাক্য আর ১,০০০ শব্দের লেখায় একই W. আর d দ্বিগুণ করলে parameter প্রায় চার গুণ। ৫১২ ঘরকে ৭ ভাগে কাটা যায় না, তাই PyTorch আগেই থামিয়ে দেয়।

৯ · খরচ — O(T²) মেপে দেখা

Parameter না বাড়লেও সময় বাড়ে। S-এ T × T ঘর, তাই গুণের কাজ T²·d, আর A-র মেমরি h·T² সংখ্যা। অন্যদিকে W দিয়ে projection T-এর সঙ্গে সমানুপাতিক। Big-O-এর ধারণা আছে DSA কোর্সের পর্ব ৪-এ। মাপা হলো CPU-তে, ১টি thread-এ, d = ২৫৬, h = ৮, B = ১; প্রতিটি T-তে ১৫ বারের দ্রুততমটি।

~/bazarghar-dl$ uv run python timing.py
torch 2.15.0+cpu, CPU, 1 thread, d=256, h=8, B=1, fastest of 15
    T        ms  ratio  proj ms    A MB
   64      0.62            0.29     0.1
  128      1.47   2.38     0.61     0.5
  256      3.53   2.40     1.13     2.0
  512     10.24   2.90     2.10     8.0
 1024     33.43   3.27     4.11    32.0
 2048    137.48   4.11     8.15   128.0
 4096    755.21   5.49    16.23   512.0
log-log slope, T=1024..4096: 2.25
log-log slope, T=64..256: 1.26

এটি একটি যন্ত্রে একটি run, তাই ধারাটি দেখুন, সংখ্যাটি নয়। ছোট T-তে T দ্বিগুণে সময় প্রায় ২.৪ গুণ, কারণ তখন projection আর অন্য খরচের ভাগ বড়, আর সেগুলো T-এর সঙ্গে সমানুপাতিক (proj ms কলামটি প্রায় দ্বিগুণ হচ্ছে)। বড় T-তে T²-এর পদ জেতে, অনুপাত ৪-এর আশেপাশে আসে। ৫.৪৯ যে ৪-এর বেশি, সেটি হয়তো cache-এ না ধরার খরচ; এটি আমরা আলাদা করে মাপিনি। শেষ কলাম: T = ৪,০৯৬-এ শুধু A-র জন্যই ৫১২ MB (২^২০ বাইটে), একটি উদাহরণে, একটি স্তরে।

হাতে হিসাব করতে পারলে অবাক হতে হয় না। A-র মেমরি = B × h × T² × প্রতি সংখ্যার বাইট। B = ৩২, h = ৮, T = ৫১২, float32 ধরলে ৩২ × ৮ × ২,৬২,১৪৪ × ৪ = ২৬,৮৪,৩৫,৪৫৬ বাইট, ঠিক ২৫৬ MiB, এক স্তরে। ১২টি স্তরে ৩ GiB-এর মতো, যদি সব একসঙ্গে রাখা হয়; backward-এর জন্য autograd অনেক মধ্যবর্তী মান রাখে বলে training-এ এমন হিসাব গুরুত্বপূর্ণ (পর্ব ১২)। T অর্ধেক করলে এই মেমরি এক-চতুর্থাংশ, B অর্ধেক করলে শুধু অর্ধেক। তাই batch ছাঁটার চেয়ে লম্বা sequence ছাঁটা বা bucketing অনেক সময় বেশি কাজের।

১০ · AI-এর যুগে — attention-এর কোড যাচাই

AI assistant multi-head attention ঝটপট লিখে দেয়, আর বেশির ভাগ ভুলে কোনো error আসে না। ভুল shape ঠিকই মিলে যায়, মান শুধু ভুল হয়।

AI-এর দেওয়া attention হাতে পেলে — যাচাইয়ের চেকলিস্ট

• ছোট batch-এ প্রতিটি ধাপের shape ছাপিয়ে নিজের হিসাবের সঙ্গে মিলিয়েছি
• A.sum(-1) সব আসল সারিতে ১, আর একটি সারি হাতে হিসাব করে মিলিয়েছি
• Pad আর ভবিষ্যৎ বদলে leak test করেছি; মান ঠিক অপরিবর্তিত
• নিজের কোড weight কপি করে nn.MultiheadAttention-এর সঙ্গে মিলিয়েছি, ১০^−৬-এর নিচে
• Mask-এর অর্থ আমার version-এর docs-এ পড়েছি; NaN আছে কি না train আর eval দুটিতেই দেখেছি
• model.eval(), torch.no_grad() ঠিক জায়গায়; একটি ছোট batch overfit করে দেখেছি
• ব্যক্তিগত ডেটা বা key AI-কে দিইনি; ফল এক run-এ নয়, কয়েকটি seed-এ দেখেছি