Fall 2026
CSC-372 Project Part 4

Gradient Descent

Tiny Stories Dataset

The TinyStories dataset, introduced by Ronen Eldan and Yuanzhi Li (Microsoft Research) in their seminal 2023 paper “TinyStories: How Small Can Language Models Be and Still Speak Coherent English?”, is a synthetic corpus designed to probe the lower bounds of language model capacity.

The dataset comprises approximately 2.1 million short stories (hundreds of words each), synthetically generated using GPT-3.5 and GPT-4. The generation process enforces strict constraints:

Vocabulary Restriction: Stories are restricted to words commonly understood by 3- to 4-year-old children (a working vocabulary of roughly 3,000-4,000 basic terms).
Controlled Prompting: Prompts systematically sample combinations of random words, target moral themes, structural twists, or specific narrative features to ensure syntactic variety while keeping concepts foundational.
Coherence Over Breadth: Because the domain discards real-world factual trivia, complex technical jargon, and convoluted sentence structures, the learning objective focuses squarely on narrative flow, pronoun resolution, simple physical reasoning, and grammar.

Prior to TinyStories, the prevailing consensus in natural language processing held that emergent linguistic capabilities such as producing grammatically sound, semantically consistent narrative text—required massive parameter scales (often tens or hundreds of billions) trained on vast web scrapes like Common Crawl. TinyStories challenged this paradigm by isolating semantic and syntactic coherence from encyclopedic world knowledge.

TinyStories shifted how researchers evaluate small language models (SLMs) and model interpretability in several ways:

Sub-100M Parameter Viability: The authors demonstrated that models with fewer than 10 million parameters—and even single-layer or small transformer networks—can produce coherent, multi-paragraph stories with near-flawless grammar and logical progression when trained on this data.
Accessible Benchmarking & Education: Full pre-training runs on TinyStories can be completed in hours on consumer-grade GPUs or modest compute setups, lowering the compute barrier for rapid prototyping, loss-curve experimentation, and classroom demonstration.
Mechanistic Interpretability: Because the models trained on TinyStories are tiny (e.g., 1M to 33M parameters) yet generate meaningful text, they serve as ideal “model organisms” for dissecting attention heads, circuit formation, and hidden-layer representations without the intractable complexity of LLMs.
Data Quality vs. Data Scale: The dataset provided early empirical evidence for the “textbooks are all you need” philosophy—showing that high-quality, synthetic, curriculum-like data can dramatically substitute for sheer raw parameter volume.

A representative story in the dataset follows a simple, archetypal narrative arc:

“Once upon a time, there was a little boy named Tim. Tim had a bright yellow ball. One day, Tim took his ball to the park. He threw the ball high into the air, but the ball landed in a tall tree. Tim was sad. A big dog named Sam walked over and saw Tim crying. Sam barked and jumped, bumping the tree trunk. The yellow ball fell down to the grass. Tim hugged Sam and smiled.”

In short, TinyStories demonstrates that fluency, grammar, and elementary reasoning are not exclusive properties of giant models, establishing an efficient sandbox for training, evaluating, and interpreting compact language architectures.

The repository justjuu/tinystories-gpt2-shards on Hugging Face provides a pre-tokenized, sharded version of the TinyStories dataset specifically formatted for training GPT-2 architectures.

While the original dataset consists of raw synthetic text files, training language models efficiently at scale requires converting text into token ID arrays and chunking them into fixed context-window sequences. By pre-encoding the stories using the standard GPT-2 byte-pair encoding (BPE) tokenizer and packaging them into manageable shards (often stored as memory-mapped binary files or Arrow tables), this repo eliminates the repetitive CPU overhead and memory bottlenecks typically involved in preprocessing raw text prior to training.

Because the data is already tokenized and structured into continuous sequence blocks, it serves as a lightweight, plug-and-play resource for rapid model experimentation. Researchers, educators, and developers building nanoGPT-style architectures, running mechanistic interpretability probes, or benchmarking custom small language models can stream or load these shards directly into PyTorch DataLoaders with zero setup time. It dramatically accelerates the pipeline from raw script to active training loop, making it especially convenient for consumer-hardware runs where pre-training cycles need to be fast, reproducible, and computationally lean.

Here, to make our lives even easier, we are only going to use only the validation part of the dataset.

The following code cell downloads the validation split of the tiny stories dataset pre-tokenized using GPT-2 tokenizer:

import numpy as np
import tiktoken
from huggingface_hub import hf_hub_download

repo = "justjuu/tinystories-gpt2-shards"

# Download the smaller validation split first.
path = hf_hub_download(
    repo_id=repo,
    filename="validation.bin",
    repo_type="dataset",
)

np.memmap creates a memory-mapped array, allowing you to read or write portions of a large file directly on disk without loading the entire file into RAM.

# The file is a continuous stream of little-endian uint16 token IDs.
tokens = np.memmap(path, dtype="<u2", mode="r")
print(f"{len(tokens):,} tokens")
print("First 20 IDs:", tokens[:20].tolist())

Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.
WARNING:huggingface_hub.utils._http:Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.

4,765,918 tokens
First 20 IDs: [32565, 13, 15899, 2497, 262, 22441, 1097, 290, 531, 11, 366, 22017, 11, 21168, 11, 534, 1097, 318, 523, 6016]

The following cell code decodes the first story using the tiktoken GPT-2 tokenizer

# 50256 is the special End of Text (<|endoftext|>) token ID in GPT-2.
# In datasets like TinyStories, this token marks the boundary between individual stories.
# np.flatnonzero(...) returns the indices of all True values.
end = np.flatnonzero(tokens == 50256)[0]

enc = tiktoken.get_encoding("gpt2")
story = enc.decode(tokens[:end].astype(int).tolist())
print(story)

Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.
WARNING:huggingface_hub.utils._http:Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.

4,765,918 tokens
First 20 IDs: [32565, 13, 15899, 2497, 262, 22441, 1097, 290, 531, 11, 366, 22017, 11, 21168, 11, 534, 1097, 318, 523, 6016]
Spot. Spot saw the shiny car and said, "Wow, Kitty, your car is so bright and clean!" Kitty smiled and replied, "Thank you, Spot. I polish it every day."

After playing with the car, Kitty and Spot felt thirsty. They found a small pond with clear water. They drank the water and felt very happy. They played together all day and became best friends.

enc.n_vocab

Data Loader and DataSet

Note that is almost identical to the DataLoader from Project Part 2.

Implement GPTDatasetV1 to split the token sequence into input windows of 10 tokens, advancing by 4 tokens each time.

Pair each window with the single token immediately following it as the target, and store both as PyTorch tensors. Note that in project part 2, both \(x_i\) and \(y_i\) were of the same length, whereas here \(y_i\) is a single token. That is the target sequence is now of length 1 instead of matching the input window length. That is the only change.

Implement __len__ to return the number of examples and __getitem__ to retrieve an input–target pair.

Create a DataLoader with a batch size of 50, shuffling enabled, drop_last=True to discard incomplete batches, and num_workers=0.

Select CUDA if available; otherwise, use the CPU.

import torch
from torch.utils.data import Dataset, DataLoader
import math

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class GPTDatasetV1(Dataset):
    def __init__(self, token_ids, context_length, stride):
      self.input_ids = []
      self.target_ids = []

      raise NotImplementedError("GPTDatasetV1 __init__ method is not implemented.")

    def __len__(self):
      return len(self.input_ids)

    def __getitem__(self, idx):
      return self.input_ids[idx], self.target_ids[idx]

batch_size  = 50
context_window  = 10
stride      = 4
shuffle     = True
drop_last   = True
num_workers = 0

dataset = GPTDatasetV1(tokens, context_window, stride)

dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        drop_last=drop_last,
        num_workers=num_workers
    )

Multi-Layer Perceptron (MLP) Model

Implement a PyTorch multilayer perceptron (MLP) that predicts the next token from a fixed window of preceding tokens. Use a token embedding layer, two hidden layers with ReLU activation, and an output layer that produces logits.

Create a class named MLP that inherits from torch.nn.Module.

1. `init()`

Its constructor should accept

input_size \(D_i\): number of tokens in the context window
hidden_size \(D\):
output_size \(D_o\): the vocabulary size of the tokenizer used

and call super().__init__().

1.1. Hyperparameters

Set the learning rate to 5e-4 and the embedding dimension to 128.

1.2. Create the embedding layer

Use torch.nn.Embedding to assign each vocabulary token a trainable vector of length 128.

Each input batch should contain integer token IDs with shape (batch_size, input_size). The embedding layer should produce a tensor with shape (batch_size, input_size, 128).

1.2. Initialize the weights and biases.

Set in_dim0 = input_size * embedding_dim.

Define the following trainable tensors using torch.nn.Parameter:

Parameter	Shape	Initialization
`Omega0`	`(in_dim0, hidden_size)`	Standard normal values multiplied by `sqrt(2 / in_dim0)`
`Beta0`	`(hidden_size,)`	Zeros
`Omega1`	`(hidden_size, hidden_size)`	Standard normal values multiplied by `sqrt(2 / hidden_size)`
`Beta1`	`(hidden_size,)`	Zeros
`Omega2`	`(hidden_size, output_size)`	Standard normal values multiplied by `sqrt(2 / hidden_size)`
`Beta2`	`(output_size,)`	Zeros

For Omega1, use hidden_size in the initialization scale because that is the number of inputs to this layer.

2. `forward()`

Implement the forward pass.

Convert the input token IDs to integer type using x.long() and retrieve their embeddings. Flatten the embeddings starting at dimension 1, preserving the batch dimension.

Compute the two hidden layers and the output:

h1 = torch.relu(x @ self.Omega0 + self.Beta0)
h2 = torch.relu(h1 @ self.Omega1 + self.Beta1)
logits = h2 @ self.Omega2 + self.Beta2

Return the logits with shape (batch_size, output_size). Do not apply an activation function to the output layer.

\[ h_1 = a[\beta_{0} + \Omega_{0} E[x]] \] \[ h_k = a[\beta_{k-1} + \Omega_{k-1} h_{k-1}] \] \[ y = \beta_k + \Omega_k h_k \]

3. `parameters()`

Return all the omega, beta parameters as well as embedding layer parameters as one list.

4. `step()`

Implement a manual gradient descent step.

Add a step() method. Inside a torch.no_grad() block, iterate through self.parameters(). For each parameter with a non-None gradient, subtract the learning rate multiplied by its gradient.

\[ \theta \leftarrow \theta - \eta \nabla_\theta L(\theta) \]

5. `zero_grad()`

Clear gradients between training iterations. Set the gradient .grad of all parameters to None.

class MLP(torch.nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        raise NotImplementedError("MLP __init__ method is not implemented.")

    def forward(self, x):
        raise NotImplementedError("MLP forward method is not implemented.")

    def parameters(self):
      raise NotImplementedError("MLP parameters method is not implemented.")

    def step(self):
        raise NotImplementedError("MLP step method is not implemented.")

    def zero_grad(self):
        raise NotImplementedError("MLP zero_grad method is not implemented.")

Softmax

Softmax is a function that converts a vector of raw scores (logits) into probabilities. For instance, given a matrix of shape \(N \times C\), where \(N\) is the number of examples and \(C\) is the number of classes, the softmax function is applied row-wise to produce a matrix where each row sums to 1.

\[ p_j = \operatorname{softmax}(z)_j = \frac{e^{z_j}}{\sum_{k=1}^{C} e^{z_k}} \]

Here, \(p_j\) represents the probability assigned to class \(j\) for a given example, and \(z_j\) is the corresponding raw score (logit).

Take log of the softmax probabilities to obtain the log-probabilities, we get:

\[ \log p_j = z_j - \log\left(\sum_{k=1}^{C} e^{z_k}\right) \]

Here we implement a numerically stable log_softmax(logits) function that returns log-probabilities for each example across its classes. Use the identity \[\log p_j = z_j - m - \log\sum_k e^{z_k-m}\]

where \(m=\max_k z_k\) is computed separately for each example. Subtracting the maximum before exponentiating prevents overflow, while computing log-probabilities directly avoids taking the logarithm of probabilities that may have rounded to zero through underflow.

import torch

def log_softmax(logits):
    # Shift logits so the largest value in each row is zero
    shifted = logits - logits.max(dim=1, keepdim=True).values

    # Compute log-probabilities directly
    log_normalizer = torch.log(
        torch.exp(shifted).sum(dim=1, keepdim=True)
    )

    return shifted - log_normalizer

Cross Entropy Loss

Here we also implement cross_entropy_from_log_probs(log_probs, targets) to compute the average cross-entropy loss for a batch.

The input log_probs contains the predicted log-probabilities for each class, and targets contains the correct class index for each example.

The log-probability of the correct class in each row is selected and then the negative mean of these values is returned.

\[ \mathcal{L} = -\frac{1}{N}\sum_{i=1}^{N}\log p_{i,y_i} \]

Here, \(N\) is the batch size, \(y_i\) is the correct class for example \(i\), and \(\log p_{i,y_i}\) is its selected value from log_probs.

def cross_entropy_from_log_probs(log_probs, targets):
    correct_log_probs = log_probs[
        torch.arange(len(targets), device=log_probs.device),
        targets
    ]

    return -correct_log_probs.mean()

Training

Instantiate your MLP using context_window as the input size, 256 hidden units per hidden layer, and enc.n_vocab as the output size, then move the model to the selected device.

Train the model for atleast five epochs, using tqdm to display progress through the dataloader.

For each batch, move the inputs and targets to the device, clear previous gradients using mlp.zero_grad(), and run the forward pass to obtain output logits.

Compute the loss using your own numerically stable cross-entropy implementation, with targets converted to integer class indices using .long().

You are not allowed to use built-in loss functions or optimizers, including torch.nn.CrossEntropyLoss, torch.nn.functional.cross_entropy, or any optimizer from torch.optim, such as Adam or SGD.

Implement the loss calculation yourself and update the model parameters through your own mlp.step() method using gradient descent.

You may use PyTorch automatic differentiation by calling loss_value.backward() to compute gradients. Record each batch’s loss using .item() and print the epoch number and mean batch loss at the end of each epoch.

from tqdm import tqdm

mlp = NotImplementedError()

for epoch in range(5):

    losses = []

    for x_i, y_i in tqdm(dataloader):
        
        raise NotImplementedError()

    print(f"Epoch: {epoch} Loss: {np.mean(losses)}")

Inference / Next Token Prediction

This code generates text one token at a time using the trained MLP.

It takes the first input sequence from a batch, moves it to the selected device, and prints the decoded sequence as the initial prompt.

Calling mlp.eval() enables evaluation mode, while torch.no_grad() disables gradient tracking during generation.

On each of the 10 iterations, the model predicts logits for the next token, and torch.argmax selects the token with the highest score.

The selected token is decoded and printed immediately, then appended to the input after removing its oldest token, keeping the context window the same length.

This process lets each prediction become part of the context for the next prediction. The final print() adds a newline.

import torch

x_i, y_i = next(iter(dataloader))
x_i = x_i[0:1].to(device)

mlp.eval()
with torch.no_grad():
    # Print the full initial prompt once
    print(enc.decode(x_i[0].tolist()), end="==>", flush=True)

    for i in range(10):  # Generate 10 tokens
        logits = mlp(x_i)

        pred_token_tensor = torch.argmax(logits, dim=-1, keepdim=True).to(torch.uint16)
        pred_token_id = pred_token_tensor.item()

        # Stream only the newly generated token
        print(enc.decode([pred_token_id]), end="", flush=True)

        # Slide the window
        x_i = torch.cat((x_i[:, 1:], pred_token_tensor), dim=1)

print()  # Add final newline

Fall 2026 CSC-372 Project Part 4