0.04B Model Architecture Design, Mathematical Modeling, and Intuition

Contents19 sections

“When constructing a skyscraper, the first step is never laying bricks on a dirt lot, but drafting an exact architectural blueprint—calculating foundation load tolerances, steel grades, and construction budgets. In large language models (LLMs), that blueprint is the Model Specification & Mathematical Architecture.”

For newcomers entering the field of large language models, dense jargon (RoPE, GQA, SwiGLU, PPL, AMP) and multi-billion-parameter figures often construct an intimidating cognitive barrier.

As the foundational bedrock of this curriculum, this guide dissects a 0.04B (35.93M parameter) modern language model from the inside out. We provide not only relatable real-world analogies, but also complete mathematical derivations and memory ledgers so you can build uncompromising engineering confidence before writing a single line of PyTorch code.

0.1 Foundational Intuition: Understanding LLMs from First Principles

0.1.1 What Is a Large Language Model? From Dictionaries to Autoregression

To silicon hardware, the entirety of human literature is merely a sequence of discrete numerical tokens:

\mathbf{x} = (x_1, x_2, x_3, \dots, x_T)

Modern Decoder-Only Transformer architectures focus mathematically on a single objective: Conditional Probability Modeling:

P(\mathbf{x}) = \prod_{t=1}^T P(x_t \mid x_1, x_2, \dots, x_{t-1})

In plain language: “Given all preceding words in the sequence, predict what word is most likely to come next.”

If the preceding context is ["Ahead", "of", "my", "bed"], the model calculates the probability of the next word being "bright" at 95%, "moonlight" at 2%, and so on.
This step-by-step sequential continuation is called Autoregression.

0.1.2 Understanding the Physical Nature of "Parameters"

Beginners frequently wonder: "What are model parameters in physical reality, and where do they reside?"

Physical Reality: In source code and GPU VRAM, parameters are ordinary floating-point tensor matrices (such as torch.float32 or torch.float16).
Intuitive Metaphor: Picture sitting before a gigantic master mixing console equipped with 36 million tunable knobs (0.04B = $3.6 \times 10^7$ parameters).
Before Training (Random Initialization): Every knob is twisted arbitrarily according to a normal distribution. Play an input note, and the monitors emit deafening white noise (unintelligible token gibberish).
During Training (Backpropagation & Gradient Updates): We feed the model tens of thousands of master classical poems (460k entries). Calculus chain rules compute prediction errors, gently nudging all 36 million knobs toward harmonious intervals.
After Training (Convergence): When all 36 million knobs reach delicate equilibrium, the model naturally produces rhythmic, structured, and evocative stanzas in response to any prompt.

0.1.3 Why 0.04B (36M)? The Golden Sweet Spot for Domain Corpora

LLM parameter scales range across immense orders of magnitude:

Toy Miniatures (< 10M): Too few parameters restrict expressive capacity, forcing rote memorization of fixed phrases that collapse under slight prompt variations.
Hyperscale Industry Models (7B / 70B / 405B): Vast capabilities, but a single forward pass demands tens to hundreds of gigabytes of VRAM with multi-million-dollar training budgets—impractical for local end-to-end learning.
Why Is 0.04B (36M) the Ideal Sweet Spot for Classical Poetry?
Chinchilla Scaling Laws and Optimal Saturation:
Pioneering empirical research (the Chinchilla Law) establishes that optimal learning requires a balanced ratio between parameter volume and training tokens.
Our purified classical poetry corpus comprises 32.76 million tokens (31,126,208 train + 1,638,222 validation).
Scaling naively to 0.1B (82M parameters) produces severe data starvation, inducing shortcut memorization and catastrophic overfitting.
In contrast, 0.04B (35.93M parameters) and 32.76M tokens establish a near 1 : 1 optimal saturation ratio, allowing the network to absorb 460k poems deeply without memorization traps.
Rapid Iteration and High Throughput:
On a single consumer GPU or Intel Arc accelerator, 0.04B delivers sustained throughput of ~6,500 to 7,500 tok/s, doubling 0.1B iteration speeds.
Running a full training epoch takes ~70 minutes, while a 500-step preview completes in just ~40 minutes.
Accessible VRAM Demands:
Peak training memory for forward, backward, and AdamW states sits comfortably at ~1.8GB VRAM, enabling full execution on thin laptops, integrated GPUs, and standard desktop cards.

0.2 Modern LLM Architecture Blueprint: Operator Topology and Tensor Flows

Contemporary open-weight architectures (LLaMA-3, Qwen-2.5, DeepSeek) have thoroughly superseded original Transformer Post-LN and absolute positional embeddings with modern operator topologies:

0.2.1 End-to-End Data Flow and Operator Graph

0.2.2 Tensor Shape and Computation Stage Flow

Stage / Operator	Operator Name	Input Tensor Shape	Output Tensor Shape	Core Computation & Physical Significance
1. Tokenization	Byte-Level BPE	Raw text string	`[B, T]` (integer IDs)	4,096 classical poetry vocabulary, 256 byte atoms, zero OOV
2. Vector Embedding	Token Embedding	`[B, T]`	`[B, T, 512]`	Maps discrete token IDs into 512-dimensional continuous latent space
3. Attention Pre-Norm	Pre-RMSNorm 1	`[B, T, 512]`	`[B, T, 512]`	Normalizes input variance to stabilize multi-layer gradient flow
4. Grouped Attention	GQA (8Q : 2KV) + RoPE	`[B, T, 512]`	`[B, T, 512]`	8 Query heads share 2 Key/Value heads; Rotary embeddings encode positions
5. Residual Addition 1	Residual Add 1	Dual `[B, T, 512]`	`[B, T, 512]`	$X = X + \text{Attention}(X)$ , preserving identity highway
6. FFN Pre-Norm	Pre-RMSNorm 2	`[B, T, 512]`	`[B, T, 512]`	Normalizes representation energy prior to non-linear expansion
7. Gated FFN	SwiGLU	`[B, T, 512]`	`[B, T, 512]`	Expands to 1,408 dims with dual-channel gating $\text{Swish}(W_{gate}x) \cdot W_{up}x$
8. Residual Addition 2	Residual Add 2	Dual `[B, T, 512]`	`[B, T, 512]`	$X = X + \text{SwiGLU}(X)$ , fusing non-linear memory representations
... Layer Stack	12x Transformer	`[B, T, 512]`	`[B, T, 512]`	Repeats attention and gating transformation across 12 stacked layers
9. Final Normalization	Final RMSNorm	`[B, T, 512]`	`[B, T, 512]`	Enforces numerical stability before vocabulary score projection
10. Logits Projection	LM Head (Weight-Tied)	`[B, T, 512]`	`[B, T, 4096]`	Reuses input embedding matrix to project hidden states into vocabulary logits
11. Decoding Sampler	Softmax / Constrained	`[B, T, 4096]`	Single Token ID	Samples next token via temperature, Top-P, and metric rhythm constraints

0.3 Hyperparameter Derivations: The Engineering Rationale

Many tutorials present hyperparameter tables without justification. A rigorous ML architect must understand: what physical constraints and mathematical trade-offs govern each value?

0.3.1 Core Hyperparameter Specification Table

Architecture Parameter	Symbol	0.04B Value	Industry Baseline	Engineering Rationale & Trade-off
Vocabulary Size	$V$	4,096	LLaMA-3 uses 128k	Small models must avoid oversized vocabularies to reserve parameter budget for hidden layers
Hidden Dimension	$d_{model}$	512	Lightweight standard	Balances subspace expressiveness while dividing cleanly into head dimensions ( $8 \times 64 = 512$ )
Layer Depth	$L$	12	Classic benchmark depth	Ensures sufficient non-linear depth and hierarchical abstraction levels
Query Heads	$n_q$	8	Per-head $d_k=64$	8 distinct multi-dimensional attention query subspaces
KV Heads	$n_{kv}$	2	GQA ratio 4	Slashes KV-Cache memory footprint and memory bandwidth by 75%
Head Dimension	$d_k$	64	$512 / 8 = 64$	Aligns perfectly with GPU Tensor Core matrix multiplication block sizes
FFN Dimension	$d_{ffn}$	1,408	$\approx \frac{8}{3} d_{model}$	Follows modern SwiGLU 2/3 parameter allocation scaling heuristics
Context Window	$T$	1,024	Extended to 8k~128k	Chinese poems span < 100 characters; 1,024 tokens easily accommodates multi-stanza sets
Weight Tying	Weight Tying	True	Small model standard	Input embeddings and output projection heads share parameters, saving 2.10M weights

0.3.2 Deep Dive: Rationales for Downscaling Vocabulary from 8192 to 4096

In frontier commercial models, LLaMA-3 uses 128k tokens and Qwen-2.5 uses 152k tokens. Why do we firmly mandate a 4,096 vocabulary for our 0.04B model?

Let us calculate the parameter overhead:

\text{Embedding Parameters} = \text{vocab\_size} \times d_{model}

With an 8192 Vocabulary:

\text{Params}_{emb} = 8,192 \times 512 \approx \mathbf{4.19\text{ M}}

In a 38M model, a single embedding matrix consumes 11.0% of the entire parameter budget! Over one-tenth of model capacity is squandered on memorizing token lookup tables instead of deep relational reasoning.

With an Optimized 4096 Vocabulary:

\text{Params}_{emb} = 4,096 \times 512 \approx \mathbf{2.10\text{ M}}

Embedding parameters drop to just 5.5%! The saved 2.1M parameter budget directly feeds the 12 Transformer hidden layers (enriching attention and SwiGLU capacity).

Empirical Character Coverage in Chinese Poetry:
Deduplicating the entire 460k poem corpus yields only 10,840 unique Chinese characters;
The top 4,000 frequent characters account for 98.81% of all token occurrences;
The remaining 1.2% rare characters are cleanly covered by Byte-BPE's 256 foundational byte atoms with zero OOV risk;
Next-token prediction narrows from 1-of-8192 to 1-of-4096, cutting entropy by 1 bit and halving Softmax memory overhead for sharper convergence!

0.4 Microscopic Accounting: Distribution of the 35.93M Parameters

Let us inspect every individual component matrix inside the 0.04B architecture with caliper precision:

Module Level	Component Matrix / Operator	Tensor Shape	Parameter Formula	Exact Parameter Count
1. Token Embedding	`token_embeddings`	`[V=4096, D=512]`	$4,096 \times 512$	2,097,152 (2.10 M)
2. Single Transformer Block	Attention RMSNorm Scale	`[D=512]`	$512$	512
	Query Projection $W_q$	`[D=512, n_q*d_k=512]`	$512 \times 512$	262,144
	Key Projection $W_k$ (GQA 2 heads)	`[D=512, n_kv*d_k=128]`	$512 \times 128$	65,536
	Value Projection $W_v$ (GQA 2 heads)	`[D=512, n_kv*d_k=128]`	$512 \times 128$	65,536
	Attention Output $W_o$	`[n_q*d_k=512, D=512]`	$512 \times 512$	262,144
	FFN RMSNorm Scale	`[D=512]`	$512$	512
	SwiGLU Gate Projection $W_{gate}$	`[D=512, d_ffn=1408]`	$512 \times 1408$	720,896
	SwiGLU Up Projection $W_{up}$	`[D=512, d_ffn=1408]`	$512 \times 1408$	720,896
	SwiGLU Down Projection $W_{down}$	`[d_ffn=1408, D=512]`	$1408 \times 512$	720,896
	Single Block Subtotal	-	-	2,819,072 (2.82 M)
3. 12-Layer Stack	12 Sequential Blocks	$12 \times \text{Single Block}$	$12 \times 2,819,072$	33,828,864 (33.83 M)
4. Final Normalization	`norm` (Final RMSNorm)	`[D=512]`	$512$	512
5. Output Head	`lm_head` (Tied Weights)	`[D=512, V=4096]`	Shared with Embedding	0 (Reused)
Total Architecture Parameters	Mini-LLaMA-0.04B Golden Architecture	-	Exact Sum	35,926,528 (35.93 M)

Physical Validation: Run in your terminal: python -c "from src.model import MiniLLaMAConfig, MiniLLaMAForCausalLM; print(sum(p.numel() for p in MiniLLaMAForCausalLM(MiniLLaMAConfig()).parameters()))", and the output displays precisely 35,926,528, exactly matching our manual calculation down to the last single parameter!

0.5 Memory and Compute Derivations: Calculating Training VRAM Footprints

Why do developers with 8GB graphics cards still encounter RuntimeError: CUDA out of memory? We divide training VRAM consumption into four distinct categories:

VRAM Category	Memory Footprint (35.9M Model)	Numerical Precision / Breakdown	Role & Optimization Advice
1. Static Weights	`71.8 MB`	`fp16` (2 bytes/param)	Model weights loaded into GPU memory for forward passes.
2. Backward Gradients	`71.8 MB`	`fp16` (2 bytes/param)	First derivatives calculated during backprop matching weight shapes.
3. AdamW Optimizer State	`431.1 MB`	`fp32` (12 bytes/param)<br>• Master fp32 weights (4B)<br>• First momentum (4B)<br>• Second momentum (4B)	The largest static memory consumer during training; essential for numerical precision.
4. Dynamic Activations	`~800 MB ~ 1.2 GB`	Scales with Batch Size & Sequence Length	Stores intermediate layer outputs for backward pass computation.
Total Steady-State VRAM	~1.5 GB ~ 2.0 GB	-	Runs smoothly on consumer laptops, integrated GPUs, and standard graphics cards.

0.5.1 Deep Dive: Why AdamW Requires 12 Bytes per Parameter

In Automatic Mixed Precision (AMP) training, forward and backward operations execute in 16-bit floats (2 bytes), but to prevent microscopic gradient updates from underflowing into zero, AdamW tracks states in full 32-bit floating point:

Master Weights: A full fp32 copy of model weights (4 bytes).
First Momentum $m_t$ : Exponential moving average of past gradients (4 bytes).
Second Momentum $v_t$ : Exponential moving average of squared gradients (4 bytes).

\text{VRAM}_{\text{AdamW}} = 35.93\text{M} \times (4 + 4 + 4) \text{ bytes} = 35.93 \times 12\text{ MB} \approx \mathbf{431.1\text{ MB}}

Notice that optimizer state alone consumes 6 times more memory than the fp16 model weights themselves!

0.5.2 Compute Derivation: FLOPs Requirements and Training Time

Every multiply-accumulate operation ( $a \times b + c$ ) in a neural network counts as 2 floating-point operations (2 FLOPs).

Forward Pass: Each parameter participates in approximately one multiply-accumulate operation:

\text{FLOPs}_{fwd} \approx 2 \times N_{\text{params}} \times T_{\text{tokens}}

Backward Pass: Backpropagation computes gradients for both activations and weight tensors, requiring double the forward compute:

\text{FLOPs}_{bwd} \approx 4 \times N_{\text{params}} \times T_{\text{tokens}}

Total Forward + Backward Compute per Token:

\text{FLOPs}_{total} \approx 6 \times N_{\text{params}} \approx 6 \times 3.59 \times 10^7 \approx \mathbf{2.16 \times 10^8 \text{ FLOPs/token}}

💡 Compute Budget Estimation: Training 3 Full Epochs on 100M Tokens

Total floating-point operations:

\text{Total Compute} = 2.16 \times 10^8 \times 10^8 \approx \mathbf{2.16 \times 10^{16} \text{ FLOPs}}

On a consumer GPU running at ~7,000 tok/s:

\text{Time} = \frac{10^8 \text{ tokens}}{7,000 \text{ tok/s}} \approx 14,285\text{ seconds} \approx \mathbf{3.96 \text{ hours}}

A quick 500-step preview run (16.38M tokens) finishes in just ~40 minutes!

0.6 Common Pitfalls and Conceptual FAQ

Q1: Why not stack layers deeper (e.g., 48 layers) with smaller hidden dimensions?

Answer: This embodies the classic trade-off between "deep & narrow" versus "shallow & wide".

Parallel Efficiency: Each Transformer block must wait for preceding layer outputs to complete (sequential execution). Excessive depth increases GPU core idle times and kernel invocation latencies.

Gradient Highway Friction: Despite Pre-LN and residual connections, signals in extremely deep networks suffer representation drift. For a 0.04B model, 12 layers with $d_{model}=512$ represents the empirically verified optimum.

Q2: Does pretraining actually instill commonsense knowledge, or just parrot text?

Answer: Autoregressive modeling appears to simply predict next characters, but accurately continuing complex passages forces the 12-layer attention network to construct structured concept representations in latent space. For example, predicting "frost" after "Ahead of bed bright moonlight" requires the model to align bed, moon, and Tang poetic imagery within high-dimensional space. This structured probabilistic modeling forms the genesis of language intelligence.

0.7 Chapter Summary and Next Steps

In this guide, we finalized the top-level architecture for our 0.04B golden model:

Specifications: $V=4096, d_{model}=512, L=12, d_{ffn}=1408, n_q=8, n_{kv}=2$ , locking in 35,926,528 parameters;
Hardware Profile: Steady-state VRAM requirements of 1.5GB ~ 2.0GB, training locally in 40 minutes to 3.5 hours;
Architecture Strategy: Fully aligned with modern open-weight architectures (RMSNorm + RoPE + SwiGLU + GQA).

The blueprints are drawn and tools prepared! In the next guide, we dive into raw byte streams and build our own tokenizer from scratch: Chapter 01 ｜ Implementing a Byte-Level BPE Tokenizer from Scratch.

REFERENCES

References

Series

Building an LLM from scratch