Author(s): Naimish Baranwal

Originally published on Towards AI.

The Rise of the SLMs: Building the Smallest Possible Language Model.

This is the first step toward building SLM(Small Language models).

This work replicates the findings reported in https://arxiv.org/abs/2305.077

The referenced research seeks to develop the smallest language model capable of processing English sentences. For coding, I took help from Andrej Karpathy’s nanoGPT repo https://github.com/karpathy/nanoGPT.

In this paper, researchers aim to create the smallest language model that can understand English words and generate fluent, consistent English text.

This is the first step towards building small language models, where I am pre-training the language model from scratch.

Why Small is the New Big in AI, and how to build one from scratch.

SLMs deliver faster inference, lower latency, and real-time performance ideal for edge devices and chatbots. They cut costs dramatically with reduced compute, energy, and memory needs, suiting SMEs and on-device AI. Privacy improves as they run locally without cloud dependency, and specialization boosts accuracy for targeted tasks.

I am using the dataset TinyStories, which is available on Hugging Face. Anyone can use the model directly from Hugging Face, available in different versions: https://huggingface.co/roneneldan/TinyStories-33M

• TinyStories-1M/3M/9M/28M/33M

I pre-train my whole model on NVIDIA GeForce RTX 5050 8GB VRAM, which took approximately one hour for 20,000 iterations. Also, for better accuracy, I ran for 50,000 iterations, which took approximately 3 hours. You can use Jupyter Notebook’s T4 GPU; it might take the same time to run, or if you have RAM, you can reduce the iterations.

Also, I use the optuna library to do the hyperparameter optimization.

Data Preparation:

In the prepare_data() function, data is loaded from Hugging Face, split into train.bin and validation.bin files, and tokenized using the GPT-2 tokenizer from tiktoken. A memory-mapped function enables efficient RAM usage during processing, while datasets are downloaded in batches for optimized performance.

In the get_batch() function, the batch size is set to 32 with a block size of 128. This means each iteration processes 128 tokens per sequence across 32 sequences, totaling 4,096 tokens per batch (32 × 128). Running 50,000 iterations processes approximately 204 million tokens—about 22% of one epoch from the TinyStories training set's roughly 900 million tokens—while still yielding decent results.

Model Architecture

The Block class implements a single transformer block comprising a LayerNorm layer, CausalSelfAttention mechanism, and a feedforward neural network (MLP).

Attention Mechanism

In the CausalSelfAttention class, causal masking is applied manually, while Flash Attention is utilized automatically when available (PyTorch 2.0+). The embedding size must be fully divisible by the number of heads; here, 384/6 = 64 dimensions per head.

Training Process

The estimate_loss() function employs gradient accumulation over 32 steps before updating weights. Every 500 iterations, it computes validation and training losses by averaging over 100 batches each.

GPT Model Assembly

The GPT class serves as the main container, combining token and positional embeddings, then passing them through 6 transformer blocks (Block class). Each block features 6 parallel causal attention heads that capture diverse perspectives—such as syntax, semantics, and local context—while the sequential blocks enable hierarchical learning.

Inference Optimization

The @torch.no_grad() decorator disables gradient tracking within inference functions for improved efficiency.

Hyperparameter optimization

Optuna systematically explores the search space, learns from each experiment, and converges on optimal settings automatically using Bayesian-inspired optimization. I use it to find the best embedding size, number of attention heads, number of transformer layers, learning rate, and dropout.

Training Results

After completing the training iterations, the best model parameters are saved in a .pt file. During inference, the system automatically loads weights from best_model_params.pt to generate responses using the optimized mode.

Here are some responses from the model:

You can find all the code and details in the following repo:

GitHub – NaimishKB/SLM–TLM-from-scratch

Contribute to NaimishKB/SLM–TLM-from-scratch development by creating an account on GitHub.

github.com

Happy Learning!!

Join thousands of data leaders on the AI newsletter. Join over 80,000 subscribers and keep up to date with the latest developments in AI. From research to projects and ideas. If you are building an AI startup, an AI-related product, or a service, we invite you to consider becoming a sponsor.

Published via Towards AI