ARX1A07/miniGPT_Project

🧠 MiniGPT — TinyStories

A compact LLaMA-inspired decoder-only Transformer trained from scratch in TensorFlow/Keras on the TinyStories dataset.

The model is designed as a small-scale experiment in implementing and training a modern decoder-only language model with components such as RoPE, RMSNorm, SwiGLU, causal self-attention, frozen pretrained embeddings, and weight tying.

---

📋 Model Configuration

Parameter	Value
Architecture	Decoder-only Transformer
Architecture style	LLaMA-inspired
Parameters	85,791,744 (~85.8M)
Transformer layers	5
Hidden dimension (d_model)	768
Attention heads	12
Head dimension	64
FFN hidden dimension	3072
Vocabulary size	50,257
Context length	256 tokens
Tokenizer	GPT-2 BPE
Framework	TensorFlow / Keras
Precision	Mixed bfloat16

---

🏗️ Model Architecture

The model follows a compact LLaMA-style decoder architecture:

GPT-2 BPE Tokenizer
        │
        ▼
Token IDs
        │
        ▼
Frozen GPT-2 Embedding Matrix
        │
        ▼
┌─────────────────────────────────────┐
│       Decoder Block × 5              │
│                                     │
│  RMSNorm                            │
│     │                               │
│     ▼                               │
│  Causal Multi-Head Self-Attention   │
│     │                               │
│     ├── Q                           │
│     ├── K ──► RoPE                  │
│     ├── V                           │
│     │                               │
│     └── Attention + Residual        │
│                                     │
│  RMSNorm                            │
│     │                               │
│     ▼                               │
│  SwiGLU Feed-Forward Network        │
│     │                               │
│     └── Residual                    │
└─────────────────────────────────────┘
        │
        ▼
Final RMSNorm
        │
        ▼
Weight-Tied Language Modeling Head
        │
        ▼
Logits [sequence_length × 50,257]

🔩 Decoder Block

Each of the 5 decoder layers operates on a hidden representation of size 768.

1. RMSNorm

The input representation is normalized before each major sub-layer using RMSNorm rather than LayerNorm.

x → RMSNorm → Attention

and subsequently:

x → RMSNorm → SwiGLU

This gives the block the pre-normalized structure commonly used by modern decoder-only architectures.

2. Causal Multi-Head Self-Attention

The model uses:

• 12 attention heads
• 64 dimensions per head
• 12 × 64 = 768 total attention dimension
• causal masking to prevent access to future tokens

The hidden representation is projected into:

Q = XWq
K = XWk
V = XWv

with attention operating independently across the 12 heads before the heads are merged back into the original 768-dimensional representation.

3. Rotary Positional Embeddings

RoPE is applied to the query and key representations inside the attention mechanism.

This provides positional information without requiring a learned positional embedding table.

4. SwiGLU Feed-Forward Network

The attention output is followed by a SwiGLU feed-forward block with an intermediate dimension of 3072.

Conceptually:

x
│
├── Linear ──► SiLU ──┐
│                     × ──► Linear ──► output
└── Linear ───────────┘

The resulting representation is returned to the residual stream.

5. Residual Connections

Each major sub-layer uses a residual connection:

x → RMSNorm → Attention → + x
                         │
                         ▼
                       x'
                         │
             RMSNorm → SwiGLU → + x'

This structure is repeated 5 times.

---

🎯 Output Projection & Weight Tying

After the final decoder block:

Hidden States
     │
     ▼
 Final RMSNorm
     │
     ▼
 Language Modeling Head
     │
     ▼
50,257 vocabulary logits

The language-modeling projection is weight-tied with the input token embedding representation, reducing redundant parameters and following the common tied-embedding approach.

The resulting tensor has the conceptual shape:

(batch, sequence_length, 50,257)

---

📊 Parameter Distribution

The final Keras model contains:

Category	Parameters
Total	85,791,744
Trainable	47,194,368
Non-trainable	38,597,376

The large non-trainable component corresponds primarily to the frozen pretrained GPT-2 embedding matrix.

---

📚 Training Configuration

Parameter	Configuration
Dataset	TinyStories
Training token budget	100M tokens
Validation token budget	5M tokens
Context length	256
Batch size	16
Epochs configured	3
Optimizer	AdamW
Weight decay	0.01
Gradient clipping	clipnorm=1.0
Maximum LR	1e-4
Minimum LR	1e-5
Warmup	2,000 steps
LR schedule	Warmup + cosine decay
Precision	Mixed bfloat16

---

📈 Evaluation

Metric	Result
Validation loss / Perplexity / Generation evaluation	Yet to be collected and evaluated

---

🛠️ Implementation & Deployment

The model itself is implemented and trained in TensorFlow/Keras.

The notebook contains the experimental training implementation and model construction.

• 📓 Training / model-construction notebook: [<COLAB LINK>](https://colab.research.google.com/drive/1DUSM24y3hcrc06BduC47kyWnciu8OGZg?usp=sharing)
• 🐳 Docker-based inference server: [<GITHUB LINK>](https://github.com/abhayramasamy/MiniGPT_Project)
• 🤗 Hugging Face model / artifacts: [<HF LINK>](https://huggingface.co/ARX1A07/miniGPT_Project/)

The Hugging Face integration and generation pipeline are intentionally documented separately from the core architecture. The Docker repository contains the deployment-oriented inference implementation.

---

⚠️ Limitations

This is a relatively small language model trained specifically on TinyStories. Its capabilities should therefore be interpreted within the scope of the dataset and parameter scale.

It should not be considered a general-purpose language model or compared directly with substantially larger pretrained LLMs.

Evaluation and generation-quality benchmarking are still pending.

---

📚 References

• Eldan, R. & Li, Y. — TinyStories: How Small Can Language Models Be and Still Speak Coherent English? (2023).
• Hugging Face — TinyStories dataset.
• Vaswani et al. — Attention Is All You Need.
• Su et al. — RoFormer: Enhanced Transformer with Rotary Position Embedding.
• Shazeer — GLU Variants Improve Transformer.

---