README.md

---
license: mit
language:
- en
datasets:
- roneneldan/TinyStories
base_model:
- openai-community/gpt2
pipeline_tag: text-generation
tags:
- text-generation-inference
- TinyStories
- miniGPT
- llama
- educational
- tensorflow
- keras
library_name: keras
---

# 🧠 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:

```text
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.

```text
x → RMSNorm → Attention
```

and subsequently:

```text
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:

```text
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:

```text
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:

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

This structure is repeated **5 times**.

---

## 🎯 Output Projection & Weight Tying

After the final decoder block:

```text
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:

```text
(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*.

---