Designing a Training System
One complete worked example — training a custom document-classification encoder from scratch — walked through every layer of the mechanics, with concrete numbers
Designing a Training System
TL;DR
This page is one worked design, not a generic checklist: training a small encoder-only transformer from scratch to classify a company's internal documents into a fixed set of categories, on a realistically-sized internal labeled dataset. Every architectural and training decision below follows from that specific task and that specific data budget.
| Property | Value |
|---|---|
| Level | Advanced |
| Reading time | ~20 minutes |
| Prerequisites | Failure Modes & Debugging |
| You will understand | How the whole track's decisions fit together in one real, concrete design |
The Requirements
A company needs to route incoming internal documents (contracts, reports, correspondence) into one of 12 categories, to feed a downstream workflow tool. They have 40,000 labeled examples collected from the existing manual process, want fast (sub-100ms) inference on CPU-only serving infrastructure, and need the model retrainable in-house as new document types appear.
| Requirement | Implication |
|---|---|
| Classification, not generation | An encoder-only architecture, not decoder-only — no need to generate text |
| 40,000 labeled examples | A meaningful but not huge dataset — architecture and regularization should match this scale, not assume web-scale data |
| Fast CPU inference | A small model, favoring parameter efficiency over raw capacity |
| Retrainable in-house | The full pipeline (not just the model) needs to be simple enough to run without a dedicated ML infra team |
Architecture Choice
Encoder-only vs decoder-only, for this task
Encoder-only transformer
RecommendedAttends to the whole input bidirectionally in one pass and produces a single representation well-suited to classification. No need to generate tokens one at a time, so no autoregressive masking and no generation-time latency concerns.
Decoder-only transformer
Built for generating text one token at a time — using one to classify means either awkwardly prompting it for a label or discarding most of what makes it decoder-shaped. More parameters and more inference latency for no benefit on a pure classification task.
A small encoder — a handful of transformer blocks (see The Transformer Architecture), a modest hidden dimension, a classification head on top of a pooled representation — sized around a few million parameters, comfortably fast enough for the CPU latency target.
Initialization and Normalization
Standard practice, chosen deliberately rather than defaulted into: layer normalization after each sub-block (stabilizes activation scale through depth, per Regularization & Generalization), and a scaled initialization scheme appropriate to the depth of the network — both chosen specifically because 40,000 examples is not a large enough dataset to reliably train its way out of a bad initialization or unstable early training.
Optimizer and Learning Rate Schedule
AdamW, with a linear warmup over the first few hundred steps followed by cosine decay to a small fraction of the peak learning rate.
Why warmup: early in training, gradients from a randomly-initialized model are
noisy and can be large — jumping straight to the target learning rate risks an
unstable first few steps (per the NaN-loss triage in Evaluating & Debugging Training).
Why cosine decay: a smoothly decreasing learning rate lets the model make large,
fast progress early and settle into a more precise minimum later, rather than
overshooting a good solution late in training with a learning rate that's still high.Regularization Strategy for This Data Size
40,000 examples across 12 classes is enough to learn real patterns, but small enough that overfitting is a genuine risk for a model with millions of parameters. Dropout (a moderate rate, e.g. 0.1) on the transformer blocks and the classification head, weight decay through AdamW's decoupled implementation, and early stopping based on the validation curve (per Evaluating & Debugging Training) are all warranted here — this is not a dataset scale where regularization can be treated as optional.
Single-GPU vs Distributed
This scale doesn't need FSDP2, or even DDP. A model of a few million parameters, with an optimizer state that easily fits alongside it, trains comfortably on a single modest GPU — even a single reasonably capable CPU could finish training in a workable timeframe given the dataset size. Reaching for distributed training here, per Distributed Training Internals, would add real communication overhead for a problem that was never memory- or compute-bound in the first place.
Evaluation Plan
A held-out validation split (a genuine split of the 40,000 examples, stratified by category so rare categories aren't accidentally underrepresented in validation), tracked every epoch alongside training loss — watching specifically for the train/validation gap widening as the early-stopping signal. Per-category accuracy and confusion-matrix inspection matter more than a single overall accuracy number here, since a model that's excellent on the 3 common categories and poor on the 9 rare ones would look deceptively good on an aggregate score alone.
Checkpointing and Reproducibility Plan
Per Production & Operations: checkpoint model and optimizer state together every epoch, log hyperparameters and metrics from the very first training run (this model will be retrained repeatedly in-house as new document types appear, so comparing a future run against this baseline matters), and pin the exact PyTorch/CUDA versions used, since this pipeline is meant to be re-run by the company's own team later, not just once.
Concept Checks
Check yourself
Why is an encoder-only architecture the right choice here instead of a decoder-only one, given decoder-only models are more commonly discussed?
Because the task is classification — producing one label from a fixed set — not text generation, and an encoder attends to the whole input bidirectionally in one pass to produce a representation well suited to that. A decoder-only model is built around autoregressive next-token generation, which adds architectural complexity and inference latency this task has no use for.
Why does this design treat regularization as necessary rather than optional, given 40,000 labeled examples isn't a tiny dataset?
Because the risk of overfitting is about the ratio between data size and model capacity, not data size alone — a model with millions of parameters can still memorize meaningful chunks of 40,000 examples without regularization, especially per rare category. Dropout, weight decay, and early stopping are cheap insurance against that risk at this scale, not a sign the dataset is considered small in an absolute sense.
Why does this design explicitly reject distributed training, when Distributed Training Internals covers FSDP2 in depth?
Because FSDP2 and DDP solve problems — a model too large for one GPU's memory, or wanting more throughput than one GPU provides — that don't apply here. A few-million-parameter model with its optimizer state fits easily on a single modest GPU, so adding distributed training would only add communication overhead without solving any actual constraint this system has.
Key Concepts Recap
| Concept | One-line summary |
|---|---|
| Requirements drove every decision | Classification task, moderate data size, CPU latency target, in-house retraining |
| Encoder-only, not decoder-only | Matches the task — no generation needed |
| AdamW with warmup + cosine decay | Stable early training, precise late-training convergence |
| Regularization treated as necessary | Dropout, weight decay, early stopping — warranted at this data-to-capacity ratio |
| No distributed training needed | The model was never memory- or compute-bound at this scale |
| Per-category evaluation | An aggregate accuracy number can hide poor performance on rare categories |
| Checkpointing for future retraining | This pipeline is meant to be re-run, not used once |
Next
You've seen every decision in this track applied to one real system. The last page collects every term used along the way: Glossary.