Word-Level LSTM Language Model from Scratch

This repository implements a Word-Level Autoregressive Language Model using Long Short-Term Memory (LSTM) cells, built entirely from scratch in Python using only NumPy. The project covers the complete pipeline: text preprocessing, learning vector representations (Word2Vec Skip-Gram), recurrent sequence modeling, BPTT optimization, and autoregressive text generation.

---

📖 Project Overview

The goal of this project is to explore and implement the core mathematical and architectural mechanics of recurrent sequence models without relying on deep learning frameworks like PyTorch or TensorFlow. By writing every layer in raw NumPy, it exposes:

1. Text Preprocessing: Cleaning corpus text, tokenization, and vocabulary creation.
2. Word Embeddings (Word2Vec Skip-Gram): Training custom dense word vectors on a corpus (Alice in Wonderland).
3. LSTM Recurrent State Transitions: Modeling long-term sequence dependencies using input, forget, output gates, and cell states.
4. Backpropagation Through Time (BPTT): Computing exact analytical gradients for all gating and projection weights.
5. Autoregressive Text Generation: Generating text token-by-token using greedy decoding.

---

📂 Repository Directory Layout

• Project Root
  • README.md: This documentation file.
  • pyproject.toml: Project metadata, dependencies, and Pyright settings.
  • requirements.txt: Pip package requirements.
  • main.py: Entrypoint file.
• src/ — Core Modules
  • lstm/ — LSTM Model & Training
    • model.py: LSTM cell calculations, forward/backward unroll, gradient clipping, and SGD optimizer.
    • train.py: Preprocessing, embedding lookup, and training loop.
    • inference.py: Text generator script using custom prompt inputs.
  • data/ — Text Pipeline
    • load_text.py: Text loader
    • clean_text.py: Removes Gutenberg headers/footers, strips illustrations, and normalizes text.
    • tokenizer.py: Tokenizes string streams into word tokens.
    • vocabulary.py: Mappings between words and unique integer tokens.
  • dataset/ — Data Loaders
    • skipgram_dataset.py: Formats training pairs for Word2Vec.
    • sequence_dataset.py: Generates input sequence context and next-word targets.
  • word2vec/ — Embedding Model
    • model.py: Implements Skip-Gram Word2Vec training.
  • utils/ — Math Helpers
    • activations.py / losses.py: Softmax and Cross-Entropy loss implementations.

---

📐 Model Architecture Diagram

Below is the computational graph of the unrolled LSTMLanguageModel processing sequential embeddings to predict the next word token:

graph TD
    subgraph Input_Sequence ["Input Sequence Word IDs"]
        x1["x_1: 'The'"]
        x2["x_2: 'White'"]
        x3["x_3: 'Rabbit'"]
        x4["x_4: 'was'"]
    end

    subgraph Embedding_Lookup ["Embedding Matrix (Word2Vec)"]
        e1["e_1 = Embed(x_1)"]
        e2["e_2 = Embed(x_2)"]
        e3["e_3 = Embed(x_3)"]
        e4["e_4 = Embed(x_4)"]
    end

    x1 --> e1
    x2 --> e2
    x3 --> e3
    x4 --> e4

    subgraph LSTM_Cell_Transitions ["Unrolled LSTM Transitions"]
        c0["Cell State: c_0 = 0"]
        h0["Hidden State: h_0 = 0"]

        c0 & h0 & e1 --> LSTM_Cell_1["LSTM Cell 1"] --> c1["c_1"] & h1["h_1"]
        c1 & h1 & e2 --> LSTM_Cell_2["LSTM Cell 2"] --> c2["c_2"] & h2["h_2"]
        c2 & h2 & e3 --> LSTM_Cell_3["LSTM Cell 3"] --> c3["c_3"] & h3["h_3"]
        c3 & h3 & e4 --> LSTM_Cell_4["LSTM Cell 4"] --> c4["c_4"] & h4["h_4"]
    end

    subgraph Output_Layer ["Output Projection & Decoded Word"]
        logits["logits = h_4 * W_y + b_y"]
        probs["probabilities = Softmax(logits)"]
        next_word["next_word = argmax(probs) -> 'going'"]
    end

    h4 --> logits
    logits --> probs
    probs --> next_word

Loading

---

🧠 LSTM Cell Mathematics

For each time step $t$ given input word embedding $e_t$, the cell state $c_t$ and hidden state $h_t$ are updated using gating mechanisms:

1. Gating & Candidate Computation

• Forget Gate (controls context retention): $$f_t = \sigma(e_t W_{xf} + h_{t-1} W_{hf} + b_f)$$
• Input Gate (controls new information flow): $$i_t = \sigma(e_t W_{xi} + h_{t-1} W_{hi} + b_i)$$
• Candidate Cell State (proposed state update): $$\tilde{c}t = \tanh(e_t W{xc} + h_{t-1} W_{hc} + b_c)$$
• Output Gate (controls hidden state output): $$o_t = \sigma(e_t W_{xo} + h_{t-1} W_{ho} + b_o)$$

2. State Updates

• Updated Cell State: $$c_t = f_t \odot c_{t-1} + i_t \odot \tilde{c}_t$$
• Updated Hidden State: $$h_t = o_t \odot \tanh(c_t)$$

where $\sigma(z) = \frac{1}{1 + e^{-z}}$ is the sigmoid activation function, and $\odot$ represents element-wise multiplication.

---

📊 Training Configuration & Results

Hyperparameters

Parameter	Value
Corpus	Alice in Wonderland
Vocabulary Size	2,867 words
Embedding Size	100 dimensions
Hidden Size	128 units
Sequence Length	4 words
Learning Rate	0.01
Epochs	20

LSTM Loss Convergence

The model demonstrates smooth loss convergence over 20 epochs without vanishing/exploding gradients, thanks to manual gradient clipping:

Epoch 01/20 | Loss : 6.5736
Epoch 05/20 | Loss : 5.2296
Epoch 10/20 | Loss : 4.6417
Epoch 15/20 | Loss : 4.2277
Epoch 20/20 | Loss : 3.8658

---

📈 Model Comparison: Vanilla RNN vs. LSTM

The following table compares the generated text quality on various prompts after training on the same corpus under equivalent conditions:

Prompt	Vanilla RNN (50 Epochs)	LSTM (20 Epochs)
Alice	Fragmented, repetitive continuation.	"Alice said to herself I should have to be no chance"
The	Severe context/topic drift.	"The Gryphon went on in a low voice voice are old"
The White	Learns phrase but drifts quickly.	"The White Rabbit was going to find that it was the first"
The Queen	Unstable grammatical continuation.	"The Queen The Queen had been to the door she had never"
She	Weak grammatical structure.	"She said the Gryphon with a great hurry voice to her"
It	Mostly incoherent context flow.	"It the King said the Gryphon with a great hurry voice"

Key Observations

• Long-Term Dependencies: The LSTM retains context much better than the Vanilla RNN. For instance, when prompted with "The White", the LSTM correctly maintains the phrase "The White Rabbit" and constructs a semantically plausible sentence, whereas the Vanilla RNN drifts off-topic.
• Coherence: Even at fewer training epochs (20 vs 50), the gating mechanism of the LSTM mitigates the vanishing gradient problem, allowing it to organize grammar structure more naturally.

---

🚀 Setup & Execution

This project uses the modern Python package manager uv for environment setup and execution.

1. Installation & Environment Sync

Initialize the virtual environment and fetch dependencies:

uv sync

2. Preprocess & Train Word2Vec Embeddings

Generate the custom 100-dimensional embeddings used as input features:

uv run playground/phase1/test_word2vec.py

3. Run the Training Pipeline

Train the LSTM model on the tokenized sequences:

uv run src/lstm/train.py

This saves the token mapping to data/vocab.pkl and the trained weights to data/lstm_weights.npz.

4. Text Generation (Inference)

Run text generation on the pre-configured prompts:

uv run src/lstm/inference.py

---

🛑 Current Limitations

• Greedy Decoding: Uses simple argmax decoding, which causes occasional repetition loops.
• Fixed Sequence Length: Trained on a fixed sequence length of 4.
• Single-Layer: Does not utilize deep stacked LSTM layers.
• Basic Vocabulary: No handling of out-of-vocabulary (<UNK>) tokens or capitalization normalization.
• No Batching: Sequences are processed individually during training.