Aesop-LLM

A custom Bengali Language Model built from scratch using PyTorch, featuring advanced transformer architectures including Multi-Head Latent Attention, RoPE positional encodings, and Mixture of Experts layers.

Overview

Aesop-LLM is a transformer-based language model designed for Bengali text generation. It implements several modern LLM techniques to achieve efficient training and inference while maintaining quality output.

Key Implementation Topics

1. Multi-Head Latent Attention (MHLA)

The model uses Multi-Head Latent Attention instead of traditional multi-head attention. This approach compresses key-value pairs through a low-rank projection before expanding them back, significantly reducing memory usage and enabling longer context handling.

• KV Compression: Projects keys and values through a compressed dimension
• Latent Space: Uses a compressed latent space for efficient attention computation
• Implementation: MultiHeadLatentAttention.py

2. Decoupled Rotary Positional Embedding (DeRoPE)

The model implements DeRoPE for positional encoding, which provides better position awareness compared to traditional positional embeddings.

• Rotary Attention: Applies rotation to query and key vectors
• Decoupled Design: Separates positional information from content
• Absolute Position Support: Supports KV caching with absolute position offsets
• Implementation: DeRoPE.py

3. Mixture of Experts (MoE)

The feed-forward network uses a Mixture of Experts architecture with SwiGLU activation.

• Top-k Gating: Selects the top-k experts for each token
• SwiGLU Activation: Combines Swish and Gated Linear Unit
• 4 Experts: Configurable number of expert networks
• Implementation: NeedClasses.py (SwiGLU, Expert, MoE classes)

4. Normalization Layers

The model uses two types of normalization:

• LayerNorm: Standard layer normalization with learnable scale and shift
• RMSNorm: Root Mean Square normalization (more efficient, no shift parameter)
• Implementation: NeedClasses.py

5. Custom BPE Tokenizer

A Byte-Pair Encoding tokenizer specifically trained for Bengali text.

• Unicode Normalization: NFKC normalization for consistent encoding
• Special Tokens: Supports <eos> (end-of-sentence) token
• Vocabulary Size: 5,515 tokens (configurable)
• Implementation: utils.py, TrainTokenizer.py

6. KV Cache for Fast Inference

Efficient inference through key-value caching:

• Token-by-Token Generation: Reuses computed keys and values
• Position Tracking: Maintains correct absolute positions for RoPE
• Implementation: TextGeneration.py

7. Training Features

• AdamW Optimizer: Weight decay regularization
• Early Stopping: Prevents overfitting with patience mechanism
• Gradient Clipping: Stabilizes training
• Sliding Window: Efficient training with stride-based batching
• Implementation: train.py

Project Structure

Aesop-LLM/
├── main.py                    # Main training script
├── train.py                   # Training loop and evaluation
├── model.py                   # AesopLLM model definition
├── TextGeneration.py          # Text generation utilities
├── Dataset.py                 # Dataset and DataLoader
├── TransformerBlock.py        # Transformer block implementation
├── MultiHeadLatentAttention.py # MHLA implementation
├── DeRoPE.py                  # Rotary positional encoding
├── NeedClasses.py             # Normalization, MoE, SwiGLU
├── utils.py                   # Text processing utilities
├── DataClean.py               # Data cleaning utilities
├── TrainTokenizer.py          # Tokenizer training
├── TestTokenizer.py           # Tokenizer testing
├── tokenizer/
│   └── bpe_tokenizer.json     # Pre-trained BPE tokenizer
├── requirements.txt           # Dependencies
└── README.md                  # This file

Setup

Prerequisites

• Python 3.8+
• PyTorch
• tokenizers library (Hugging Face)

Installation

1. Clone the repository

   cd Aesop-LLM

2. Install dependencies

   pip install torch tokenizers

3. Prepare your data

   • Place your Bengali text file as merged.txt in the project root
   • Or update the path in main.py to point to your data

Model Configuration

Default configuration in main.py:

Parameter	Default	Description
vocab_size	5515	Vocabulary size
context_length	256	Maximum sequence length
emb_dim	200	Embedding dimension
num_heads	8	Number of attention heads
num_layers	6	Number of transformer layers
dropout	0.1	Dropout probability
rope_dim	24	RoPE embedding dimension
compressed_dim	50	KV compression dimension

Usage

Training the Model

Run the main training script:

python main.py

This will:

1. Load and preprocess the Bengali text data
2. Split into 90% training and 10% validation
3. Train the model for 1 epoch (configurable)
4. Save checkpoints to model/AesopLLM_checkpoint.pth

Text Generation

Use the trained model for text generation:

from TextGeneration import generate_text, generate_text_kv
from tokenizers import Tokenizer
from model import AesopLLM
import torch

# Load tokenizer
bpe_tokenizer = Tokenizer.from_file("tokenizer/bpe_tokenizer.json")

# Load model
CONFIG = {
    "vocab_size": 5515,
    "context_length": 256,
    "emb_dim": 200,
    "num_heads": 8,
    "num_layers": 6,
    "dropout": 0.1,
    "rope_dim": 24,
    "compressed_dim": 50,
    "bias_qkv": False
}
model = AesopLLM(CONFIG)
checkpoint = torch.load("model/AesopLLM_checkpoint.pth")
model.load_state_dict(checkpoint['model_state_dict'])

# Generate text (basic)
generated = generate_text(
    model=model,
    tokenizer=bpe_tokenizer,
    prompt="আজকে সকালে",
    max_length=100,
    top_k=25,
    temperature=0.1,
    device='cuda'
)

# Generate text with KV cache (faster)
generated_kv = generate_text_kv(
    model=model,
    tokenizer=bpe_tokenizer,
    prompt="আজকে সকালে",
    max_length=100,
    top_k=25,
    temperature=0.1,
    device='cuda'
)

Training a New Tokenizer

If you need to train a tokenizer on your own data:

from utils import train_and_save_bpe_tokenizer, read_text_file

text = read_text_file("your_data.txt")
train_and_save_bpe_tokenizer(
    text,
    vocab_size=5515,
    out_dir="tokenizer/bpe_tokenizer.json"
)

Hyperparameters

Training Hyperparameters

Parameter	Recommended Value	Description
Learning Rate	5e-5	Start with small LR for large models
Weight Decay	0.01	L2 regularization
Batch Size	10	Training batch size
Max Length	256	Sequence length
Stride	5	Sliding window stride
Epochs	1-10	Number of training epochs

Generation Hyperparameters

Parameter	Recommended Value	Description
max_length	100-500	Tokens to generate
top_k	25-50	Top-k sampling parameter
temperature	0.1-1.0	Lower = more deterministic

License

See LICENSE file for details.