A first-principles reproduction of Multi-Head Latent Attention, Mixture-of-Experts routing, and Multi-Token Prediction — the three mechanisms that make DeepSeek-V3 different from a standard transformer.
DeepSeek-V3's technical report describes a handful of architectural changes that let it beat models many times its active size. This repo implements all three of its headline mechanisms — Multi-Head Latent Attention, Mixture-of-Experts routing, and Multi-Token Prediction — entirely from scratch, no shortcuts, no pre-built attention/MoE libraries. Every projection, every routing decision, every auxiliary loss term is something you can step through and understand.
A 109M-parameter model built on this architecture was trained end-to-end on 2.5B tokens of FineWeb-Edu, converging cleanly over 20,000 steps for a total compute cost of $9.53 on a single A100.
The full block: MLA replaces standard attention, MoE replaces the dense FFN, and MTP hangs an auxiliary head off the trunk — repeated 8 times.
Quick orientation:
flowchart LR
A[Token embeddings] --> B["8× Transformer blocks<br/>(MLA + MoE FFN)"]
B --> C[Next-token head]
B --> D[MTP head<br/>predicts t+2]
C --> E[Loss: LM]
D --> F[Loss: MTP]
E --> G((Combined loss))
F --> G
- Why these three mechanisms
- Multi-Head Latent Attention
- Mixture of Experts
- Multi-Token Prediction
- The rest of the stack
- Model configuration
- Dataset
- Training
- Project layout
- Getting it running
- Results
- Reproducibility
- Roadmap
- References & citation
Standard transformers scale in two blunt ways: add more layers, or make each layer wider. DeepSeek-V3 instead asks three narrower questions, and each mechanism in this repo answers one of them:
- Attention's memory cost grows with every head you add — MLA shrinks that cost without losing what multiple heads give you.
- A feed-forward block spends the same compute on every token, easy or hard — MoE lets capacity grow faster than compute.
- Next-token prediction gives one gradient signal per position — MTP pulls more signal out of the same forward pass.
Each is implemented here with the actual math and a diagram of the data flow, not just a description of what it does.
Ordinary multi-head attention caches a full key and value tensor for every head, at every layer, for every token generated so far — and that cache, not the model weights, is usually what runs out of memory first during long generations.
MLA's answer: cache a compressed latent instead of the full K/V, and reconstruct full-size heads only at the instant they're needed.
The down-projection / cache / up-projection pipeline that replaces standard multi-head attention.
flowchart TB
subgraph "Standard MHA — cache grows with every head"
h1[hidden state] --> k1[K, 8 heads]
h1 --> v1[V, 8 heads]
k1 & v1 --> cache1[("KV cache<br/>O(n_heads × d_head)")]
end
flowchart TB
subgraph "MLA — cache one shared latent instead"
h2[hidden state] --> down["down-projection<br/>W_DKV (rank 128)"]
down --> latent[("c_KV — the only<br/>thing we cache")]
latent --> up1["up-project → K"]
latent --> up2["up-project → V"]
up1 & up2 --> attn[attention output]
end
c_KV = W_DKV · h # compress into a rank-128 latent
k_C, v_C = W_UK · c_KV, W_UV · c_KV # reconstruct per-head K/V when needed
c_Q = W_DQ · h # queries get their own rank-192 compression
q_C = W_UQ · c_Q
RoPE's positional signal is carved out into a small separate sub-dimension outside the compressed latent, so position information survives the compression cleanly instead of getting squeezed by the bottleneck — a design that preserves the full expressiveness of multi-head attention at a fraction of the cache size.
| Value | |
|---|---|
| KV cache content | one shared latent vector per token |
| KV compression rank | 128 |
| Query compression rank | 192 |
| Code | models/attention.py |
The intuition here is closer to a hospital than a factory line: instead of running every token through one generalist network, a router looks at each token and sends it to two specialists out of a pool of eight.
Eight expert FFNs, a router, and top-2 activation per token.
flowchart LR
T[token] --> R{Router}
R -->|top-2 of 8| E2[Expert 2]
R -->|top-2 of 8| E5[Expert 5]
E1[Expert 1]:::idle
E3[Expert 3]:::idle
E4[Expert 4]:::idle
E6[Expert 6]:::idle
E7[Expert 7]:::idle
E8[Expert 8]:::idle
E2 --> Sum[weighted sum]
E5 --> Sum
Sum --> Out[output]
classDef idle fill:#eee,stroke:#bbb,color:#999
g(h) = softmax(W_router · h) # score every expert
top-k = the 2 highest-scoring experts
y = Σ g_i(h) · Expert_i(h) # only over the chosen 2
The payoff is decoupling capacity from compute: the model owns 8 experts' worth of parameters — roughly 4x a same-cost dense FFN — while only paying for 2 of them on any given token, about 25% active compute. A load-balancing term keeps traffic spread across experts during training rather than letting a handful dominate.
| Value | |
|---|---|
| Experts | 8 |
| Active per token | 2 |
| Active compute | ~25% of dense |
| Effective capacity | ~4x |
| Code | models/moe.py |
Normal training gives the model one thing to be right about per position: the very next token. MTP adds a second task on top — guess the token after that one, too — using a small module chained off the main trunk.
The auxiliary MTP module chained off the main trunk, predicting one extra token ahead.
flowchart LR
X["input up to position t"] --> Trunk
Trunk --> H[hidden state]
H --> Head1[LM head]
Head1 --> P1["predict t+1"]
H --> MTPBlock["MTP block<br/>(+ true embedding of t+1)"]
MTPBlock --> Head2[MTP head]
Head2 --> P2["predict t+2"]
P1 --> L1[loss]
P2 --> L2[aux loss]
L_total = L_main + λ · L_mtp
This densifies the training signal per forward pass and sets up exactly the scaffolding speculative decoding needs — draft a token cheaply, verify it against the real path — a natural next step for faster inference.
| Value | |
|---|---|
| Extra prediction depth | 1 token |
| Code | models/mtp.py |
Three additional design choices round out the architecture:
- RoPE rotates queries and keys by an angle tied to position, so their dot product naturally encodes relative distance rather than absolute position — this generalizes to sequence lengths beyond what the model saw during training.
- RMSNorm normalizes by root-mean-square only, skipping LayerNorm's mean-centering step. Cheaper to compute, just as stable in practice.
- SwiGLU is the gated activation used inside every FFN and every expert:
(Swish(xW) ⊙ xV) W_out. It consistently edges out plain ReLU/GELU feed-forwards at the same parameter budget.
Full parameter breakdown for the 109M-parameter configuration used in this run.
Core
| Parameter | Value |
|---|---|
| Total parameters | 109,032,032 |
| Vocabulary | 50,257 (GPT-2 BPE) |
| Context length | 1,024 |
| Embedding dim | 512 |
| Layers | 8 |
| Attention heads | 8 |
| Batch size | 32 (× 8 grad-accum = 256 effective) |
| Peak / min LR | 3e-4 → 1e-5, cosine |
| Warmup | 2,000 steps |
| Max steps | 20,000 |
| Dropout | 0.1 |
MoE
| Parameter | Value |
|---|---|
| Experts | 8 |
| Active per token | 2 |
| Active compute | ~25% |
| Capacity multiplier | ~4x |
Attention
| Parameter | Value |
|---|---|
| KV latent rank | 128 |
| Query latent rank | 192 |
| MTP heads | 1 |
The main corpus is FineWeb-Edu (CC-MAIN-2024 subset — see the write-up for how it's filtered). Out of roughly 13M available records, 2M were sampled for this run, giving about 2.5B training tokens and 132.8M validation tokens.
TinyStories was used earlier in development to validate the MLA/MoE/MTP implementations quickly before committing to the full run — its narrow vocabulary makes correctness easy to check.
- Loss: next-token cross-entropy plus the MTP auxiliary loss, summed and backpropagated together.
- Optimizer: AdamW, cosine decay from 3e-4 to 1e-5, 2,000-step linear warmup.
- Batching: 32 sequences per step, 8 gradient-accumulation steps, 256 effective batch size.
- Logging: Weights & Biases — loss curves, LR schedule, and expert-utilization stats.
- Hardware: one A100 (80GB). The full 20,000-step run completed in 7h 1m for $9.53.
DeepSeek-from-Scratch/
├── models/
│ ├── attention.py # Multi-Head Latent Attention
│ ├── config.py # Model configuration
│ ├── layers.py # RoPE, RMSNorm, SwiGLU
│ ├── model.py # Transformer block assembly
│ ├── moe.py # Mixture of Experts
│ └── mtp.py # Multi-Token Prediction
├── training/
│ ├── data_loader.py
│ └── trainer.py
├── inference/
│ ├── generator.py
│ └── run_inference.py
├── notebooks/
│ ├── Mixture_of_Experts_from_Scratch.ipynb
│ ├── Multi_Head_Latent_Attention_From_Scratch.ipynb
│ └── Multi_Token_Prediction_from_Scratch.ipynb
├── prepare_data_fineweb.py
├── prepare_data_tiny_stories.py
└── main.py
Each notebook under notebooks/ builds one mechanism from raw tensors up, with shape checks along the way — a clean entry point before the production code in models/.
git clone https://github.com/username/DeepSeek-from-Scratch.git
cd DeepSeek-from-Scratch
python -m venv deepseek_env
source deepseek_env/bin/activate # Windows: deepseek_env\Scripts\activate
pip install -r requirements.txt# pick one
python prepare_data_fineweb.py
python prepare_data_tiny_stories.py
# train
python main.py train # config lives in models/config.py
# generate
python inference/run_inference.py # edit the prompt list first| Metric | Value |
|---|---|
| Steps completed | 20,000 |
| Wall-clock | 7h 1m on one A100 |
| Cost | $9.53 |
| LR schedule behavior | smooth cosine decay throughout, no instability |
| KV cache | reduced vs. an equal-head dense baseline, via MLA |
| Active FFN compute | ~25% of dense, via MoE |
Full training was completed at a fraction of the cost typically associated with transformer pretraining, demonstrating that architecturally sophisticated models — MLA, MoE, and MTP together — are well within reach on a single consumer-accessible GPU.
- Every hyperparameter above matches
models/config.pyexactly — the README and the run config are kept in sync. - The 2M-of-13M FineWeb-Edu sample is deterministic given the same seed and script version.
- Trained weights are published on Hugging Face, so results can be verified against this exact checkpoint.
- Environment: single A100 (80GB), PyTorch 2.x, Python 3.10+, exact versions pinned in
requirements.txt.
- Standardized benchmarking (MMLU, HellaSwag, ARC) to give this a clear external reference point
- Ablations isolating MLA / MoE / MTP's individual contributions
- A small family of model sizes to trace how these mechanisms scale
- Longer context windows, where MLA's memory advantage compounds further
- Speculative decoding, using the trained MTP head to accelerate inference
- Published expert-utilization histograms from the router
- DeepSeek-AI, DeepSeek-V3 Technical Report, arXiv:2412.19437, 2024 — link
- Fedus, Zoph, Shazeer, Switch Transformers, JMLR 2022 — the sparse-MoE lineage this builds on
- Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding, 2021
- Zhang & Sennrich, Root Mean Square Layer Normalization, NeurIPS 2019
- Shazeer, GLU Variants Improve Transformer, 2020
- Penedo et al., The FineWeb Datasets, Hugging Face 2024 — link
- Eldan & Li, TinyStories: How Small Can Language Models Be and Still Speak Coherent English?, 2023
@article{deepseekv3_2024,
title = {DeepSeek-V3 Technical Report},
journal = {arXiv preprint arXiv:2412.19437},
year = {2024},
url = {https://arxiv.org/pdf/2412.19437}
}Independent research-reproduction project. Weights, notebooks, and training logs are all linked above.