RNN

25WORD_CATEGORIES = {▶
26    "NOUN_HUMAN":  ["boy", "girl", "man", "woman", "king", "queen"],▶
27    "NOUN_ANIMAL": ["cat", "dog", "lion", "bird", "fish", "mouse"],
28    "NOUN_FOOD":   ["bread", "cake", "apple", "milk", "cookie"],
29    "NOUN_OBJECT": ["book", "rock", "glass", "plate", "ball"],
30    "VERB_INTRAN": ["sleep", "walk", "think", "swim"],
31    "VERB_TRAN":   ["chase", "eat", "see", "break", "read", "move"],
32    "VERB_PERCEPT":["smell", "hear", "watch"],▶
33}
34

50TEMPLATES = [▶
51    # 1.  NOUN_HUMAN VERB_INTRAN
52    [("DET+CAT", "NOUN_HUMAN"), ("CAT", "VERB_INTRAN")],
53    # 2.  NOUN_HUMAN VERB_TRAN NOUN_OBJECT
54    [("DET+CAT", "NOUN_HUMAN"), ("CAT", "VERB_TRAN"), ("DET+CAT", "NOUN_OBJECT")],
55    # ...16 templates total...
56]
57
58def generate_sentence(template: list) -> str:
59    """Fill one template with random word choices."""
60    tokens = []
61    for slot in template:
62        if slot[0] == "DET+CAT":
63            word = random.choice(WORD_CATEGORIES[slot[1]])▶
64            tokens.append(pick_det())
65            tokens.append(word)
66        elif slot[0] == "CAT":
67            word = random.choice(WORD_CATEGORIES[slot[1]])▶
68            tokens.append(word)
69    return " ".join(tokens)▶
70

119def generate_corpus(num_sentences: int, seed: int = 42):▶
120    random.seed(seed)
121    np.random.seed(seed)
122    sentences = []
123    for _ in range(num_sentences):
124        tidx = int(np.random.choice(len(TEMPLATES), p=normalized_weights))▶
125        sentence = generate_sentence(TEMPLATES[tidx])
126        sentences.append(sentence)
127    return sentences, template_counts
128
129# Usage:
130# python generate_corpus.py --num-sentences 600 --seed 42
131corpus_text = " ".join(sentences)▶
132

157def load_corpus(path: str) -> str:
158    """Read corpus file, lowercase, strip extra whitespace."""
159    text = Path(path).read_text(encoding="utf-8")
160    text = text.lower().strip()
161    # Collapse multiple spaces to single space
162    import re
163    text = re.sub(r" +", " ", text)
164    return text
165
166
167

167def build_vocab(text: str) -> Tuple[List[str], Dict[str, int], Dict[int, str]]:
168    """Extract unique characters, sort, create char↔index mappings."""
169    chars = sorted(set(text))▶
170    char2idx = {c: i for i, c in enumerate(chars)}▶
171    idx2char = {i: c for i, c in enumerate(chars)}▶
172    return chars, char2idx, idx2char▶
173
174
175

175def encode(text: str, char2idx: Dict[str, int]) -> List[int]:
176    """Map each character to its index."""
177    return [char2idx[c] for c in text]▶
178
179
180

180def split_data(data: List[int], val_fraction: float = 0.1) -> Tuple[List[int], List[int]]:▶
181    """
182    Split encoded corpus into train/val.
183
184    Uses the last val_fraction of the corpus for validation, matching the▶
185    temporal ordering of the data (avoids data leakage from future context).
186    """
187    split_idx = int(len(data) * (1.0 - val_fraction))▶
188    return data[:split_idx], data[split_idx:]
189
190
191# ---------------------------------------------------------------------------
192# Activation functions
193# ---------------------------------------------------------------------------
194
195

70class CharRNN:
71    """Vanilla RNN for character-level language modeling.
72
73    All parameters are raw nn.Parameter tensors. No nn.Module subclass,
74    keeping it completely manual so every operation is visible in the code
75    walkthrough.
76
77    Architecture:
78        x_t (V,1) --[W_xh]--> a_t = W_xh @ x_t + W_hh @ h_{t-1} + b_h▶
79        h_t = activation(a_t)  ← H-dimensional hidden state
80        o_t = W_hy @ h_t + b_y  ← output logits (V,1)▶
81        p_t = softmax(o_t)       ← probability over next character
82    """
83
84    def __init__(self, vocab_size: int, hidden_size: int):
85        self.vocab_size = vocab_size    # V ≈ 27 (characters a-z + space)
86        self.hidden_size = hidden_size  # H = 64 (default)
87
88        # Input-to-hidden: (H, V)
89        self.W_xh = nn.Parameter(torch.empty(hidden_size, vocab_size))▶
90        # Hidden-to-hidden recurrent: (H, H)
91        self.W_hh = nn.Parameter(torch.empty(hidden_size, hidden_size))▶
92        # Hidden-to-output: (V, H)
93        self.W_hy = nn.Parameter(torch.empty(vocab_size, hidden_size))▶
94        # Biases: initialized to zeros in constructor
95        self.b_h = nn.Parameter(torch.zeros(hidden_size, 1))▶
96        self.b_y = nn.Parameter(torch.zeros(vocab_size, 1))
97
98    def parameters(self) -> List[nn.Parameter]:
99        return [self.W_xh, self.W_hh, self.W_hy, self.b_h, self.b_y]▶
100
101    def count_parameters(self) -> int:
102        return sum(p.numel() for p in self.parameters())
103
104
105# ---------------------------------------------------------------------------
106# Initialization
107# ---------------------------------------------------------------------------
108
109

109def init_parameters(model: CharRNN, method: str = "orthogonal", seed: int = 42) -> None:
110    """
111    Initialize model parameters.
112
113    Args:
114        method: W_hh initialization strategy.
115            "orthogonal" (default): spectral radius = 1.0, optimal for gradient flow
116            "xavier"              : Xavier uniform for W_hh (not RNN-specific)
117            "random"              : N(0, 0.01) random, worst case for gradient flow
118        seed: torch manual seed
119
120    Why orthogonal for W_hh:
121        An orthogonal matrix has |eigenvalue| = 1 for all eigenvalues, so the
122        Jacobian product in BPTT neither grows nor shrinks. This places the
123        network at the optimal spectral radius for gradient flow, the
124        RNN-specific initialization insight that Xavier alone does not provide.
125    """
126    torch.manual_seed(seed)
127    V, H = model.vocab_size, model.hidden_size
128
129    # Xavier uniform for input projection: keeps variance stable across layers
130    bound_xh = math.sqrt(6.0 / (V + H))
131    nn.init.uniform_(model.W_xh, -bound_xh, bound_xh)
132
133    # W_hh initialization (method-dependent for ablation experiments)
134    if method == "orthogonal":
135        nn.init.orthogonal_(model.W_hh)▶
136    elif method == "xavier":
137        bound_hh = math.sqrt(6.0 / (H + H))
138        nn.init.uniform_(model.W_hh, -bound_hh, bound_hh)
139    elif method == "random":
140        nn.init.normal_(model.W_hh, mean=0.0, std=0.01)
141    else:
142        raise ValueError(f"Unknown init method: {method}. Choose: orthogonal/xavier/random")
143
144    # Xavier uniform for output projection
145    bound_hy = math.sqrt(6.0 / (H + V))
146    nn.init.uniform_(model.W_hy, -bound_hy, bound_hy)
147
148    # Biases stay zero (already initialized in constructor, reset here for safety)
149    nn.init.zeros_(model.b_h)
150    nn.init.zeros_(model.b_y)
151
152
153# ---------------------------------------------------------------------------
154# Corpus preprocessing
155# ---------------------------------------------------------------------------
156
157

224def forward(
225    model: CharRNN,
226    inputs: List[int],
227    h_prev: np.ndarray,
228    activation: str = "tanh",
229    use_embedding: bool = False,
230    embedding: Optional[np.ndarray] = None,
231) -> dict:
232    """
233    Forward pass through seq_len time steps.
234
235    Args:
236        model: CharRNN with parameters
237        inputs: list of seq_len+1 character indices (input[0..T-1] → target[1..T])
238        h_prev: hidden state from previous chunk, shape (H, 1)
239        activation: activation function name ('tanh'/'relu'/'sigmoid')
240        use_embedding: if True, use learned embedding instead of one-hot
241        embedding: embedding matrix (H_emb, V) if use_embedding=True
242
243    Returns dict with:
244        xs: one-hot or embedded input vectors {t: (V_or_E, 1)}
245        xs_indices: original character indices {t: int}
246        hs: hidden states {t: (H,1)}, hs[-1] = h_prev
247        os: output logits {t: (V,1)}
248        ps: softmax probabilities {t: (V,1)}
249        targets: list of target character indices
250        loss: scalar total (unnormalized) cross-entropy loss
251
252    Note on loss normalization:
253        loss here is the SUM over all time steps (unnormalized). capture_step
254        divides by seq_len to get per-character bits. Gradient magnitudes in
255        backward() reflect this unnormalized loss and scale linearly with
256        seq_len, important to understand when comparing experiments with
257        different chunk lengths (Plan 7, Experiment 3).
258    """
259    act_fn, _ = _activation_fn(activation)
260
261    W_xh = model.W_xh.data.numpy()
262    W_hh = model.W_hh.data.numpy()
263    W_hy = model.W_hy.data.numpy()
264    b_h  = model.b_h.data.numpy()
265    b_y  = model.b_y.data.numpy()
266
267    xs, xs_indices, hs, os, ps = {}, {}, {}, {}, {}
268    hs[-1] = h_prev.copy()
269    loss = 0.0
270
271    seq_len = len(inputs) - 1  # predict inputs[1..T] from inputs[0..T-1]
272
273    for t in range(seq_len):
274        xs_indices[t] = inputs[t]
275
276        if use_embedding and embedding is not None:
277            # Learned embedding lookup: (E, 1)
278            xs[t] = embedding[:, inputs[t]].reshape(-1, 1)
279        else:
280            # One-hot encoding: (V, 1)
281            xs[t] = np.zeros((model.vocab_size, 1))
282            xs[t][inputs[t]] = 1.0
283
284        # h_t = activation(W_xh @ x_t + W_hh @ h_{t-1} + b_h)
285        pre_act = W_xh @ xs[t] + W_hh @ hs[t - 1] + b_h
286        hs[t] = act_fn(pre_act)
287
288        # o_t = W_hy @ h_t + b_y
289        os[t] = W_hy @ hs[t] + b_y
290
291        # softmax with numerical stability: subtract max before exp
292        exp_o = np.exp(os[t] - np.max(os[t]))
293        ps[t] = exp_o / np.sum(exp_o)
294
295        # cross-entropy: -log(p[target])
296        target_t = inputs[t + 1]
297        loss += -np.log(ps[t][target_t, 0] + 1e-12)▶
298
299    targets = inputs[1:]
300    return {
301        "xs": xs,
302        "xs_indices": xs_indices,
303        "hs": hs,
304        "os": os,
305        "ps": ps,
306        "targets": targets,
307        "loss": loss,
308    }
309
310
311# ---------------------------------------------------------------------------
312# Backward pass: Manual BPTT
313# ---------------------------------------------------------------------------
314
315

224def forward(
225    model: CharRNN,
226    inputs: List[int],
227    h_prev: np.ndarray,
228    activation: str = "tanh",
229    use_embedding: bool = False,
230    embedding: Optional[np.ndarray] = None,
231) -> dict:
232    """
233    Forward pass through seq_len time steps.
234
235    Args:
236        model: CharRNN with parameters
237        inputs: list of seq_len+1 character indices (input[0..T-1] → target[1..T])
238        h_prev: hidden state from previous chunk, shape (H, 1)
239        activation: activation function name ('tanh'/'relu'/'sigmoid')
240        use_embedding: if True, use learned embedding instead of one-hot
241        embedding: embedding matrix (H_emb, V) if use_embedding=True
242
243    Returns dict with:
244        xs: one-hot or embedded input vectors {t: (V_or_E, 1)}
245        xs_indices: original character indices {t: int}
246        hs: hidden states {t: (H,1)}, hs[-1] = h_prev
247        os: output logits {t: (V,1)}
248        ps: softmax probabilities {t: (V,1)}
249        targets: list of target character indices
250        loss: scalar total (unnormalized) cross-entropy loss
251
252    Note on loss normalization:
253        loss here is the SUM over all time steps (unnormalized). capture_step
254        divides by seq_len to get per-character bits. Gradient magnitudes in
255        backward() reflect this unnormalized loss and scale linearly with
256        seq_len, important to understand when comparing experiments with
257        different chunk lengths (Plan 7, Experiment 3).
258    """
259    act_fn, _ = _activation_fn(activation)
260
261    W_xh = model.W_xh.data.numpy()
262    W_hh = model.W_hh.data.numpy()
263    W_hy = model.W_hy.data.numpy()
264    b_h  = model.b_h.data.numpy()
265    b_y  = model.b_y.data.numpy()
266
267    xs, xs_indices, hs, os, ps = {}, {}, {}, {}, {}
268    hs[-1] = h_prev.copy()
269    loss = 0.0
270
271    seq_len = len(inputs) - 1  # predict inputs[1..T] from inputs[0..T-1]
272
273    for t in range(seq_len):
274        xs_indices[t] = inputs[t]
275
276        if use_embedding and embedding is not None:
277            # Learned embedding lookup: (E, 1)
278            xs[t] = embedding[:, inputs[t]].reshape(-1, 1)
279        else:
280            # One-hot encoding: (V, 1)
281            xs[t] = np.zeros((model.vocab_size, 1))
282            xs[t][inputs[t]] = 1.0
283
284        # h_t = activation(W_xh @ x_t + W_hh @ h_{t-1} + b_h)
285        pre_act = W_xh @ xs[t] + W_hh @ hs[t - 1] + b_h
286        hs[t] = act_fn(pre_act)
287
288        # o_t = W_hy @ h_t + b_y▶
289        os[t] = W_hy @ hs[t] + b_y
290
291        # softmax with numerical stability: subtract max before exp
292        exp_o = np.exp(os[t] - np.max(os[t]))
293        ps[t] = exp_o / np.sum(exp_o)
294
295        # cross-entropy: -log(p[target])
296        target_t = inputs[t + 1]
297        loss += -np.log(ps[t][target_t, 0] + 1e-12)
298
299    targets = inputs[1:]
300    return {
301        "xs": xs,
302        "xs_indices": xs_indices,
303        "hs": hs,
304        "os": os,
305        "ps": ps,
306        "targets": targets,
307        "loss": loss,
308    }
309
310
311# ---------------------------------------------------------------------------
312# Backward pass: Manual BPTT
313# ---------------------------------------------------------------------------
314
315

224def forward(
225    model: CharRNN,
226    inputs: List[int],
227    h_prev: np.ndarray,
228    activation: str = "tanh",
229    use_embedding: bool = False,
230    embedding: Optional[np.ndarray] = None,
231) -> dict:
232    """
233    Forward pass through seq_len time steps.
234
235    Args:
236        model: CharRNN with parameters
237        inputs: list of seq_len+1 character indices (input[0..T-1] → target[1..T])
238        h_prev: hidden state from previous chunk, shape (H, 1)
239        activation: activation function name ('tanh'/'relu'/'sigmoid')
240        use_embedding: if True, use learned embedding instead of one-hot
241        embedding: embedding matrix (H_emb, V) if use_embedding=True
242
243    Returns dict with:
244        xs: one-hot or embedded input vectors {t: (V_or_E, 1)}
245        xs_indices: original character indices {t: int}
246        hs: hidden states {t: (H,1)}, hs[-1] = h_prev
247        os: output logits {t: (V,1)}
248        ps: softmax probabilities {t: (V,1)}
249        targets: list of target character indices
250        loss: scalar total (unnormalized) cross-entropy loss
251
252    Note on loss normalization:
253        loss here is the SUM over all time steps (unnormalized). capture_step
254        divides by seq_len to get per-character bits. Gradient magnitudes in
255        backward() reflect this unnormalized loss and scale linearly with
256        seq_len, important to understand when comparing experiments with
257        different chunk lengths (Plan 7, Experiment 3).
258    """
259    act_fn, _ = _activation_fn(activation)
260
261    W_xh = model.W_xh.data.numpy()
262    W_hh = model.W_hh.data.numpy()
263    W_hy = model.W_hy.data.numpy()
264    b_h  = model.b_h.data.numpy()
265    b_y  = model.b_y.data.numpy()
266
267    xs, xs_indices, hs, os, ps = {}, {}, {}, {}, {}
268    hs[-1] = h_prev.copy()
269    loss = 0.0
270
271    seq_len = len(inputs) - 1  # predict inputs[1..T] from inputs[0..T-1]
272
273    for t in range(seq_len):
274        xs_indices[t] = inputs[t]
275
276        if use_embedding and embedding is not None:
277            # Learned embedding lookup: (E, 1)
278            xs[t] = embedding[:, inputs[t]].reshape(-1, 1)
279        else:
280            # One-hot encoding: (V, 1)
281            xs[t] = np.zeros((model.vocab_size, 1))
282            xs[t][inputs[t]] = 1.0
283
284        # h_t = activation(W_xh @ x_t + W_hh @ h_{t-1} + b_h)
285        pre_act = W_xh @ xs[t] + W_hh @ hs[t - 1] + b_h
286        hs[t] = act_fn(pre_act)
287
288        # o_t = W_hy @ h_t + b_y
289        os[t] = W_hy @ hs[t] + b_y
290
291        # softmax with numerical stability: subtract max before exp
292        exp_o = np.exp(os[t] - np.max(os[t]))
293        ps[t] = exp_o / np.sum(exp_o)
294
295        # cross-entropy: -log(p[target])
296        target_t = inputs[t + 1]
297        loss += -np.log(ps[t][target_t, 0] + 1e-12)
298
299    targets = inputs[1:]
300    return {
301        "xs": xs,
302        "xs_indices": xs_indices,
303        "hs": hs,
304        "os": os,
305        "ps": ps,
306        "targets": targets,
307        "loss": loss,
308    }
309
310
311# ---------------------------------------------------------------------------
312# Backward pass: Manual BPTT
313# ---------------------------------------------------------------------------
314
315