Language Model Fundamentals

Input:  sequence x₁ … x_T,  initial state h₀ = 0
Params: W_xh (D→H),  W_hh (H→H),  W_hy (H→V),  biases b_h, b_y

for t = 1 … T:
    aₜ  ←  W_xh · xₜ  +  W_hh · hₜ₋₁  +  b_h
    hₜ  ←  tanh(aₜ)                 # new hidden state
    ŷₜ  ←  softmax(W_hy · hₜ + b_y)   # output distribution

return h₁ … h_T, ŷ₁ … ŷ_T

import numpy as np

class VanillaRNN:
    def __init__(self, vocab, hidden=64, seed=0):
        rng = np.random.default_rng(seed)
        s   = 0.01                      # small init to avoid early saturation
        self.Wxh = rng.normal(0, s, (hidden, vocab))
        self.Whh = rng.normal(0, s, (hidden, hidden))
        self.Why = rng.normal(0, s, (vocab, hidden))
        self.bh  = np.zeros((hidden, 1))
        self.by  = np.zeros((vocab,  1))
        self.H   = hidden

    def forward(self, inputs, h_prev):
        """inputs: list of one-hot column vectors. Returns states + caches."""
        h = h_prev
        xs, hs, ys, ps = {}, {-1: h_prev}, {}, {}
        for t, x in enumerate(inputs):
            xs[t] = x
            h     = np.tanh(self.Wxh @ x + self.Whh @ hs[t-1] + self.bh)
            y     = self.Why @ h + self.by
            p     = np.exp(y - y.max()); p /= p.sum()   # stable softmax
            hs[t], ys[t], ps[t] = h, y, p
        return xs, hs, ps

# A character-level RNN trained on text can generate plausible characters,
# but it forgets context beyond ~10-20 steps. The next section shows why.

def backward(self, xs, hs, ps, targets):
    dWxh = np.zeros_like(self.Wxh); dWhh = np.zeros_like(self.Whh)
    dWhy = np.zeros_like(self.Why); dbh = np.zeros_like(self.bh); dby = np.zeros_like(self.by)
    dh_next = np.zeros((self.H, 1))

    for t in reversed(range(len(targets))):  # walk backwards in time
        dy = ps[t].copy(); dy[targets[t]] -= 1   # softmax-CE gradient: p - y
        dWhy += dy @ hs[t].T;  dby += dy

        dh  = self.Why.T @ dy + dh_next         # grad flowing into h_t
        draw = (1 - hs[t]**2) * dh            # tanh'(a) = 1 - tanh²(a)
        dbh += draw
        dWxh += draw @ xs[t].T
        dWhh += draw @ hs[t-1].T
        dh_next = self.Whh.T @ draw            # propagate to previous step

    # Gradient clipping is ESSENTIAL for RNNs (prevents explosion)
    for d in [dWxh, dWhh, dWhy, dbh, dby]:
        np.clip(d, -5, 5, out=d)
    return dWxh, dWhh, dWhy, dbh, dby

import numpy as np

def gradient_norm_over_time(W_hh, T=100, H=64):
    """Track how a gradient shrinks as it flows back through T steps."""
    grad  = np.random.randn(H, 1)       # gradient at the final step
    norms = [np.linalg.norm(grad)]
    for _ in range(T):
        # Each backward step: multiply by W_hhᵀ and a tanh' factor (~0.5 avg)
        grad = (W_hh.T @ grad) * 0.5
        norms.append(np.linalg.norm(grad))
    return np.array(norms)

rng = np.random.default_rng(0)
H   = 64

# Case 1: small weights → vanishing
W_small = rng.normal(0, 0.5/np.sqrt(H), (H, H))
norms_v = gradient_norm_over_time(W_small)
print(f"After 100 steps (small W): {norms_v[-1]/norms_v[0]:.2e} of original")
# After 100 steps (small W): 3.7e-31 of original  -- completely vanished

# Case 2: large weights → exploding
W_large = rng.normal(0, 3.0/np.sqrt(H), (H, H))
norms_e = gradient_norm_over_time(W_large)
print(f"After 100 steps (large W): {norms_e[-1]/norms_e[0]:.2e} of original")
# After 100 steps (large W): 1.4e+18 of original  -- exploded to NaN territory

# Conclusion: a gradient cannot survive 100 steps of multiplication by W_hh.
# Information from t=1 cannot influence the loss at t=100. The RNN is 'amnesiac'.

# Concatenate previous hidden state and current input: [hₜ₋₁ ; xₜ]

fₜ  ←  σ(W_f · [hₜ₋₁, xₜ] + b_f)        # forget gate
iₜ  ←  σ(W_i · [hₜ₋₁, xₜ] + b_i)        # input gate
c̃ₜ  ←  tanh(W_c · [hₜ₋₁, xₜ] + b_c)      # candidate cell content
oₜ  ←  σ(W_o · [hₜ₋₁, xₜ] + b_o)        # output gate

cₜ  ←  fₜ ⊙ cₜ₋₁  +  iₜ ⊙ c̃ₜ            # ← THE GRADIENT HIGHWAY
hₜ  ←  oₜ ⊙ tanh(cₜ)                    # gated output

import numpy as np

def sigmoid(x): return 1 / (1 + np.exp(-np.clip(x, -30, 30)))

class LSTMCell:
    def __init__(self, input_dim, hidden, seed=0):
        rng = np.random.default_rng(seed)
        Z   = input_dim + hidden          # concatenated [h, x] size
        s   = 1.0 / np.sqrt(Z)
        # One weight matrix per gate; rows map [h_prev, x] → gate value
        self.Wf = rng.normal(0, s, (hidden, Z)); self.bf = np.ones((hidden,1))  # bf=1: start by remembering
        self.Wi = rng.normal(0, s, (hidden, Z)); self.bi = np.zeros((hidden,1))
        self.Wc = rng.normal(0, s, (hidden, Z)); self.bc = np.zeros((hidden,1))
        self.Wo = rng.normal(0, s, (hidden, Z)); self.bo = np.zeros((hidden,1))
        self.H  = hidden

    def step(self, x, h_prev, c_prev):
        z  = np.vstack([h_prev, x])       # concatenate
        f  = sigmoid(self.Wf @ z + self.bf)   # forget gate
        i  = sigmoid(self.Wi @ z + self.bi)   # input gate
        c_tilde = np.tanh(self.Wc @ z + self.bc)  # candidate
        o  = sigmoid(self.Wo @ z + self.bo)   # output gate

        c  = f * c_prev + i * c_tilde       # cell update (the highway)
        h  = o * np.tanh(c)                 # gated hidden output
        return h, c, (f, i, c_tilde, o, z)  # cache for backprop

# Note bf initialised to 1.0: the standard 'forget gate bias' trick.
# It makes the LSTM remember by default at the start of training,
# dramatically improving learning of long-range dependencies (Jozefowicz 2015).

rₜ  ←  σ(W_r · [hₜ₋₁, xₜ] + b_r)          # reset gate
zₜ  ←  σ(W_z · [hₜ₋₁, xₜ] + b_z)          # update gate

h̃ₜ  ←  tanh(W_h · [rₜ ⊙ hₜ₋₁, xₜ] + b_h)   # candidate state
hₜ  ←  (1 − zₜ) ⊙ hₜ₋₁  +  zₜ ⊙ h̃ₜ        # interpolate old & new

import numpy as np

class GRUCell:
    def __init__(self, input_dim, hidden, seed=0):
        rng = np.random.default_rng(seed); Z = input_dim + hidden; s = 1/np.sqrt(Z)
        self.Wr = rng.normal(0, s, (hidden, Z)); self.br = np.zeros((hidden,1))
        self.Wz = rng.normal(0, s, (hidden, Z)); self.bz = np.zeros((hidden,1))
        self.Wh = rng.normal(0, s, (hidden, Z)); self.bh = np.zeros((hidden,1))

    def step(self, x, h_prev):
        z_cat = np.vstack([h_prev, x])
        r = sigmoid(self.Wr @ z_cat + self.br)    # reset gate
        z = sigmoid(self.Wz @ z_cat + self.bz)    # update gate
        # Candidate uses reset-gated previous state
        h_tilde = np.tanh(self.Wh @ np.vstack([r * h_prev, x]) + self.bh)
        h = (1 - z) * h_prev + z * h_tilde       # interpolate
        return h

# 3 matrices vs LSTM's 4. On a 256-hidden, 300-input cell:
#   GRU:  3 × 256 × 556 = 427k params
#   LSTM: 4 × 256 × 556 = 569k params   (33% more)

import torch; import torch.nn as nn

# A 2-layer bidirectional LSTM for sequence tagging
bilstm = nn.LSTM(
    input_size=300,       # word embedding dimension
    hidden_size=256,      # hidden units per direction
    num_layers=2,         # stacked depth
    bidirectional=True,    # forward + backward
    batch_first=True,
)

x       = torch.randn(32, 50, 300)  # (batch, seq_len, emb_dim)
out, (h, c) = bilstm(x)
print(out.shape)   # (32, 50, 512) -- 512 = 256 × 2 directions

# A tagging head maps each 512-d state to tag logits
tagger = nn.Linear(512, 9)  # e.g. 9 NER tags (BIO scheme)
logits = tagger(out)           # (32, 50, 9)

ENCODER  (reads the source sequence)
  h^enc ← 0
  for t = 1 … T_src:
    h^enc ← LSTM_enc(x_t, h^enc)
  context ← h^enc_{T_src}       # final state = fixed-size summary

DECODER  (generates the target sequence)
  h^dec ← context                # initialise from encoder
  y_0 ← <START>
  for t = 1 … T_tgt:
    h^dec ← LSTM_dec(embed(y_{t-1}), h^dec)
    y_t   ← argmax softmax(W · h^dec)  # emit next token
    if y_t == <END>: break

import torch; import torch.nn as nn

class Encoder(nn.Module):
    def __init__(self, vocab, emb=256, hidden=512):
        super().__init__()
        self.embed = nn.Embedding(vocab, emb)
        self.lstm  = nn.LSTM(emb, hidden, batch_first=True)
    def forward(self, src):
        out, (h, c) = self.lstm(self.embed(src))
        return out, (h, c)     # out used later for attention

class Decoder(nn.Module):
    def __init__(self, vocab, emb=256, hidden=512):
        super().__init__()
        self.embed = nn.Embedding(vocab, emb)
        self.lstm  = nn.LSTM(emb, hidden, batch_first=True)
        self.proj  = nn.Linear(hidden, vocab)
    def forward(self, tok, state):
        out, state = self.lstm(self.embed(tok), state)
        return self.proj(out), state

# Training uses 'teacher forcing': feed the TRUE previous target token,
# not the model's own (possibly wrong) prediction. This stabilises early training
# but creates exposure bias -- the model never practises recovering from its
# own mistakes (a problem scheduled sampling and RL fine-tuning later address).

import torch; import torch.nn as nn; import torch.nn.functional as F

class BahdanauAttention(nn.Module):
    def __init__(self, hidden):
        super().__init__()
        self.W1 = nn.Linear(hidden, hidden)    # projects decoder state
        self.W2 = nn.Linear(hidden, hidden)    # projects encoder states
        self.v  = nn.Linear(hidden, 1, bias=False)

    def forward(self, dec_state, enc_outputs):
        # dec_state: (B, H)   enc_outputs: (B, T_src, H)
        dec = self.W1(dec_state).unsqueeze(1)    # (B, 1, H)
        enc = self.W2(enc_outputs)              # (B, T_src, H)
        scores = self.v(torch.tanh(dec + enc)).squeeze(-1)  # (B, T_src)
        weights = F.softmax(scores, dim=-1)       # (B, T_src) attention
        context = torch.bmm(weights.unsqueeze(1), enc_outputs)  # (B,1,H)
        return context.squeeze(1), weights     # context + weights for viz

# The attention weights are interpretable: plotting the (T_tgt × T_src)
# weight matrix for a translation reveals soft word alignments -- e.g.
# 'la maison bleue' → 'the blue house' shows the diagonal-ish alignment
# with the adjective-noun swap clearly visible.

import torch; import torch.nn as nn

model   = nn.LSTM(256, 512, batch_first=True)
opt     = torch.optim.Adam(model.parameters(), lr=1e-3)
bptt    = 35                            # truncation length
hidden  = None

for chunk in data_chunks:            # each chunk is (B, 35, 256)
    # Detach hidden state: carry VALUES forward but cut the gradient graph
    if hidden is not None:
        hidden = tuple(h.detach() for h in hidden)

    out, hidden = model(chunk, hidden)
    loss = compute_loss(out, targets)

    opt.zero_grad()
    loss.backward()                    # gradient only flows within this chunk
    nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)  # ESSENTIAL
    opt.step()

# Why detach? Without it, the graph grows unboundedly across chunks and
# you'd backprop through the entire history -- defeating the truncation.