TinyLM

✓ Status: Complete

A learning lab for teaching myself to build a character-level language model from scratch.

Outcome

This project achieved its original goal: building a character-level language model from first principles while understanding why each major component exists and how the architecture naturally evolves from the problem.

Progress

TinyLM is complete and consists of four progressively more capable implementations:

• Stage 1 — Count-Based Language Model
• Stage 2 — Linear Neural Language Model
• Stage 3 — Embeddings
• Stage 4 — Multi-Layer Perceptron (MLP)

Each stage introduces exactly one new architectural idea while preserving everything learned previously.

Stage 1: Count-Based Language Model

Before building a neural network, I built a complete count-based character language model from first principles. The goal was to understand the language modeling problem before introducing machine learning.

Architecture

Raw text
    │
    ▼
Tokenizer
    │
    ▼
Token IDs
    │
    ▼
Dataset Builder
    │
    ▼
(Context, Target) Training Examples
    │
    ▼
CountLanguageModel
    │
    ▼
Probability Distribution
    │
    ▼
Generator
    │
    ▼
Generated Token IDs
    │
    ▼
Tokenizer.decode()
    │
    ▼
Generated Text

Components

• Tokenizer

  • Builds a fixed vocabulary from a corpus.
  • Encodes text into token IDs.
  • Decodes token IDs back into text.

• Dataset Builder

  • Converts a sequence of token IDs into (context, target) training examples.
  • Uses a fixed-size sliding context window.
  • Pads the beginning of a sequence with <START> tokens.

• CountLanguageModel

  • Learns how often each target token follows each context.
  • Converts counts into probability distributions.
  • Cannot generalize beyond contexts seen during training.

• Generator

  • Repeatedly queries the model for the next-token probability distribution.
  • Samples the next token.
  • Updates the context window and continues generation.

Key Concepts Learned

• A language model predicts one token at a time.
• Long sequences are generated by repeatedly predicting the next token.
• A corpus can be transformed into many supervised training examples.
• Probabilities are conditional on the current context.
• Sampling from a probability distribution produces non-deterministic text.
• Exact-count models memorize rather than generalize.

Limitations Discovered

• The model cannot predict unseen contexts.
• The vocabulary is fixed after training.
• Generation currently has no explicit end-of-sequence token, revealing the need for an <END> token in the next iteration.

Stage 2: Neural Language Model

The count table was replaced with a learnable neural network. Rather than memorizing counts, the model learns a function that maps a context to a probability distribution over the next token.

Architecture

Raw text
    │
    ▼
Tokenizer
    │
    ▼
Token IDs
    │
    ▼
Dataset Builder
    │
    ▼
(Context, Target) Training Examples
    │
    ▼
One-Hot Encoding
    │
    ▼
Flatten
    │
    ▼
Linear Layer
    │
    ▼
Logits
    │
    ▼
Softmax
    │
    ▼
Probability Distribution
    │
    ▼
Cross-Entropy Loss
    │
    ▼
Reverse-Mode Autodiff
    │
    ▼
Gradient Descent

Components

• LanguageModel
  • Maps a fixed-size context to logits.
  • Owns learnable weights and biases.
• Softmax
  • Converts logits into a probability distribution.
• Cross-Entropy Loss
  • Measures how much probability the model assigned to the correct next token.
• Autograd
  • Tracks computations and propagates gradients through the computation graph.
• Training Loop
  • Performs forward pass, loss computation, backpropagation, and parameter updates.

Key Concepts Learned

• A count table can be replaced by a learnable function.
• Logits represent unnormalized preferences over the vocabulary.
• Softmax converts logits into probabilities.
• Cross-entropy provides a learning signal based on the probability assigned to the correct token.
• Reverse-mode autodiff computes gradients automatically.
• Gradient descent updates parameters to reduce loss.

Limitations Discovered

• One-hot vectors cannot express similarity between tokens.
• Input dimensionality grows with vocabulary size.
• The model is still purely linear.

Stage 3: Embeddings

One-hot vectors were replaced with learned embeddings. Instead of representing each token as a sparse identity vector, every token now owns a learned dense vector that is optimized together with the rest of the model.

Architecture

Context
    │
    ▼
Embedding Lookup
    │
    ▼
Flatten
    │
    ▼
Linear Layer
    │
    ▼
Logits
    │
    ▼
Softmax
    │
    ▼
Cross-Entropy Loss
    │
    ▼
Gradient Descent

Key Concepts Learned

• Embeddings are simply learnable parameter vectors.
• An embedding table is a lookup table indexed by token ID.
• Embeddings reduce the input dimensionality from context_size × vocabulary_size to context_size × embedding_size.
• Embeddings are trained jointly with the language model.
• The overall training pipeline remains unchanged; only the input representation changes.

Limitations Discovered

• The model is still fundamentally a linear model.
• More expressive relationships require a non-linear hidden layer.

Stage 4: Multi-Layer Perceptron (MLP)

The single linear layer was replaced with a two-layer neural network separated by a non-linear activation function. This allows the model to learn relationships that cannot be represented by a purely linear transformation.

Architecture

Context
    │
    ▼
Embedding Lookup
    │
    ▼
Flatten
    │
    ▼
Linear Layer
    │
    ▼
Tanh
    │
    ▼
Linear Layer
    │
    ▼
Logits
    │
    ▼
Softmax
    │
    ▼
Cross-Entropy Loss
    │
    ▼
Gradient Descent

Components

• EmbeddingTable
  • Stores one learned embedding vector for every token.
• Linear
  • Reusable linear transformation used for both the hidden and output layers.
• Tanh
  • Introduces non-linearity between the two linear layers.
• LanguageModel
  • Composes the embedding table, hidden layer, activation function, and output layer.

Key Concepts Learned

• Stacking linear layers alone does not increase model capacity.
• Non-linear activation functions allow the network to learn more complex relationships.
• Extracting a reusable Linear layer simplifies the overall architecture.
• The language model becomes a composition of reusable building blocks.

Outcome

The MLP model successfully trains end-to-end, the loss converges during training, and the trained model correctly generates the learned sequence ("hello").

Conclusion

TinyLM is intentionally complete.

Its purpose was to understand how a language model evolves from memorization to learning by introducing one architectural idea at a time.

By the end of the project, the following questions have been answered:

• How does a count-based language model work?
• How can memorization be replaced by a learnable neural network?
• Why do embeddings replace one-hot vectors?
• Why are non-linear activation functions necessary?
• How do gradient descent and backpropagation train a language model?

This project intentionally stops before transformers. The next project starts from a clean slate to build a transformer architecture from first principles, using TinyLM as the conceptual foundation rather than extending it further.