Building Google’s Gemma 3 From Scratch

What Are We Building and Why?

Fig 1: Complete Gemma 3 architecture: 164.6M trainable parameters across 18 transformer blocks.

The Problem This Blog Solves

There are thousands of tutorials showing you how to load a pre-trained model in three lines of Python and generate text. That is useful for building demos. But it teaches you absolutely nothing about what is happening inside. When your model generates garbage, when inference is too slow for production, when you need to modify the architecture for a specific domain — you are completely helpless. You are a consumer, not an engineer.

This blog exists to make you an engineer. Over the next 22 chapters, we will build a complete language model from scratch — not a toy model, but the exact same architecture that Google uses in their production Gemma 3 models. We will write every single component by hand in PyTorch. No pre-built transformer libraries. No importing someone else’s weights. We begin with 164.6 million random floating-point numbers and end with an artificial intelligence system that writes coherent children’s stories with characters, dialogue, emotions, and moral lessons.

By the end, you will understand every matrix multiplication, every normalization step, every attention score, and every gradient update that transforms random noise into language understanding. More importantly, you will understand why each of these operations exists — what problem it solves, what would break without it, and how the researchers who invented it were thinking.

What Exactly Is a Language Model?

Strip away all the hype, and a language model is a mathematical function that takes a sequence of words as input and outputs a probability for every possible next word. If the input is “The cat sat on the,” the model outputs something like: “mat” with probability 0.15, “floor” with 0.12, “chair” with 0.08, “roof” with 0.03, and so on for all 256,128 words in its vocabulary.

That is it. That is the entire job. Every impressive capability you have ever seen from ChatGPT, Claude, Gemini, or any other AI — conversations, translations, code generation, poetry, reasoning — all of it emerges from this single capability: predicting the next word, applied over and over again.

When you ask a chatbot “What is the capital of France?”, the model is not “looking up” the answer in a database. It is predicting, word by word, what text is most likely to follow your question based on patterns it absorbed from billions of training examples. The prediction for the next word after “The capital of France is” has overwhelmingly high probability for “Paris” because that pattern appeared thousands of times in training data.

Why Build From Scratch?

Understanding. When a model produces strange outputs, when fine-tuning fails, when quantization degrades quality — debugging requires understanding the internal mechanics. The difference between an ML engineer who can debug a transformer and one who cannot is the difference between someone who built one from scratch and someone who only called APIs.

Modification. Real-world applications often require architectural changes. Maybe you need a model that processes images alongside text. Maybe you need one that runs on a microcontroller with 512KB of memory. All of these require understanding the architecture well enough to modify it.

Depth of Knowledge. In job interviews, technical discussions, and research, the depth of understanding that comes from implementation is immediately apparent. The person who can explain why RoPE uses rotation rather than addition, or why GeGLU outperforms ReLU, demonstrates a fundamentally different level of understanding.

The Seven Innovations We Will Implement

Gemma 3 is not “just another transformer.” It contains seven key innovations over the original 2017 Transformer architecture, each solving a specific problem:

1. Multi-Query Attention: Shares a single Key and Value across all attention heads, keeping only Queries separate. Reduces the KV cache by 4× during inference, enabling deployment on smaller hardware with only ~0.5% quality loss.

2. Sliding Window Attention: Restricts each token to attending only to its nearest 512 neighbors instead of the full sequence. Reduces attention cost from O(n²) to O(n×512) — a 64× saving. Only 3 of 18 layers use full global attention.

3. Dual-Base RoPE: Uses fast rotation (base=10,000) for 15 sliding-window layers focusing on local grammar, and slow rotation (base=1,000,000) for 3 global layers tracking long-range dependencies. Each layer gets the position resolution it needs.

4. QK Normalization: Applies RMSNorm to both Q and K before computing attention, preventing score magnitudes from growing uncontrollably and keeping softmax distributions well-behaved.

5. (1+γ) RMSNorm: Initializes normalization scale to 0 and uses (1+γ) as the actual scale. At initialization this is identity, providing exceptional stability during early training.

6. GeGLU Feed-Forward: Replaces ReLU with a gated mechanism where one projection computes activation and another controls how much of each feature passes through — fundamentally more expressive per-input feature selection.

7. √dim Embedding Scaling: Multiplies embeddings by √640 ≈ 25.3 to bring their magnitude to the same scale as the residual stream after 18 layers of accumulation.

Our Model vs Google’s Production Model

How Language Models Actually Think: Next-Word Prediction

The Mathematical Foundation

A language model learns a conditional probability distribution. Given a sequence of tokens x₁, x₂, …, xₜ, it outputs the probability of every possible next token:

Here f_θ is our neural network (the transformer) with learnable parameters θ. The function takes a sequence of tokens and outputs a vector of 256,128 raw numbers called logits — one per vocabulary token. The softmax function converts these raw logits into a valid probability distribution:

Why Exponential?

Simple division (zᵢ/Σzⱼ) fails because logits can be negative, giving negative “probabilities.” Taking absolute values loses ordering information: a logit of −5 would get the same probability as +5. The exponential function maps all real numbers to strictly positive values, preserves ordering, is smooth and differentiable everywhere, and maximizes entropy subject to matching expected logit values.

The Autoregressive Generation Process

When generating text, the model works token by token. Each new token depends on ALL previous tokens — this is why it is called autoregressive:

def generate(model, prompt_tokens, max_new_tokens=200, temperature=0.7):
    tokens = prompt_tokens.clone()
    for _ in range(max_new_tokens):
        logits = model(tokens)                        <span class="cmt"># Forward pass</span>
        logits = logits[:, -1, :] / temperature       <span class="cmt"># Scale last position</span>
        probs = torch.softmax(logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)
        tokens = torch.cat([tokens, next_token], dim=1)
        if next_token.item() == eos_token_id:
            break
    return tokens

Cross-Entropy Loss — How We Measure “Wrong”

During training, we know the actual next word. Cross-entropy loss measures how far the model’s predictions are from reality:

If the model assigns probability 0.9 to the correct token: L = −log(0.9) = 0.105 (small loss). If it assigns 0.01: L = −log(0.01) = 4.605 (large loss). If probability approaches 0: L approaches infinity.

From information theory: cross-entropy measures the expected number of extra bits needed to encode reality using the model’s probability distribution. Our model’s final loss of 1.78 means ~1.78 extra bits per token compared to the true distribution.

Perplexity: The Human-Interpretable Metric

Perplexity answers: “On average, how many words is the model effectively choosing between at each step?” Our final model achieved perplexity 5.96, meaning at each position the model narrows 256,128 vocabulary tokens down to about 6 equally likely candidates.

Tokens, Embeddings, and Vocabulary — Converting Text to Numbers

The Fundamental Problem: Machines Cannot Read

Neural networks perform mathematical operations: addition, multiplication, matrix products. They cannot operate on the letter “A” or the word “cat.” We need a systematic way to convert arbitrary text into sequences of numbers that preserve meaningful relationships.

Fig 2: The three-stage pipeline: text → token IDs → embedding vectors → scaled embeddings.

Stage 1: Tokenization — Why Not Just Use Words?

Problem 1: Infinite vocabulary. New words appear constantly. A word-level vocabulary cannot handle anything it has not seen during training.

Problem 2: Morphological blindness. “run,” “running,” “runner,” and “runs” would be four completely separate tokens with no shared representation.

Problem 3: Memory explosion. English has over 170,000 words in common use. The embedding table alone would consume most of the model’s parameters.

Gemma 3 uses SentencePiece with a vocabulary of 256,128 subword tokens. Common words like “the” get one token. Rare words get split into multiple subword pieces. This gives a fixed, finite vocabulary that can represent any text.

Stage 2: Embedding Lookup

Each token ID maps to a dense vector of 640 floating-point numbers through a simple table lookup. This table has dimensions [256,128 × 640], containing 163,921,920 parameters — 99.6% of our model’s trainable parameters.

Initially random, through training semantically similar tokens develop similar vectors. “Cat” ends up close to “kitten” in 640-dimensional space. This geometric organization of meaning is called an embedding space.

Stage 3: √dim Scaling — Gemma’s Crucial Innovation

The transformer’s residual stream accumulates values through 18 layers of addition. Without scaling, the initial embeddings would be tiny compared to accumulated residual values. Multiplying by √dim brings the embedding magnitude to the same scale as the residual stream.

class GemmaEmbedding(nn.Module):
    def __init__(self, vocab_size=256128, dim=640):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, dim)
        self.scale = dim ** 0.5                       <span class="cmt"># √640 ≈ 25.3</span>

    def forward(self, token_ids):
        x = self.embedding(token_ids)                 <span class="cmt"># Lookup: (batch, seq, 640)</span>
        return x * self.scale                         <span class="cmt"># Scale to match residual stream</span>

Architecture Overview — The Complete Picture

Before diving into individual components, let us understand the complete data flow. A transformer block processes a sequence of token embeddings through two sub-layers: attention (tokens communicate with each other) and feed-forward (each token thinks independently). Gemma 3 stacks 18 of these blocks.

Fig 3: Data flow through a single Gemma 3 transformer block — every operation annotated.

Pre-Norm: Why Normalize Before, Not After

The original 2017 Transformer applied normalization after each sub-layer (post-norm). GPT-2 switched to pre-norm, and virtually all modern architectures follow. Pre-norm creates a clean residual stream:

With post-norm, gradients must flow through normalization during backpropagation. With pre-norm, the residual connection provides a “gradient highway” — gradients flow directly from loss to any layer without passing through normalization or activation functions.

Layer Configuration: 15 Sliding + 3 Global

Not all 18 layers are identical. 15 layers use sliding window attention (nearest 512 tokens), while every 6th layer (5, 11, 17) uses full global attention.

RMSNorm — The (1+γ) Innovation from First Principles

Why Normalize at All?

Without normalization, activation magnitudes drift unpredictably through layers. Some dimensions explode while others shrink. This “internal covariate shift” makes training unstable.

BatchNorm → LayerNorm → RMSNorm

BatchNorm (2015) normalizes across the batch dimension. Works for CNNs but fails for language models because batch statistics are noisy with variable-length sequences.

LayerNorm (2016) normalizes across the feature dimension. For each token, computes mean and variance of its 640-dimensional vector, then centers and scales.

RMSNorm (Zhang & Sennrich, 2019) removes the mean subtraction — only scales by root-mean-square. Simpler, faster, and works just as well:

Gemma’s Innovation: (1+γ) Initialization

Standard RMSNorm initializes γ to 1.0. Gemma 3 initializes γ to 0.0 and uses (1+γ) as the actual scale factor. At initialization: scale = (1+0) = 1, making the layer a pure identity function. The network starts as close to identity as possible and only learns modifications — exceptional stability in the first thousand training steps.

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.zeros(dim))  <span class="cmt"># γ init to 0, NOT 1!</span>

    def forward(self, x):
        rms = torch.sqrt(torch.mean(x.float() ** 2, dim=-1, keepdim=True) + self.eps)
        x_normalized = x.float() / rms
        return (x_normalized * (1.0 + self.weight)).to(x.dtype)  <span class="cmt"># (1+γ) scaling</span>

Fig 4: Left: RMSNorm behavior with (1+γ) at initialization vs after training. Right: GeGLU vs ReLU activation.

RoPE — Rotary Position Encoding from First Principles

The Position Problem

Attention is position-agnostic by default. If you give it tokens in any order, it computes the same scores. But word order obviously matters: “the dog bit the man” means something entirely different from “the man bit the dog.”

Historical Approaches and Their Problems

Sinusoidal encoding (2017): Fixed sine/cosine waves added to embeddings. Cannot easily learn relative positions.

Learned position embeddings (GPT-2): Separate learnable embedding per position. Cannot generalize beyond trained length; positions are absolute.

RoPE’s Elegant Solution: Encode Position Through Rotation

RoPE (Su et al., 2021) rotates each pair of dimensions in Q and K vectors by an angle proportional to the token’s position:

The key property: when computing the dot product of a rotated query at position i with a rotated key at position j, the angles cancel to depend only on (i−j), not absolute positions. The model naturally learns relative position relationships.

Fig 5: RoPE rotation concept, dual-base frequencies, and position sensitivity by layer type.

Dual Base — Gemma 3’s Innovation

base = 10,000 for sliding-window layers → fast rotation → fine-grained local position.

base = 1,000,000 for global layers → slow rotation → distinguish tokens thousands of positions apart.

def precompute_rope_frequencies(head_dim, max_seq_len, base=10000.0):
    freqs = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
    positions = torch.arange(max_seq_len)
    angles = torch.outer(positions, freqs)
    return torch.cos(angles), torch.sin(angles)

def apply_rope(x, cos, sin):
    x1, x2 = x[..., ::2], x[..., 1::2]             <span class="cmt"># Even/odd dimension pairs</span>
    return torch.cat([
        x1 * cos - x2 * sin,                        <span class="cmt"># Rotate</span>
        x1 * sin + x2 * cos
    ], dim=-1)

Multi-Query Attention — Complete Derivation

Attention from First Principles

Attention answers: “For each token, which other tokens should I pay attention to, and how much?” It works in three steps:

Step 1 — Project: Transform each token into Query (“What am I looking for?”), Key (“What do I contain?”), Value (“What information do I carry?”).

Step 2 — Score: Compute compatibility between every query-key pair via dot product, divided by √d_k to prevent softmax saturation.

Step 3 — Aggregate: Apply softmax to get attention weights, then multiply by values.

Fig 6: Multi-Head vs Multi-Query attention, sliding window mask, and compute savings.

Why Multiple Heads?

A single head can only learn one type of relationship per layer. Language has many simultaneous relationships: syntactic, semantic, positional, referential. Our model uses 4 heads, each with head_dim = 256, attending to different relationship types in parallel.

Why Multi-Query? The KV Cache Problem

During inference, the model needs Keys and Values for ALL previous tokens at every step. Multi-Query shares a single K and V across all 4 heads, reducing cache by 4× with only ~0.5% quality loss.

class MultiQueryAttention(nn.Module):
    def __init__(self, dim=640, n_heads=4, head_dim=256):
        super().__init__()
        self.n_heads, self.head_dim = n_heads, head_dim
        self.scale = head_dim ** -0.5

        self.wq = nn.Linear(dim, n_heads * head_dim, bias=False)  <span class="cmt"># 4 separate Q</span>
        self.wk = nn.Linear(dim, head_dim, bias=False)             <span class="cmt"># 1 shared K</span>
        self.wv = nn.Linear(dim, head_dim, bias=False)             <span class="cmt"># 1 shared V</span>
        self.wo = nn.Linear(n_heads * head_dim, dim, bias=False)
        self.q_norm = RMSNorm(head_dim)
        self.k_norm = RMSNorm(head_dim)

    def forward(self, x, mask=None, rope_cos=None, rope_sin=None):
        B, S, _ = x.shape
        q = self.wq(x).view(B, S, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.wk(x).view(B, S, 1, self.head_dim).transpose(1, 2)
        v = self.wv(x).view(B, S, 1, self.head_dim).transpose(1, 2)

        q, k = self.q_norm(q), self.k_norm(k)       <span class="cmt"># QK normalization</span>
        q, k = apply_rope(q, rope_cos, rope_sin), apply_rope(k, rope_cos, rope_sin)
        k = k.expand(-1, self.n_heads, -1, -1)      <span class="cmt"># Broadcast to all heads</span>
        v = v.expand(-1, self.n_heads, -1, -1)

        scores = (q @ k.transpose(-2, -1)) * self.scale
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        out = (torch.softmax(scores, dim=-1) @ v)
        return self.wo(out.transpose(1, 2).reshape(B, S, -1))

Sliding Window Attention — From O(n²) to O(n)

The Quadratic Problem

Standard self-attention computes scores between every pair of tokens — an n×n matrix, O(n²) computation. For 32,768 tokens: over one billion attention scores per layer. Most of this is wasted.

The Observation: Attention Is Local

Research consistently shows that attention weights concentrate on nearby tokens. Grammar relationships are within 5–10 positions, phrase structure within 20–50, paragraph coherence within 200–500. Long-range dependencies exist but are sparse.

Sliding Window: Restrict to Local Context

Each token attends only to its nearest w = 512 neighbors. The attention mask becomes a band matrix. Computation drops from O(n²) to O(n×w):

But What About Long-Range Dependencies?

Every 6th layer (5, 11, 17) uses full global attention. These 3 layers handle long-range dependencies, while 15 sliding layers handle local structure cheaply. Information propagates globally through the stack: layer 5 broadcasts long-range info → layers 6–10 refine locally → layer 11 re-synchronizes → and so on.

def create_sliding_window_mask(seq_len, window_size=512):
    mask = torch.tril(torch.ones(seq_len, seq_len))  <span class="cmt"># Causal</span>
    for i in range(seq_len):
        mask[i, :max(0, i - window_size + 1)] = 0   <span class="cmt"># Restrict to window</span>
    return mask

<span class="cmt"># In the model forward pass:</span>
for idx, layer in enumerate(self.layers):
    is_global = (idx % 6 == 5)                       <span class="cmt"># Layers 5, 11, 17</span>
    mask = global_mask if is_global else sliding_mask
    rope_base = 1_000_000 if is_global else 10_000
    x = layer(x, mask=mask, rope_base=rope_base)

GeGLU Feed-Forward Networks — Gated Feature Selection

The Feed-Forward Layer’s Job

After attention lets tokens communicate, the feed-forward network lets each token think independently. It processes each token’s 640-dimensional representation through a wider 2,048-dimensional space, applies a non-linearity, and projects back. This is where the model stores “factual knowledge” — associations and patterns from training.

The Evolution: ReLU → GELU → GeGLU

ReLU (2017 Transformer)

If positive, pass through; if negative, zero out. Simple but crude — no middle ground. Creates “dead neurons” that output zero for all inputs.

GELU (GPT-2, BERT)

Smooths the transition using a Gaussian CDF. Eliminates dead neurons, but applies the same non-linearity uniformly.

GeGLU (Shazeer, 2020) — What Gemma 3 Uses

Two parallel projections: one computes “content” (what information to carry), the other computes a “gate” (how much of each feature to let through). The gate is input-dependent — different inputs activate different feature combinations.

class GeGLUFeedForward(nn.Module):
    def __init__(self, dim=640, hidden_dim=2048):
        super().__init__()
        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
        self.up_proj   = nn.Linear(dim, hidden_dim, bias=False)
        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)

    def forward(self, x):
        gate = torch.nn.functional.gelu(self.gate_proj(x))  <span class="cmt"># What to let through</span>
        up   = self.up_proj(x)                               <span class="cmt"># The information</span>
        return self.down_proj(gate * up)                     <span class="cmt"># Gated → project back</span>

GeGLU requires 3 projection matrices instead of 2 — 50% more FFN parameters. But it typically improves perplexity by 0.3–0.5 points, well worth the cost.

Putting It All Together — The Full Transformer Block

The Complete Data Flow

class GemmaTransformerBlock(nn.Module):
    def __init__(self, dim=640, n_heads=4, head_dim=256,
                 ffn_hidden=2048, is_global=False):
        super().__init__()
        self.is_global = is_global
        self.attn_norm = RMSNorm(dim)
        self.attention = MultiQueryAttention(dim, n_heads, head_dim)
        self.ffn_norm = RMSNorm(dim)
        self.feed_forward = GeGLUFeedForward(dim, ffn_hidden)

    def forward(self, x, mask, rope_cos, rope_sin):
        x = x + self.attention(self.attn_norm(x), mask, rope_cos, rope_sin)
        x = x + self.feed_forward(self.ffn_norm(x))
        return x

The Full Gemma 3 Model

class Gemma3Model(nn.Module):
    def __init__(self, vocab_size=256128, dim=640, n_layers=18,
                 n_heads=4, head_dim=256, ffn_hidden=2048):
        super().__init__()
        self.embedding = GemmaEmbedding(vocab_size, dim)
        self.layers = nn.ModuleList([
            GemmaTransformerBlock(
                dim, n_heads, head_dim, ffn_hidden,
                is_global=(i % 6 == 5)               <span class="cmt"># Layers 5, 11, 17</span>
            ) for i in range(n_layers)
        ])
        self.final_norm = RMSNorm(dim)
        self.output_proj = nn.Linear(dim, vocab_size, bias=False)
        self.output_proj.weight = self.embedding.embedding.weight  <span class="cmt"># Weight tying</span>

    def forward(self, token_ids, targets=None):
        x = self.embedding(token_ids)
        for layer in self.layers:
            base = 1_000_000 if layer.is_global else 10_000
            cos, sin = precompute_rope_frequencies(256, x.size(1), base)
            x = layer(x, mask, cos.to(x.device), sin.to(x.device))
        logits = self.output_proj(self.final_norm(x))
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) if targets is not None else None
        return logits, loss

Parameter Count Breakdown

The Training Pipeline — DataLoader, Optimizer, Scheduler

The TinyStories Dataset

We train on TinyStories (Eldan & Li, 2023), ~471 million tokens across ~2.1 million short children’s stories generated by GPT-3.5/GPT-4. Each story features characters, dialogue, emotions, and moral lessons — designed to evaluate what small language models can learn.

The DataLoader

class TinyStoriesDataLoader:
    def __init__(self, data, batch_size=32, seq_len=512):
        self.data = data
        self.batch_size, self.seq_len = batch_size, seq_len
        self.pos = 0

    def get_batch(self):
        B, S = self.batch_size, self.seq_len
        buf = self.data[self.pos : self.pos + B * S + 1]
        x = buf[:-1].view(B, S)                     <span class="cmt"># Input:  tokens [0..n-1]</span>
        y = buf[1:].view(B, S)                      <span class="cmt"># Target: tokens [1..n]</span>
        self.pos += B * S
        if self.pos + B * S + 1 > len(self.data):
            self.pos = 0
        return x.cuda(), y.cuda()

AdamW Optimizer

Adam with decoupled weight decay. Maintains running averages of gradient mean and squared gradients per parameter. Settings: lr=3e-4, β₁=0.9, β₂=0.95, weight_decay=0.1, ε=1e-8.

Cosine Schedule with Warmup

Warmup (steps 0–500): LR ramps linearly from 0 to 3e-4. Lets optimizer moment estimates stabilize before aggressive learning.

Cosine decay (steps 500–13,000): LR follows cosine from 3e-4 down to 3e-5. Aggressive learning mid-training, gentle refinement at end.

Mixed Precision Training & Gradient Accumulation

bfloat16 Mixed Precision

Float32 uses 32 bits per number. bfloat16 uses 16 — halving memory and doubling throughput. We use mixed precision: forward/backward in bf16, optimizer maintains float32 master copy for small gradient updates.

Why bfloat16 over float16? bfloat16 has the same exponent range as float32 (8 bits), just less mantissa precision. No overflow/underflow issues — simpler and more stable.

Gradient Accumulation

We want effective batch size 128, but GPU fits only 32. Run 4 micro-batches of 32, accumulating gradients, then update once. Mathematically identical to batch 128.

accumulation_steps = 4                               <span class="cmt"># 4 × 32 = effective batch 128</span>

for step in range(max_steps):
    optimizer.zero_grad()
    for micro_step in range(accumulation_steps):
        x, y = train_loader.get_batch()
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            logits, loss = model(x, y)
            loss = loss / accumulation_steps          <span class="cmt"># Scale for averaging</span>
        loss.backward()                              <span class="cmt"># Accumulate gradients</span>
    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
    optimizer.step()

Gradient Clipping

Clips gradients to max norm 0.5. Prevents catastrophic weight updates from occasional large gradients. Preserves gradient direction but limits step size.

The Training Run — From Random Noise to Language

Fig 7: Training/validation loss curves, perplexity progression, and cosine learning rate schedule.

The Journey: Loss Curve Milestones

What Emergence Looks Like

Step 0: “thr%%k si &&jjjjj wqq the the the.” Pure gibberish.

Step 500: “the the was a a little the was.” Common tokens but no grammar.

Step 1,000: “The little was very happy. She the to her.” Grammar emerges.

Step 5,000: “Once upon a time, there was a little girl named Lily.” Story templates.

Step 13,000: “Once upon a time, there was a little girl named Lily. She loved to play in the garden with her dog, Max. One sunny day, Lily found a beautiful flower. ‘Look, Max!’ she said. ‘It is so pretty!’ Max wagged his tail and barked happily. Lily picked the flower and brought it to her mom. Her mom smiled and said, ‘What a lovely gift, Lily.’ Lily felt very happy.” Full coherence.

Monitoring & Diagnosing Training Problems

Key Signals to Watch

Training loss: Should decrease steadily. Spikes are normal but transient.

Validation loss: Should track training loss. Divergence = overfitting.

Gradient norm: Should be stable (~0.3–0.5 after warmup). Spikes = instability.

Common Problems and Fixes

Our training was remarkably smooth — no spikes, no divergence. Directly attributable to (1+γ) RMSNorm initialization, QK normalization, and cosine warmup schedule. Good architecture makes training easy.

Inference — The Generation Algorithm

Temperature and Sampling

Temperature T controls creativity. T<1 sharpens the distribution (more deterministic). T>1 flattens it (more random). Top-k restricts sampling to the k most likely tokens. Our best config: T=0.7, top_k=50.

Example Output (T=0.7, top_k=50)

“Once upon a time, there was a little girl named Lily. She loved
to play in the garden with her dog, Max. One sunny day, Lily found
a beautiful flower. 'Look, Max!' she said. 'It is so pretty!' Max
wagged his tail and barked happily. Lily picked the flower and
brought it to her mom. Her mom smiled and said, 'What a lovely
gift, Lily.' Lily felt very happy.”

The KV Cache Optimization

During generation, recomputing attention for all previous tokens at every step is wasteful. The KV cache stores previously computed Keys and Values, so each new token only needs to compute its own Q/K/V and attend to the cached history. This reduces generation from O(n²) to O(n) per token.

Multi-Query Attention makes this especially efficient: we only cache 1 K and 1 V per layer instead of 4 each, cutting cache memory by 4×.

Results — Honest Benchmarking

Fig 8: Perplexity comparison across models and parameter efficiency analysis.

Fig 9: Resource efficiency comparison — our model achieves best efficiency across all metrics.

Honest Limitations

Coherence degrades beyond ~100 tokens — characters drift, plots loop. No complex multi-step reasoning. Vocabulary limited to children’s stories. These limitations motivate the entropy-based innovation in Part V.

Deployment to HuggingFace

Publishing the Model

We package the trained model for easy reuse by the community. The HuggingFace Hub provides versioned model hosting, automatic download, and standardized loading interfaces.

What We Upload

Model weights: The final checkpoint at step 13,000 — all 164.6M parameters saved as a PyTorch state dict.

Configuration: A JSON file specifying all hyperparameters (dim=640, n_layers=18, n_heads=4, etc.) so the architecture can be reconstructed exactly.

Tokenizer: The SentencePiece model file so users can encode/decode text identically.

Model card: Documentation covering training data, evaluation results, intended use, and limitations.

<span class="cmt"># Upload to HuggingFace</span>
from huggingface_hub import HfApi
api = HfApi()
api.upload_folder(
    folder_path="./model_checkpoint",
    repo_id="G3nadh/gemma3-270m-tinystories",
    repo_type="model"
)

<span class="cmt"># Load and generate in 5 lines</span>
from huggingface_hub import hf_hub_download
path = hf_hub_download("G3nadh/gemma3-270m-tinystories", "pytorch_model.bin")
model = load_model(path, device="cuda")
print(generate_text(model, "Once upon a time", max_tokens=200, temperature=0.7))

Model available at: huggingface.co/G3nadh/gemma3-270m-tinystories

The Coherence Problem — Why Small Models Struggle

The ~100 Token Wall

Our model generates beautiful, coherent text for the first 60–100 tokens. Then something breaks. Characters change names mid-sentence. The plot loops back to the beginning. New characters appear from nowhere. The story loses direction and devolves into repetitive patterns.

This is not a bug — it is a fundamental limitation of small language models. The model’s 640-dimensional hidden state simply cannot maintain all the information needed for a long narrative: character names, relationships, plot progress, emotional arcs, setting details, unresolved conflicts.

Why Existing Solutions Don’t Work for Small Models

Longer context windows: Our model supports 32,768 tokens, but the issue is not context length — it is the model’s ability to use that context effectively. A small model with a long window is like a student with a 1,000-page textbook but poor reading comprehension.

RecurrentGPT, Re3, DOC: These iterative generation methods work well but all require 7B+ parameter models. They rely on the model’s ability to follow complex instructions, maintain outlines, and self-evaluate — capabilities that 164M parameter models simply do not have.

Simple chunking (fixed 80 tokens): Generate 80 tokens, stop, paste the last sentence as a prompt, continue. This helps but the chunk boundaries are arbitrary — you might cut mid-thought, mid-sentence, or right when the model is on a creative roll.

The Research Gap

Nobody has addressed coherent long-form generation specifically for sub-500M parameter models. The field assumes you need a large model for long text. We challenge this assumption.

Entropy-Based Adaptive Scene Cutting — The Core Innovation

Shannon Entropy as a Coherence Signal

At every generated token, the model outputs a probability distribution over the vocabulary. Shannon entropy measures the “spread” of this distribution:

When confident: entropy is LOW (probability concentrates on few tokens). When confused: entropy is HIGH (probability spreads across many tokens). We use this as a real-time coherence signal.

Fig 10: Entropy over tokens with adaptive cut points — spikes trigger scene boundaries.

The Complete Pipeline

1. Seed with story prompt.

2. Generate tokens, computing entropy at each step.

3. Track rolling average (window = 5).

4. When entropy > threshold (3.0) OR spike > 1.8× rolling average: trim to last complete sentence, end scene.

5. Extract last 2–3 sentences as context bridge (carries character names, situation, emotional state).

6. Seed next scene with bridge.

7. Repeat until target length or natural ending.

Fig 11: Context bridges carrying character and narrative state between adaptive scenes.

def generate_adaptive_scene(model, prompt, config):
    tokens = tokenize(prompt)
    entropy_history = []
    for step in range(config.max_tokens):
        logits = model(tokens)
        probs = softmax(logits[-1])
        entropy = -sum(p * log2(p) for p in probs if p > 0)
        entropy_history.append(entropy)
        rolling_avg = mean(entropy_history[-5:])

        if step >= config.min_tokens:
            if entropy > config.threshold:             <span class="cmt"># Threshold cut</span>
                return cut("threshold")
            if entropy > 1.8 * rolling_avg:            <span class="cmt"># Spike cut</span>
                return cut("spike")
            if all(e > 2.0 for e in entropy_history[-5:]):
                return cut("sustained")                <span class="cmt"># Sustained confusion</span>

        tokens.append(sample(probs, T=0.7, top_k=50))
    return cut("max_tokens")

Why This Works

Entropy is an internal signal from the model about its own confidence. We are not imposing external heuristics — we are listening to the model tell us where it is losing coherence. Easy scenes → model stays confident longer → MORE tokens (100–120). Hard scenes → confusion rises early → FEWER tokens (40–60). Scene length adapts to the model’s actual capability at that moment.

Benchmarking the Innovation

Main Results

Ablation Study — What Matters Most

Future Applications and Research Directions

Immediate: Academic Publication

Short paper (6–8 pages) presenting entropy-based adaptive scene cutting. Comprehensive ablation, human evaluation, and GPT-4 judge across 100 prompts. Title candidates: “Know When to Stop: Entropy-Based Scene Cutting for Coherent Story Generation with Small LMs” or “Scene Cards: Adaptive Iterative Generation for Sub-500M Language Models.”

Medium-Term: Edge AI & Multi-Language

164.6M parameters fits on mobile devices. Combined with adaptive cutting, this enables on-device story generation for educational apps without cloud APIs. The entropy mechanism is language-agnostic (measures confidence, not language features), so the framework transfers directly to other languages — especially important for low-resource settings.

Long-Term: Entropy as a General Coherence Signal

Beyond stories: paragraph boundaries in essays? Section transitions in documentation? Turn boundaries in dialogue? All unexplored applications of the same principle. Combined with knowledge distillation, RAG, and RLHF — each combination is a potential follow-up publication.

Academic References