Transformer block — residual, LayerNorm আর feed-forward

Query, key আর value দিয়ে attention এসেছে পর্ব ৩২ (Attention — query, key, value আর alignment)-এ। তার matrix রূপ আর multi-head-এর হিসাব পর্ব ৩৩ (Self-attention আর multi-head attention)-এ, আর order-এর তথ্য যোগ করা পর্ব ৩৪ (Positional encoding)-এ। সেগুলো এখানে আবার শেখানো হবে না। LayerNorm-এর সূত্র আছে পর্ব ১৪-এ, warmup পর্ব ১৫-এ, dropout পর্ব ১৬-এ, residual-এর প্রথম গণিত পর্ব ১৭-এ, GELU পর্ব ৩-এ। Block-এর একটি সরল ছবি দেখেছেন AI-এর ভিত্তি কোর্সের পর্ব ৭০-এ। এবার সেই ছবির প্রতিটি তির আর বাক্স খুলে দেখা হবে।

১ · একটি block-এ চারটি কাজ — কে কী করে

একটি Transformer block ঠিক দুটি sublayer দিয়ে তৈরি: multi-head attention আর একটি ছোট feed-forward network (FFN). প্রতিটির চারপাশে থাকে residual connection আর LayerNorm. প্রথম sublayer শুধু token-দের মধ্যে তথ্য সরায়; দ্বিতীয়টি প্রতিটি token-এ আলাদা হিসাব করে। একটি অন্যটির অভাব পূরণ করে।

আধুনিক pre-norm রূপে একটি encoder block এই দুটি লাইনেই শেষ:

x ← x + Attn(LN(x))
x ← x + FFN(LN(x))

RNN-এর সঙ্গে তুলনা করলে একটি সুবিধা চোখে পড়ে। RNN-এ প্রতিটি ধাপ আগের ধাপের জন্য অপেক্ষা করে (পর্ব ২৮), আর এখানে একটি block-এর সব token একসঙ্গে হিসাব হয়। দুই দূরের token-এর মধ্যে তথ্য যেতে অনেক ধাপ লাগে না, এক attention-ই যথেষ্ট। দাম হলো T² হিসাব, আর order নিজে থেকে আসে না, যা পর্ব ৩৪-এর বিষয়।

এখানে x-এর আকার (B, T, d): B হলো batch-এ কয়টি উদাহরণ, T কয়টি token, d প্রতিটি token-এর vector-এর মাপ। LN হলো LayerNorm, Attn হলো multi-head attention. লক্ষ্য করুন, আকার বদলায় না। এই একটি গুণের জন্যই block স্তূপের মতো সাজানো যায়: একটি block-এর output হুবহু পরেরটির input.

২ · ভেতরের তিনটি উপাদান — LayerNorm, FFN আর dropout

LayerNorm. প্রতিটি token-এর d মাপের vector নিজের গড় বিয়োগ করে, নিজের standard deviation দিয়ে ভাগ হয়, তারপর শেখা γ দিয়ে গুণ আর β যোগ হয়। Batch-এর অন্য উদাহরণ বা অন্য token এখানে আসে না। তাই ১ আকারের batch-এও এটি একইভাবে চলে, আর padding-এর দৈর্ঘ্য এর ওপর প্রভাব ফেলে না। এ কারণেই sequence-এ BatchNorm-এর বদলে এটি চলে। Parameter: γ আর β মিলে 2d. নতুন অনেক model এর বদলে RMSNorm নেয়, যাতে গড় বিয়োগ বা β নেই (পর্ব ১৪).

FFN. দুটি linear স্তর আর মাঝে একটি activation:

FFN(x) = W₂ · GELU(W₁x + b₁) + b₂

W₁ d থেকে r·d-তে নেয়, W₂ ফিরিয়ে আনে। r হলো expansion factor, প্রচলিত মান ৪. একে “position-wise” বলা হয়, কারণ T-এর প্রতিটি ঘরে একই weight আলাদাভাবে চলে। তাহলে T যাই হোক, FFN-এর parameter বদলায় না।

FFN কেন দরকার? Attention-এর output হলো V-গুলোর ভারিত যোগফল। Weight ঠিক থাকলে তা V-এর ওপর linear. Softmax-এর non-linearity শুধু কাকে কতটা নেওয়া হবে, সেখানে। একটি token-এর ভেতরের feature মিশিয়ে নতুন feature বানানোর কাজটি মূলত FFN করে। আর চওড়া মাঝের স্তরের কারণে parameter-এর বড় অংশও এখানেই থাকে; একটু পরে হিসাব করে দেখবেন।

Activation. মূল প্রবন্ধে ReLU, পরে GPT-ধারার অনেক model-এ GELU. আরেকটি রূপ SwiGLU (Shazeer ২০২০-এর GLU রূপগুলোর পরীক্ষা; রিপোর্ট অনুযায়ী Llama-ধারায় ব্যবহৃত): SwiGLU(x) = W₃(SiLU(W₁x) ⊙ W₂x), এখানে ⊙ মানে ঘর ধরে গুণ। Matrix তিনটি, তাই মাঝের মাপ 8d/3 রাখলে parameter প্রায় একই থাকে: 3·d·(8d/3) = 8d² আর 2·d·4d = 8d². কোনটি কোন কাজে ভালো, তা নিয়ে সর্বজনীন রায় নেই।

Dropout কোথায়. মূল প্রবন্ধে প্রতিটি sublayer-এর output-এ, residual যোগের ঠিক আগে, আর embedding ও positional encoding-এর যোগফলে। PyTorch-এর nn.MultiheadAttention-এ attention weight-এও dropout দেওয়া যায়। FFN-এর মাঝের স্তরের পরেও PyTorch-এর layer একটি বসায়। লক্ষ্য করুন, residual পথের ওপর কখনো dropout নেই: পুরোনো মান অক্ষত থাকে, বাদ পড়তে পারে শুধু নতুন অংশ। ছোট dataset-এ dropout সাহায্য করে; রিপোর্ট অনুযায়ী অনেক বড় LLM training-এ dropout খুব কম বা শূন্য রাখা হয়, কারণ সেখানে ডেটাই প্রচুর। এটি উদাহরণ ধরে নিন, নিয়ম নয়।

৩ · NumPy-তে একটি block — প্রতিটি shape ছাপিয়ে

এবার পাতার ভেতরেই একটি pre-norm block চালান। Model কিছু শেখেনি, weight এলোমেলো; দেখার বিষয় শুধু কোন array কোন আকারে কোথায় যায়। কোড ছোট রাখতে linear স্তরের bias বাদ, আর LayerNorm-এর শেখা γ = ১, β = ০ ধরা হয়েছে। প্রথমবার চালাতে Pyodide আর SciPy নামার জন্য কয়েক সেকেন্ড লাগতে পারে। পরের ব্লকটি এই ফাংশনগুলোই আবার ব্যবহার করবে।

import numpy as np
from scipy.special import erf
rng = np.random.default_rng(42)
B, T, d, h, r = 2, 5, 16, 4, 4
dk, TRACE = d // h, True
def show(name, a):
    if TRACE: print(f"{name:24}{a.shape}")
def ln(x):                                  # LayerNorm, γ=1, β=0
    return (x - x.mean(-1, keepdims=True)) / (x.std(-1, keepdims=True) + 1e-5)
def gelu(z): return 0.5 * z * (1 + erf(z / np.sqrt(2)))
def init(rng, s):                           # s = weight-এর standard deviation
    return dict(qkv=rng.normal(0, s, (d, 3*d)), o=rng.normal(0, s, (d, d)),
                f1=rng.normal(0, s, (d, r*d)), f2=rng.normal(0, s, (r*d, d)))
def mha(z, w):
    n, t = z.shape[:2]
    qkv = z @ w["qkv"];                                 show("z @ W_qkv", qkv)
    Q, K, V = [a.reshape(n, t, h, dk).transpose(0, 2, 1, 3)
               for a in np.split(qkv, 3, -1)];          show("Q, K, V (B,h,T,dk)", Q)
    s = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(dk);      show("scores (B,h,T,T)", s)
    a = np.exp(s - s.max(-1, keepdims=True))
    a = a / a.sum(-1, keepdims=True);                   show("weights (B,h,T,T)", a)
    c = (a @ V).transpose(0, 2, 1, 3).reshape(n, t, d); show("concat (B,T,d)", c)
    return c @ w["o"]
def block(x, w, res=True, norm=True):       # এক pre-norm block
    N = ln if norm else (lambda a: a)
    a = mha(N(x), w);                                   show("attention output", a)
    x = x + a if res else a;                            show("x + attn", x)
    u = gelu(N(x) @ w["f1"]);                           show("FFN hidden (B,T,r*d)", u)
    f = u @ w["f2"];                                    show("FFN output", f)
    return x + f if res else f
x = rng.normal(size=(B, T, d));                         show("input x", x)
y = block(x, init(rng, 1 / np.sqrt(d)));                show("output", y)
print("LN(x): mean", ln(x)[0, 0].mean().round(4), "std", ln(x)[0, 0].std().round(4))
print("|x| =", np.linalg.norm(x).round(2), " |y - x| =", np.linalg.norm(y - x).round(2))

input x                 (2, 5, 16)
z @ W_qkv               (2, 5, 48)
Q, K, V (B,h,T,dk)      (2, 4, 5, 4)
scores (B,h,T,T)        (2, 4, 5, 5)
weights (B,h,T,T)       (2, 4, 5, 5)
concat (B,T,d)          (2, 5, 16)
attention output        (2, 5, 16)
x + attn                (2, 5, 16)
FFN hidden (B,T,r*d)    (2, 5, 64)
FFN output              (2, 5, 16)
output                  (2, 5, 16)
LN(x): mean 0.0 std 1.0
|x| = 10.93  |y - x| = 19.93

৪ · PyTorch-এ নিজের block — nn.TransformerEncoderLayer-এর সঙ্গে মিলিয়ে

নিজের লেখা কোড ঠিক কি না, তা জানার সবচেয়ে কড়া উপায় হলো একই weight বসিয়ে একটি পরীক্ষিত implementation-এর সঙ্গে সংখ্যা মেলানো। নিচের class-গুলো পর্ব ৯-এর ধাঁচে লেখা। এটি সত্যিকারের run, CPU-তে; torch 2.15.0, Python 3.14, লেখার সময়ের স্থিতিশীল সংস্করণ।

import torch, torch.nn as nn

class MHA(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.h, self.dk = h, d // h
        self.qkv = nn.Linear(d, 3 * d)           # d থেকে 3d, পরে ৩ ভাগ
        self.out = nn.Linear(d, d)

    def forward(self, y, m=None, mask=None):     # m=None মানে self-attention
        m = y if m is None else m
        B, T, S, d = y.shape[0], y.shape[1], m.shape[1], y.shape[2]
        Q = self.qkv(y).chunk(3, -1)[0]          # Q আসে y থেকে
        _, K, V = self.qkv(m).chunk(3, -1)       # K, V আসে m থেকে
        Q = Q.view(B, T, self.h, self.dk).transpose(1, 2)
        K = K.view(B, S, self.h, self.dk).transpose(1, 2)
        V = V.view(B, S, self.h, self.dk).transpose(1, 2)
        s = Q @ K.transpose(-2, -1) / self.dk ** 0.5
        if mask is not None:                     # True মানে দেখা নিষেধ
            s = s.masked_fill(mask, float("-inf"))
        z = (s.softmax(-1) @ V).transpose(1, 2).reshape(B, T, d)
        return self.out(z)

class Block(nn.Module):                          # pre-norm encoder block
    def __init__(self, d, h, r=4):
        super().__init__()
        self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.att = MHA(d, h)
        self.ffn = nn.Sequential(nn.Linear(d, r * d), nn.GELU(),
                                 nn.Linear(r * d, d))

    def forward(self, x, mask=None):
        x = x + self.att(self.n1(x), mask=mask)
        return x + self.ffn(self.n2(x))

এবার PyTorch-এর layer বানিয়ে তার weight এই class-এ কপি করা। nn.MultiheadAttention তিনটি projection একটিই matrix in_proj_weight-এ রাখে, আকার (3d, d); ঠিক সে কারণে আমাদের qkv-ও এক টুকরো।

ref = nn.TransformerEncoderLayer(d, h, r * d, dropout=0.0, activation="gelu",
                                 batch_first=True, norm_first=True).eval()
mine = Block(d, h, r).eval()

def copy_mha(a, b):
    a.qkv.weight.copy_(b.in_proj_weight); a.qkv.bias.copy_(b.in_proj_bias)
    a.out.load_state_dict(b.out_proj.state_dict())

with torch.no_grad():
    copy_mha(mine.att, ref.self_attn)
    mine.n1.load_state_dict(ref.norm1.state_dict())
    mine.n2.load_state_dict(ref.norm2.state_dict())
    mine.ffn[0].load_state_dict(ref.linear1.state_dict())
    mine.ffn[2].load_state_dict(ref.linear2.state_dict())

~/bazarghar-dl$ uv run python match.py
torch 2.15.0+cpu | parameter: 49984 49984
shape: (8, 12, 64) (8, 12, 64)
max |diff| float32, mask ছাড়া : 0.000e+00
max |diff| float32, causal    : 4.768e-07
max |diff| float64, mask ছাড়া : 0.000e+00
max |diff| float64, causal    : 8.882e-16
decoder: parameter 66752 66752, max |diff| 4.768e-07
dropout=0.1, train(): একই input দুই বার, পার্থক্য 1.019
dropout=0.1, eval() : একই input দুই বার, পার্থক্য 0.000

৫ · Pre-norm বনাম post-norm — ক্রম আর training-এর স্থিরতা

একই তিনটি কাজ, ক্রম আলাদা। Post-norm: LN(x + f(x)), মূল ২০১৭ নকশা। Pre-norm: x + f(LN(x)). দেখতে ছোট বদল, কিন্তু x-এর পথ পুরো আলাদা হয়ে যায়। নিচের ছোট উদাহরণে d = ২, x = [১, ৩], আর sublayer হিসেবে একটি কাল্পনিক f, যা দুই ঘর অদলবদল করে অর্ধেক করে: f([a, b]) = [b/২, a/২].

গণিতে বললে, pre-norm block-এর Jacobian (output-এর সঙ্গে input-এর আংশিক অন্তরকলজের matrix; মেশিন লার্নিং কোর্সের পর্ব ১১) হলো I + F′, আর L টি block-এর গুণফলে একটি বিশুদ্ধ I পদ টিকে থাকে। Post-norm-এ প্রতিটি উৎপাদকের সামনে LN-এর Jacobian-ও বসে। এতে gradient-এর পথ আর পরিষ্কার থাকে না। Xiong ও সহকর্মীরা (২০২০) mean field বিশ্লেষণে দেখান, post-norm-এ শুরুতে output-এর কাছের parameter-এর gradient বড় হতে পারে, যা উঁচু learning rate-এ training অস্থির করে আর warmup তা সামলায়। Pre-norm-এ তাঁরা warmup ছাড়াই তুলনীয় ফল পান।

নিজে মেপে দেখা. কাজটি ছোট: ১৬ অঙ্কের একটি অর্ডার নম্বরের অঙ্কগুলো ছোট থেকে বড় ক্রমে সাজিয়ে লেখা। প্রতিটি output ঘরের জন্য ১০টি অঙ্ক থেকে বাছাই, তাই অনুমান-স্তরের loss ln ১০ ≈ ২.৩০৩. Stack-এ L টি nn.TransformerEncoderLayer, d = ৬৪, h = ৪, r = ৪, dropout ০, B = ৩২, Adam (β = ০.৯, ০.৯৮). Warmup মানে প্রথম ২০০ ধাপে learning rate রৈখিকভাবে বাড়ানো। প্রতি নকশায় ৩টি seed, ৩০০ ধাপ, CPU-তে। Pre-norm stack-এ শেষে একটি LayerNorm আছে।

def batch(g):
    x = torch.randint(0, 10, (32, 16), generator=g)
    return x, x.sort(dim=1).values           # কাজ: অঙ্ক সাজিয়ে লেখা

class Net(nn.Module):
    def __init__(s, d, h, L, pre):
        super().__init__()
        s.tok, s.pos = nn.Embedding(10, d), nn.Embedding(16, d)
        s.blocks = nn.ModuleList(nn.TransformerEncoderLayer(d, h, 4 * d, 0.0,
            "gelu", batch_first=True, norm_first=pre) for _ in range(L))
        s.norm = nn.LayerNorm(d) if pre else nn.Identity()
        s.head = nn.Linear(d, 10)
    def forward(s, x):
        h = s.tok(x) + s.pos(torch.arange(16))
        for b in s.blocks: h = b(h)
        return s.head(s.norm(h))
# লুপে: for pg in opt.param_groups: pg["lr"] = lr * min(1, (t + 1) / warm)

~/bazarghar-dl$ uv run python prepost.py 24 0.003 3 300
--- pre   warm=  0 | L=24 lr=0.003 | 3 seed
step    0 loss  2.446 [ 2.419, 2.466] acc 0.104 [0.083,0.137]
step   25 loss  1.019 [ 0.978, 1.060] acc 0.577 [0.566,0.592]
step   50 loss  0.666 [ 0.624, 0.737] acc 0.735 [0.702,0.757]
step  100 loss  0.395 [ 0.320, 0.448] acc 0.837 [0.811,0.875]
step  200 loss  0.119 [ 0.095, 0.163] acc 0.960 [0.940,0.974]
step  300 loss  0.132 [ 0.073, 0.216] acc 0.954 [0.922,0.975]
  max grad norm, first 200 steps: 4.8 5.4 6.9
…
--- post  warm=  0 | L=24 lr=0.003 | 3 seed
step    0 loss  2.479 [ 2.393, 2.540] acc 0.103 [0.089,0.130]
step   25 loss  2.317 [ 2.303, 2.340] acc 0.104 [0.099,0.112]
step   50 loss  2.310 [ 2.306, 2.317] acc 0.100 [0.099,0.103]
step  100 loss  2.310 [ 2.305, 2.313] acc 0.100 [0.095,0.107]
step  200 loss  2.310 [ 2.308, 2.311] acc 0.101 [0.095,0.107]
step  300 loss  2.308 [ 2.306, 2.311] acc 0.094 [0.084,0.103]
  max grad norm, first 200 steps: 4.9 8.0 3.3
…
…

L = ২৪, learning rate ৩×১০^−৩ (৩ seed-এর গড়)

L = ৬, learning rate ৩×১০^−৩ (৩ seed-এর গড়)

L = ২৪, learning rate ৩×১০^−৪ (৩ seed-এর গড়)

তিনটি টেবিল একসঙ্গে পড়লে ফলটি “pre-norm ভালো, post-norm খারাপ”-এর চেয়ে সূক্ষ্ম। প্রথমত, warmup post-norm-কে বাঁচাল: একই গভীরতায় আটকে না গিয়ে ০.৪২০-তে এল, যদিও pre-norm-এর পেছনে। দ্বিতীয়ত, L = ৬-এ চারটি নকশাই ভালো শিখল; সমস্যা গভীরতার সঙ্গে আসে। তৃতীয়ত, learning rate দশ ভাগের এক ভাগ করলে L = ২৪-এও post-norm warmup ছাড়া শিখল, আর এই একটি run-এ ধাপ ৩০০-তে তার গড় loss সবচেয়ে কম (০.০৫৯), যদিও পরিসর pre-norm-এর সঙ্গে মিশে যায়। অর্থাৎ post-norm “ভুল” নয়; তার নিরাপদ learning rate-এর সীমা এখানে নিচে। Pre-norm এই সীমাটি বাড়িয়ে দেয়। আর warmup pre-norm-কে এখানে সাহায্য করেনি (ধাপ ৩০০-তে L = ২৪, ৩×১০^−৩-এ ০.১৩২ বনাম ০.২১৫), কারণ ৩০০ ধাপের মধ্যে ২০০ ধাপ কম গতিতে গেছে।

৬ · Block কেন ভালো শেখে — gradient সংখ্যায় দেখা

এখন আগের NumPy ফাংশনগুলো দিয়েই ৩২ block-এর একটি stack বানানো হবে, দুই ধরনের weight std নিয়ে: ০.২৫ (d = ১৬-এ ১/√d) আর ০.১. চারটি রূপ পাশাপাশি: দুটিই আছে, শুধু LN, শুধু residual, কোনোটিই নেই। শেষের output-কে একটি স্থির এলোমেলো vector-এ প্রজেক্ট করে একটি সংখ্যা (loss-এর বিকল্প) নেওয়া হবে। k-তম block-এর পরের activation-এর ওপর সেই সংখ্যার gradient norm মাপা হবে finite difference দিয়ে: প্রতিটি উপাদানে ±ε বসিয়ে বাকি block-গুলো আবার চালানো। PyTorch পাতায় চলে না, তাই এই পথ। সবগুলো উপাদান একটি বড় batch-এ রাখায় এটি দ্রুত চলে। ডেস্কটপে একই রকম block-এ (T = ৬, L = ১৬) এই পদ্ধতির ফল PyTorch autograd-এর সঙ্গে পাঁচ significant অঙ্ক পর্যন্ত মিলেছে (যেমন ১১.৫৭৯ দুই দিকেই)।

import numpy as np
import matplotlib.pyplot as plt
TRACE, L = False, 32
ks = list(range(0, L + 1, 8))

def gnorm(W, c, k, a, res, norm):           # k-তম activation-এর gradient norm
    n, e = a.size, 1e-5 * max(1.0, np.abs(a).max())
    p = np.tile(a.reshape(-1), (2 * n, 1)); i = np.arange(n)
    p[i, i] += e; p[n + i, i] -= e
    y = p.reshape(2 * n, 4, d)
    for w in W[k:]: y = block(y, w, res, norm)
    v = ((ln(y) if norm else y) * c).sum((1, 2))
    return np.linalg.norm((v[:n] - v[n:]) / (2 * e))

fig, axes = plt.subplots(1, 2, figsize=(8, 3.4), sharey=True)
for ax, s in zip(axes, (0.25, 0.1)):        # 0.25 = 1/sqrt(16)
    rng = np.random.default_rng(42)
    W = [init(rng, s) for _ in range(L)]
    x0, c = rng.normal(size=(1, 4, d)), rng.normal(size=(4, d))
    print(f"\nweight std = {s}; k = {ks}")
    for name, res, norm in [("residual + LN", 1, 1), ("LN only", 0, 1),
                            ("residual only", 1, 0), ("neither", 0, 0)]:
        a, g = x0, []
        for k in range(L + 1):
            if k in ks: g.append(gnorm(W, c, k, a, res, norm))
            if k < L: a = block(a, W[k], res, norm)
        print(f"{name:14}" + "".join(f"{v:9.1e}" for v in g))
        ax.semilogy(ks, g, marker="o", label=name)
    ax.set_title(f"weight std = {s}")
    ax.set_xlabel("activation after block k (0 = input)")
axes[0].set_ylabel("gradient norm (log)"); axes[0].legend(fontsize=8)
plt.show()

weight std = 0.25; k = [0, 8, 16, 24, 32]
residual + LN   1.3e+01  1.7e+00  1.1e+00  8.9e-01  7.8e-01
LN only         1.5e-01  3.0e-01  5.2e+00  1.3e+01  5.1e+00
residual only   2.8e+14  5.1e+09  5.0e+06  9.5e+03  7.8e+00
neither         2.1e+04  7.2e+02  1.9e+02  6.5e+01  7.8e+00

weight std = 0.1; k = [0, 8, 16, 24, 32]
residual + LN   8.7e+00  6.8e+00  5.5e+00  5.0e+00  4.7e+00
LN only         5.5e+00  2.7e+00  4.5e+00  4.6e+01  4.2e+01
residual only   1.7e+01  1.4e+01  1.2e+01  9.4e+00  7.8e+00
neither         2.9e-51  5.8e-39  1.1e-25  4.7e-13  7.8e+00

এই পরিমাপ থেকে একটি প্রচলিত init কৌশলের যুক্তিও বোঝা যায়। Residual শাখা শুরুতে যত ছোট, block তত identity-র কাছে, আর সিগন্যাল তত অক্ষত। GPT-2 প্রতিবেদনে residual স্তরের weight-কে স্তরের সংখ্যার ১/√N দিয়ে ছোট করার কথা আছে; ReZero আর LayerScale একই লক্ষ্যে শেখা যায় এমন একটি ছোট গুণক বসায়। সব কৌশলের মূল কথা এক: network শুরু করুক প্রায় “কিছু না করে”, তারপর দরকারি বদলটুকু শিখুক।

৭ · Parameter, MAC আর মেমরি — হাতে হিসাব

পরীক্ষা আর ইন্টারভিউতে প্রায়ই প্রশ্ন আসে: এই model-এ parameter কত? উত্তর মুখস্থ না করে একবার গুনে নিন। নিচে d = ৬৪, h = ৪, r = ৪, আর সব linear স্তরে bias আছে (PyTorch-এর default).

সাধারণ সূত্র, r = ৪ আর bias ধরে: attention 4(d²+d), FFN 2rd² + rd + d, LayerNorm 4d. বড় d-তে মোট প্রায় 12d². T, B বা h কোথাও নেই। এই সূত্রগুলো নিচের ব্লকে একটি টেবিল বানায়। সঙ্গে MAC আর মেমরিও আছে।

MAC প্রতি token, এক block-এ: Q, K, V 3d²; স্কোর h·T·d_k = T·d; weighted sum আরও T·d; output projection d²; FFN 2rd². যোগফল (4 + 2r)d² + 2Td, এখানে softmax আর LayerNorm-এর ছোট খরচ ধরা হয়নি। পুরো sequence-এর জন্য T দিয়ে গুণ করুন। দুটি অংশ সমান হয় যখন T = (4 + 2r)d/2, মানে r = ৪-এ T = ৬d. এর নিচের T-তে খরচের বেশি অংশ T² নয়, projection আর FFN-এর।

def parts(d, r):
    att = 3 * (d*d + d) + (d*d + d)
    ffn = (d * r*d + r*d) + (r*d * d + d)
    return att, ffn, 2 * (d + d)

print("  d  h  r  L      att      ffn    LN    block      all L")
for d, h, r, L in [(64,4,4,2), (128,4,4,4), (256,8,4,6), (512,8,4,6), (768,12,4,12)]:
    a, f, n = parts(d, r); b = a + f + n
    print(f"{d:3d} {h:2d} {r:2d} {L:2d} {a:8,d} {f:8,d} {n:5,d} {b:8,d} {L*b:10,d}")

print("\nএকই d = 512, r = 4, আলাদা h:",
      [(h, sum(parts(512, 4))) for h in (1, 8, 16)])

mac = lambda d, r, T: ((4 + 2*r) * d*d, 2 * T * d)
print("\nMAC প্রতি token, d = 256, r = 4")
for T in (64, 512, 1536, 4096):
    p, q = mac(256, 4, T)
    print(f"T = {T:4d}: rest {p:9,d}  attention {q:9,d}  share {q/(p+q):.0%}")

print("\nDepth বনাম width, T = 128, r = 4 (প্রায় একই budget)")
for d, L in [(128, 16), (256, 4)]:
    P = L * sum(parts(d, 4)); M = L * sum(mac(d, 4, 128))
    print(f"d = {d}, L = {L:2d}: parameter {P:9,d}  MAC/token {M:9,d}")

B, h, T, d, r = 8, 8, 512, 256, 4
mb = lambda n: f"{n * 4 / 1e6:6.1f} MB"           # float32 = 4 byte
print(f"\nমেমরি, B={B} h={h} T={T} d={d}, এক block, float32")
print("parameter          ", mb(sum(parts(d, r))))
print("Q, K, V            ", mb(3 * B * T * d))
print("attention weights  ", mb(B * h * T * T))
print("FFN hidden         ", mb(B * T * r * d))

  d  h  r  L      att      ffn    LN    block      all L
 64  4  4  2   16,640   33,088   256   49,984     99,968
128  4  4  4   66,048  131,712   512  198,272    793,088
256  8  4  6  263,168  525,568 1,024  789,760  4,738,560
512  8  4  6 1,050,624 2,099,712 2,048 3,152,384 18,914,304
768 12  4 12 2,362,368 4,722,432 3,072 7,087,872 85,054,464

একই d = 512, r = 4, আলাদা h: [(1, 3152384), (8, 3152384), (16, 3152384)]

MAC প্রতি token, d = 256, r = 4
T =   64: rest   786,432  attention    32,768  share 4%
T =  512: rest   786,432  attention   262,144  share 25%
T = 1536: rest   786,432  attention   786,432  share 50%
T = 4096: rest   786,432  attention 2,097,152  share 73%

Depth বনাম width, T = 128, r = 4 (প্রায় একই budget)
d = 128, L = 16: parameter 3,172,352  MAC/token 3,670,016
d = 256, L =  4: parameter 3,159,040  MAC/token 3,407,872

মেমরি, B=8 h=8 T=512 d=256, এক block, float32
parameter              3.2 MB
Q, K, V               12.6 MB
attention weights     67.1 MB
FFN hidden            16.8 MB

পুরো model-এর একটি হিসাব. ২০১৭ সালের প্রবন্ধের মাপে ধরুন: d = ৫১২, ৬টি encoder আর ৬টি decoder block, প্রবন্ধের ইংরেজি–জার্মান কাজের প্রায় ৩৭,০০০ token-এর ভাগ করা vocabulary, আর embedding ও head একই matrix. Encoder ৬ × ৩১,৫২,৩৮৪ = ১,৮৯,১৪,৩০৪. Decoder ৬ × ৪২,০৪,০৩২ = ২,৫২,২৪,১৯২. Embedding ৩৭,০০০ × ৫১২ = ১,৮৯,৪৪,০০০. মোট প্রায় ৬.৩ কোটি। এটি আমাদের সূত্রের হিসাব, প্রবন্ধের জানানো সংখ্যা নয়; তবে এই পদ্ধতিতেই যেকোনো model-এর আকার অনুমান করা যায়।

Depth বনাম width. L বাড়ালে তথ্য আরও বার token-দের মধ্যে ঘোরে, এবং একটি হিসাবের ওপর আরেকটি হিসাব গড়ে ওঠে। d বাড়ালে প্রতি token-এ বেশি তথ্য ধরে, কিন্তু parameter বর্গে বাড়ে। কোনো একক নিয়ম নেই। Stochastic depth (পর্ব ১৬) বা residual-এর ছোট শুরু খুব গভীর stack-এ সাহায্য করে; আবার gradient checkpointing মেমরি বাঁচায়, সময় নেয়। নিজের কাজে কয়েকটি ছোট কনফিগারেশন validation-এ তুলনা করাই সৎ পথ।

৮ · Decoder block আর cross-attention

Decoder block-এ encoder-এর দুটির পাশে একটি তৃতীয় sublayer যুক্ত হয়। প্রথম sublayer হলো causal mask-সহ self-attention, যাতে কোনো ঘর পরের token না দেখে। তৃতীয়টি FFN. মাঝেরটি cross-attention, যা পর্ব ৩০-এর fixed context-এর সমস্যার সমাধান: এখন decoder-এর প্রতিটি ঘর encoder-এর পুরো output, মানে memory, দেখতে পারে। নিচে d = ২, এক head, projection বাদ (Q, K, V সরাসরি বসানো), T_src = ৩, T_tgt = ২.

বাস্তবে দুটি কথা মনে রাখা দরকার। প্রথমত, encoder একবার চলে আর memory বানিয়ে রাখে; decoder প্রতিটি নতুন token লেখার সময় সেই একই memory পড়ে। দ্বিতীয়ত, input-এ padding থাকলে memory-র সেই ঘরগুলোও বাদ দিতে হয়। তাই cross-attention-এ causal mask নয়, লাগে padding mask (PyTorch-এ memory_key_padding_mask). পর্ব ৩০-এর padding-এর সমস্যাটি এখানে অন্য রূপে ফিরে আসে: এই mask ভুলে গেলে কোনো error হয় না, শুধু decoder pad-এর দিকেও তাকায়।

সত্যিকারের run-এ মাপগুলো এবার বড় করে দেখুন। নিচে আমাদের DecBlock pre-norm, তিন sublayer, আর উপরের MHA-ই ব্যবহার করে। Q একটি input থেকে, K আর V অন্যটি থেকে নেওয়ার জন্য শুধু m দেওয়া হয়।

def forward(self, y, mem, causal):
    y = y + self.sa(self.n1(y), mask=causal)   # masked self-attention
    y = y + self.ca(self.n2(y), m=mem)         # Q: y, K ও V: mem
    return y + self.ffn(self.n3(y))

~/bazarghar-dl$ uv run python encdec.py
src (2, 9, 64) -> memory (2, 9, 64)
tgt (2, 5, 64)
  masked self-attn-এর পরে (2, 5, 64)
  cross-attn-এর পরে       (2, 5, 64)
  FFN-এর পরে              (2, 5, 64)
parameter: encoder block 49984 | decoder block 66752
cross weights (B, h, T_tgt, T_src): (2, 4, 5, 9)
এক সারির যোগফল: 1.0
T_src =  3 -> decoder output (2, 5, 64) | parameter 66752
T_src = 40 -> decoder output (2, 5, 64) | parameter 66752
d=512, h=8, ffn=2048: 3152384 4204032

৯ · Weight tying — একই matrix দুই জায়গায়

Language model-এর শুরুতে embedding matrix (V, d), শেষে LM head (d → V). দুটি একই আকারের। Embedding বলে “এই token-এর vector কী”, head বলে “এই vector কোন token-এর কাছাকাছি”। তাই একই matrix দুই কাজে লাগানো স্বাভাবিক চিন্তা। একে বলে weight tying. PyTorch-এ এক লাইন: head.weight = tok.weight; এখানে V মানে vocabulary-র আকার, আগের attention-এর V নয়।

~/bazarghar-dl$ uv run python tying.py
V    d  L     untied       tied     saved     %
   100   64  2    129,280    122,880     6,400   5.0
  5000  128  4  2,106,112  1,466,112   640,000  30.4
 32000  256  6 21,188,608 12,996,608 8,192,000  38.7
 50000  768 12 162,052,608 123,652,608 38,400,000  23.7
same object: True
state_dict keys has both: True True
named_parameters has head.weight: False
wrong count over state_dict: 2,106,112 | right: 1,466,112
one grad tensor shared: torch.Size([5000, 128]) True

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

AI assistant একটি Transformer block দ্রুত লিখে দেয়, আর কোডটি প্রায় সবসময় চলে। ভুলগুলো error দেয় না; তাই সেগুলো ধরতে হয় মেপে।

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

• ছোট batch-এ প্রতিটি sublayer-এর shape ছাপিয়ে মিলিয়েছি; T_src আর T_tgt আলাদা রেখেছি
• model.eval() আর torch.no_grad() দিয়ে dropout বন্ধ করে পরীক্ষিত layer-এর সঙ্গে সংখ্যা মিলিয়েছি
• Parameter সূত্র আর model.parameters() এক; পরের token বদলে causal mask পরীক্ষা করেছি
• শুরুর loss ln V-এর কাছে; একটি ছোট batch overfit হয়
• সহজ baseline-এর সঙ্গে, কয়েকটি seed-এ তুলনা করেছি; test ডেটা দিয়ে কিছু বাছাই করিনি
• Argument আর default নিজের সংস্করণের docs-এ দেখেছি, প্রবন্ধের দাবি মূল সূত্রে মিলিয়েছি
• ব্যক্তিগত ডেটা বা key AI-কে দিইনি; “model বোঝে” ধরনের দাবি যাচাই ছাড়া মানিনি