Pico-LLM

Built three autoregressive language models from scratch in PyTorch — k-gram MLP, LSTM, and Transformer with self-attention and KV-cache decoding. Designed a complete training and inference pipeline with next-token prediction and nucleus (top-p) sampling to benchmark generation quality and efficiency.

What This Is

Pico-LLM is a from-scratch exploration of how language models learn to predict the next token — and how architectural choices shape what they learn. Built as part of NYU's CSCI-GA.2565 (Machine Learning) course with other three team members, this project goes beyond the required core tasks to investigate attention interpretability, normalization strategies, positional encoding effects, and scaling behavior across custom datasets.

Rather than importing a pretrained model and fine-tuning, this project builds three distinct architectures from the ground up and puts them through the same training and evaluation pipeline to compare how they behave.

The three models:

• K-gram MLP — A fixed-context window model that concatenates the previous k token embeddings and passes them through an MLP to predict the next token. Implemented as a sequence-to-sequence model with a sliding window operation. Tested with both one-hot and nn.Embedding inputs to demonstrate why learned embeddings outperform sparse representations.
• LSTM — A recurrent model that maintains a hidden state across the full sequence, capturing long-range dependencies without an explicit context window. Used as the baseline architecture provided in the starter code.
• Transformer — A causal decoder-only Transformer with multi-head self-attention, RMSNorm (not LayerNorm), positional embeddings, and KV-cache decoding for efficient autoregressive generation. Architecture references GPT-2, Llama 3, and Karpathy's minGPT.

All three are trained on TinyStories and a simple synthetic sequence dataset (3seqs.txt), with an optional custom Wikipedia corpus for extended evaluation.

What's Inside

pico-llm/
├── pico-llm.py              # Main training + inference script (all three models)
├── core_tasks.md             # Full experimental writeup with plots and analysis
├── 3seqs.txt                 # Synthetic sequence data for sanity checks
└── trained_outputs/          # Training curves, attention maps, generation samples
    ├── outputs_embedding/
    ├── outputs_onehot/
    ├── outputs_3seqs_fullpattern/
    ├── outputs_tinystories_full/
    ├── outputs_wiki_512/
    ├── outputs_wiki_1024/
    ├── outputs_tiny_underfit/
    ├── outputs_tiny_overfit/
    ├── outputs_postnorm_tinystories/
    ├── outputs_no_positional_emd_postnorm_tinystories/
    ├── plots_interpretability/
    └── hyperparams/

Core Tasks

These were the required implementations for the course:

1. Baseline Sanity Check (LSTM)

Verified the training loop, data pipeline, and generation code work as expected on both TinyStories and 3seqs.txt. Both train and test loss decrease over epochs, confirming a working LSTM baseline.

2. K-gram MLP Implementation

Built the KGramMLPSeqModel with a sliding window .forward() routine. Compared one-hot vs. nn.Embedding inputs — one-hot training loss barely moves, while the embedding version converges quickly, demonstrating why learned representations matter.

3. Nucleus (Top-p) Sampling

Implemented top-p sampling inside generate_text. Evaluated across multiple values (greedy, 0.2, 0.5, 0.75, 0.95, 1.0) on both datasets. Lower values produce near-deterministic output; higher values introduce diversity at the cost of occasional incoherence.

4. Transformer Implementation

Built a causal decoder-only Transformer from scratch with multi-head self-attention, RMSNorm, skip connections, positional embeddings, and KV-cache decoding. Trained on both 3seqs.txt (near-zero loss) and TinyStories (signs of overfitting in later epochs).

Optional Extensions

Beyond the core requirements, the following experiments were completed:

Custom Dataset (Wikipedia)

Trained the Transformer on a 30K-line subset of the HuggingFace Wikipedia corpus to study behavior on longer, information-dense, factual text. Observed faster convergence with larger embeddings (1024 vs. 512) and clear overfitting patterns distinct from TinyStories.

Overfitting vs. Underfitting Analysis

Systematically demonstrated both regimes by varying epochs and training steps. Underfitting (3 epochs, 100 steps/epoch) produces incoherent, broken sentences. Overfitting (15 epochs, 2500 steps/epoch) produces fluent but memorized, low-diversity output. Includes both quantitative (loss curves) and qualitative (generated text) analysis.

Hyperparameter Sweeps

Studied embedding dimensionality (512 vs. 1024) — larger embeddings converge faster due to greater expressive power. Also analyzed Transformer depth/width (16 heads, 8 blocks) on both TinyStories and Wikipedia, showing how model capacity relative to dataset size drives overfitting behavior.

Positional Embedding Effects

Compared Transformer with and without positional embeddings. Without them, the causal mask preserves temporal direction but not position — the model degenerates into a bag-of-words distribution, producing heavily repetitive output. With them, generation is coherent and narratively structured.

Attention Head Interpretability

Visualized attention patterns across Layer 3's four heads and found clear specialization:

• Head 0 — Grammatical structure and local context establishment
• Head 1 — Subject-verb-object relationship tracking
• Head 2 — Named entity recognition and action dependencies
• Head 3 — High-level cross-sentence dependency linking

Pre-Norm vs. Post-Norm

Benchmarked both normalization strategies. Post-norm achieves lower training and test loss on small datasets like TinyStories by stabilizing gradient flow and preventing the residual branch from dominating early in training.

Quick Start

Sanity check (CPU, fast):

python3 pico-llm.py \
  --tinystories_weight 0.0 \
  --input_files 3seqs.txt \
  --prompt "0 1 2 3 4" \
  --block_size 32 \
  --embed_size 128 \
  --max_steps_per_epoch 10 \
  --device_id cpu \
  --output_dir outputs_sanity

Full training on TinyStories (GPU recommended):

python3 pico-llm.py \
  --tinystories_weight 1.0 \
  --prompt "Once upon a time" \
  --block_size 512 \
  --embed_size 512 \
  --kgram_k 3 \
  --max_steps_per_epoch 750 \
  --kgram_chunk_size 16 \
  --batch_size 16 \
  --epochs 10 \
  --learning_rate 3e-4 \
  --test_fraction 0.1 \
  --device_id cuda:0 \
  --output_dir outputs_tinystories

Requirements

• Python 3.8+
• PyTorch
• NumPy
• Matplotlib

Full Writeup

See core_tasks.md for the complete experimental log — training curves, generation samples, attention heatmaps, hyperparameter sweeps, and detailed analysis for every experiment.

Context

Built for NYU CSCI-GA.2565 (Machine Learning). The course provided starter code with an LSTM baseline and a training loop — everything else (k-gram MLP, Transformer, nucleus sampling, all optional extensions) was implemented from scratch.