# Annotated Deep Learning Paper Implementations — Administrator & Developer Guide

For usage-facing content (installation, import paths, running experiments, troubleshooting), see the companion User Guide.

---

## 1. Architecture Overview

### What This Project Does

`labml_nn` is an open-source Python package that provides annotated, educational PyTorch implementations of 60+ deep learning research papers. It is structured as both a pip-installable library (package name `labml-nn`, import name `labml_nn`) and a literate-programming documentation site rendered at [nn.labml.ai](https://nn.labml.ai).

The repository serves a dual purpose:

1. **Library:** Every implementation is a standard `nn.Module` that can be imported and composed into larger models
2. **Documentation:** Every Python file contains extensive docstrings with LaTeX math and explanatory prose. A build pipeline converts these annotated sources into side-by-side HTML (code + explanation) published to [nn.labml.ai](https://nn.labml.ai)

The project depends on the `labml` ecosystem (a separate set of packages by the same authors) for experiment tracking, configuration management, and training loop infrastructure.

### Component Architecture

```text
labml_nn (pip package)
├── transformers/          22+ transformer variants (MHA, XL, RoPE, ViT, etc.)
├── diffusion/             DDPM, DDIM, Stable Diffusion, Latent Diffusion
├── gan/                   Original, DCGAN, CycleGAN, WGAN, WGAN-GP, StyleGAN2
├── optimizers/            Adam, AMSGrad, RAdam, AdaBelief, Sophia, Noam, AdamWarmup
├── normalization/         BatchNorm, LayerNorm, InstanceNorm, GroupNorm, WeightStd, DeepNorm
├── sampling/              Greedy, Temperature, Top-k, Nucleus
├── rl/                    PPO (with GAE), DQN (Dueling + Prioritized Replay)
├── neox/                  GPT-NeoX 20B inference, fine-tuning, int8 quantization
├── lora/                  Low-Rank Adaptation with GPT-2
├── rwkv/                  RWKV recurrence-based language model
├── graphs/                GAT, GATv2
├── cfr/                   Counterfactual Regret Minimization (Kuhn Poker)
├── resnet/                ResNet with bottleneck blocks
├── conv_mixer/            ConvMixer (patch-based convolution)
├── unet/                  U-Net segmentation
├── capsule_networks/      Capsule Networks with dynamic routing
├── lstm/                  LSTM cell and multi-layer
├── recurrent_highway_networks/  RHN cell
├── hypernetworks/         HyperLSTM
├── sketch_rnn/            Sketch RNN (stroke generation)
├── distillation/          Knowledge distillation
├── adaptive_computation/  PonderNet
├── uncertainty/           Evidential deep learning
├── activations/           Fuzzy Tiling Activations, Swish
├── scaling/               Zero3 memory optimization
├── helpers/               Datasets, training loops, metrics, schedules
├── experiments/           Shared experiment configs (CIFAR-10, MNIST, NLP autoregression)
└── utils/                 Tokenizer configs, map-style datasets
```

> [!NOTE]
> The `labml_nn` package has a flat hierarchy where each research area is a top-level subpackage. There is no shared base class beyond PyTorch's `nn.Module` — implementations are intentionally independent so readers can study one paper without understanding others.

### Technology Stack

| Layer | Technology |
|---|---|
| Language | Python 3.7+ |
| Deep learning framework | PyTorch ≥1.10 |
| Experiment tracking | labml (custom, same authors) |
| Configuration | labml.configs (declarative BaseConfigs) |
| Training infrastructure | labml-helpers (training loops, data loaders, metrics) |
| Tensor operations | einops, numpy |
| Distributed training | fairscale (pipeline parallelism) |
| Documentation rendering | Custom literate-programming pipeline → HTML |
| Package distribution | PyPI (`labml-nn` v0.5.1) |
| License | MIT |

### Data Flow

The typical workflow for using `labml_nn`:

```text
1. User imports nn.Module subclass from labml_nn.*
2. Module is instantiated with paper-specific hyperparameters
3. Module is composed into a larger model or used standalone
4. Experiment configs (BaseConfigs subclasses) wire together:
   - Model architecture
   - Dataset loaders (MNIST, CIFAR-10, text, etc.)
   - Optimizer selection
   - Training loop (via TrainValidConfigs)
5. labml tracks metrics, saves checkpoints, provides dashboard
```

---

## 2. System Requirements

| Component | Minimum | Recommended |
|---|---|---|
| Python | 3.7 | 3.10+ |
| PyTorch | 1.10 | Latest stable |
| CUDA | Not required | 11.x+ for GPU training |
| RAM | 4 GB | 16+ GB |
| GPU VRAM | — | 8 GB (most experiments), 48 GB (GPT-NeoX) |
| Disk | 200 MB (package) | 50+ GB (GPT-NeoX checkpoints) |

### Required Dependencies

From `setup.py`:

```python
install_requires=[
    'labml',
    'torch',
    'torchtext',
    'torchvision',
    'einops',
    'numpy',
    'fairscale',
]
```

From `requirements.txt` (pinned versions for development):

```text
torch>=1.10
torchvision>=0.11
torchtext>=0.11
labml>=0.4.147
labml-helpers>=0.4.84
numpy>=1.19
matplotlib>=3.0.3
einops>=0.3.0
gym[atari]
opencv-python
Pillow>=6.2.1
faiss
```

> [!IMPORTANT]
> `setup.py` and `requirements.txt` diverge: `setup.py` omits `labml-helpers`, `matplotlib`, `gym`, `opencv-python`, `Pillow`, and `faiss`. These are needed for specific experiments but not for core library usage. When deploying for full experiment reproduction, use `requirements.txt`.

---

## 3. Configuration Reference

### The labml Configuration System

`labml_nn` uses the `labml.configs` system for experiment configuration. This is a declarative system where:

- Configuration classes inherit from `BaseConfigs`
- Parameters are declared as class attributes with type annotations and optional defaults
- Alternative implementations are registered via `@option` decorators
- Configurations are composed via `aggregate()` to define named presets

**Pattern:**

```python
from labml.configs import BaseConfigs, option, aggregate

class MyConfigs(BaseConfigs):
    d_model: int = 512
    dropout: float = 0.1
    model: nn.Module  # computed from other configs

@option(MyConfigs.model, 'transformer')
def _model_transformer(c: MyConfigs):
    return TransformerModel(d_model=c.d_model, dropout=c.dropout)
```

### Key Configuration Classes

| Class | Location | Purpose |
|---|---|---|
| `TransformerConfigs` | `labml_nn/transformers/configs.py` | Transformer model assembly |
| `FeedForwardConfigs` | `labml_nn/transformers/configs.py` | FFN layer with GLU variant support |
| `OptimizerConfigs` | `labml_nn/optimizers/configs.py` | Optimizer selection and hyperparameters |
| `OptimizerConfigs` | `labml_nn/helpers/optimizer.py` | Legacy optimizer config |
| `DeviceConfigs` | `labml_nn/helpers/device.py` | CPU/GPU device selection |
| `MNISTConfigs` | `labml_nn/helpers/datasets.py` | MNIST data loading |
| `CIFAR10Configs` | `labml_nn/helpers/datasets.py` | CIFAR-10 data loading |
| `TrainValidConfigs` | `labml_nn/helpers/trainer.py` | Training loop orchestration |
| `TrainingLoopConfigs` | `labml_nn/helpers/trainer.py` | Low-level loop control |
| `NLPAutoRegressionConfigs` | `labml_nn/experiments/nlp_autoregression.py` | Autoregressive text experiments |
| `NLPClassificationConfigs` | `labml_nn/experiments/nlp_classification.py` | Text classification experiments |
| `CFRConfigs` | `labml_nn/cfr/__init__.py` | CFR game solver |
| `RWKVConfigs` | `labml_nn/rwkv/configs.py` | RWKV model configuration |
| `PipelineParallelTrainerConf` | `labml_nn/neox/utils/trainer.py` | Multi-GPU pipeline training |

### TransformerConfigs Reference

| Parameter | Type | Default | Description |
|---|---|---|---|
| `d_model` | `int` | — | Embedding dimension |
| `heads` | `int` | — | Attention heads |
| `d_ff` | `int` | 2048 | FFN hidden size |
| `dropout` | `float` | 0.1 | Dropout rate |
| `n_layers` | `int` | — | Number of layers |
| `n_src_vocab` | `int` | — | Source vocabulary size |
| `n_tgt_vocab` | `int` | — | Target vocabulary size |

### FeedForwardConfigs Reference

| Parameter | Type | Default | Description |
|---|---|---|---|
| `d_model` | `int` | — | Input features |
| `d_ff` | `int` | 2048 | Hidden layer size |
| `dropout` | `float` | 0.1 | Dropout rate |
| `activation` | `nn.Module` | ReLU | Activation (ReLU or GELU via @option) |
| `is_gated` | `bool` | `False` | Enable gated FFN |
| `bias1` | `bool` | `True` | First layer bias |
| `bias2` | `bool` | `True` | Second layer bias |
| `bias_gate` | `bool` | `False` | Gate layer bias |
| `glu_variant` | `str` | `'none'` | GLU preset |

### MNISTConfigs Reference

| Parameter | Type | Default | Description |
|---|---|---|---|
| `dataset_name` | `str` | `'MNIST'` | Dataset identifier |
| `train_batch_size` | `int` | 64 | Training batch size |
| `valid_batch_size` | `int` | 1024 | Validation batch size |
| `train_loader_shuffle` | `bool` | `True` | Shuffle training data |
| `valid_loader_shuffle` | `bool` | `False` | Shuffle validation data |

### CIFAR10Configs Reference

| Parameter | Type | Default | Description |
|---|---|---|---|
| `dataset_name` | `str` | `'CIFAR10'` | Dataset identifier |
| `train_batch_size` | `int` | 64 | Training batch size |
| `valid_batch_size` | `int` | 1024 | Validation batch size |

---

## 4. Deployment & Scaling

### PyPI Package

The package is published to PyPI as `labml-nn` version 0.5.1.

**Package metadata from `setup.py`:**

| Field | Value |
|---|---|
| Package name | `labml_nn` |
| Version | `0.5.1` |
| Authors | Varuna Jayasiri, Nipun Wijerathne |
| License | MIT |
| Python classifiers | Python 3, Scientific/Engineering, AI |
| Excluded packages | `labml`, `labml.*`, `labml_samples`, `labml_helpers`, `test` |

### Documentation Website

The documentation site at [nn.labml.ai](https://nn.labml.ai) is generated from the annotated Python source:

1. Python files contain docstrings with markdown + LaTeX
2. A build pipeline (in `utils/`) processes these into HTML
3. HTML files live in the `docs/` directory (2000+ files across English + translations)
4. The `translate_cache/` directory contains cached translations (Chinese, Japanese, Sinhala)

**Site generation utilities:**

| Script | Location | Purpose |
|---|---|---|
| `diagrams.py` | `utils/diagrams.py` | Processes SVG diagrams, adds CSS classes |
| `papers_list.py` | `utils/papers_list.py` | Collects paper titles from module docstrings |
| `sitemap.py` | `utils/sitemap.py` | Generates sitemap.xml from docs/ tree |

> [!TIP]
> The `utils/diagrams.py` script recursively processes SVG files in the `docs/` directory, applying consistent styling classes for backgrounds, strokes, and fills. Run it after updating any diagram.

### Multi-GPU Training

The GPT-NeoX implementation includes pipeline-parallel training via `PipelineParallelTrainerConf`:

```python
from labml_nn.neox.utils.trainer import PipelineParallelTrainerConf

conf = PipelineParallelTrainerConf()
# Splits model layers across available GPUs
```

The `fairscale` dependency provides the underlying distributed primitives.

---

## 5. Operations

### Experiment Tracking with labml

All experiments use `labml` for tracking. The `labml` library:

- Records metrics via `tracker.save()`
- Manages experiment artifacts and checkpoints
- Provides a web dashboard for monitoring
- Stores data in `lab.get_data_path()` (configurable)

### Training Loop Architecture

The training infrastructure centers on `TrainValidConfigs` in `labml_nn/helpers/trainer.py`:

```text
TrainingLoopConfigs
  └── TrainValidConfigs
        ├── init()     — create model, optimizer, data loaders
        ├── step()     — single train/valid step (abstract)
        ├── run_step() — handles forward, loss, backward, optimizer step
        ├── run()      — main training loop
        └── sample()   — generate samples (optional)
```

> [!NOTE]
> The `TrainingLoop` class in `labml_nn/helpers/trainer.py` supports graceful interruption: catching `SIGINT` (Ctrl+C) raises `StopIteration('SIGINT')`, allowing the training loop to save state before exiting.

> [!TIP]
> To add a new experiment, subclass `TrainValidConfigs` (or one of its domain-specific children like `NLPAutoRegressionConfigs`) and implement `init()` and `step()`. The training loop, metric tracking, and checkpoint management come for free.

### GPT-NeoX Checkpoint Management

GPT-NeoX model weights are downloaded from `https://mystic.the-eye.eu/public/AI/models/GPT-NeoX-20B/slim_weights/` via `labml_nn/neox/checkpoint.py`:

- Checkpoints are stored at `lab.get_data_path() / 'neox' / 'slim_weights'` (or `neox_fast/`)
- The model has 44 transformer layers by default
- Each layer's weights are split across 2 checkpoint files that get merged
- Merging handles dimension-0 concatenation (attention weights), dimension-1 concatenation (FFN), and duplicate parameters (layer norms)

### Fine-Tuning Infrastructure

The `FineTuner` class in `labml_nn/neox/utils/finetune.py` provides a parameter-freezing abstraction:

```python
class FineTuner:
    def get_trainable_params(self)     # parameters to train
    def set_trainable_params(self)     # freeze everything else
    def state_dict(self)               # save only trainable params
    def load_state_dict(self, state)   # restore trainable params
```

`FineTuneBiases` is the provided concrete implementation that trains only bias parameters.

---

## 6. Performance

### Memory Optimization

The `labml_nn.scaling.zero3` module implements Zero-3 style memory optimization for training large models. This partitions optimizer states, gradients, and parameters across data-parallel workers.

### LLM.int8() Quantization

For GPT-NeoX inference on consumer GPUs, `labml_nn.neox.utils.llm_int8` provides 8-bit quantization:

- Requires the `bitsandbytes` package
- Significantly reduces GPU memory requirements
- Trades slight accuracy for dramatically lower VRAM usage

### Adam FP16

The `labml_nn.optimizers.adam_fp16` module provides an FP16-aware Adam variant:

- `AdamFP16` maintains master weights in FP32 while training in FP16
- `GradScalerFP16` extends PyTorch's `GradScaler` for mixed-precision training

---

## 7. Security

### Supply Chain

The package depends on PyTorch ecosystem packages (torch, torchvision, torchtext), the labml ecosystem, and standard scientific Python (numpy, einops). All dependencies are sourced from PyPI.

### Checkpoint Downloads

The GPT-NeoX implementation downloads model weights from `https://mystic.the-eye.eu/public/AI/models/GPT-NeoX-20B/slim_weights/`. These are loaded via `torch.load()` with pickle deserialization.

> [!WARNING]
> `torch.load()` uses Python's `pickle` module by default, which can execute arbitrary code during deserialization. Only load checkpoints from trusted sources. The GPT-NeoX checkpoints are hosted by Eleuther AI.

### No Authentication or Network Services

The `labml_nn` package does not expose any network services, APIs, or authentication mechanisms. It is a pure computation library. The only network activity is:

- Downloading datasets (via torchvision/labml, on first use)
- Downloading GPT-NeoX checkpoints (on first use)
- labml experiment tracking (optional, to labml.ai dashboard)

---

## 8. Developer Internals

### Repository Structure

```text
annotated_deep_learning_paper_implementations/
├── setup.py                          Package definition (v0.5.1)
├── requirements.txt                  Full dev/experiment dependencies
├── readme.md                         GitHub README
├── labml_nn/                         Main package (all implementations)
│   ├── __init__.py                   Package docstring = implementation catalog
│   ├── transformers/                 Transformer implementations (24 variants)
│   │   ├── __init__.py               Module catalog + re-exports
│   │   ├── configs.py                TransformerConfigs, FeedForwardConfigs
│   │   ├── mha.py                    Multi-Head Attention
│   │   ├── models.py                 Encoder, Decoder, EncoderDecoder, Generator
│   │   ├── feed_forward.py           Position-wise FFN
│   │   ├── positional_encoding.py    Sinusoidal positional encoding
│   │   ├── label_smoothing_loss.py   Label smoothing cross-entropy
│   │   ├── xl/                       Transformer-XL + relative MHA
│   │   ├── rope/                     Rotary Positional Embeddings
│   │   ├── alibi/                    ALiBi attention bias
│   │   ├── retro/                    RETRO retrieval-augmented
│   │   ├── compressive/              Compressive Transformer
│   │   ├── gpt/                      GPT-2 architecture
│   │   ├── glu_variants/             GLU feed-forward variants
│   │   ├── knn/                      kNN-LM
│   │   ├── feedback/                 Feedback Transformer
│   │   ├── switch/                   Switch Transformer (MoE)
│   │   ├── fast_weights/             Fast Weights Transformer
│   │   ├── fnet/                     FNet (Fourier mixing)
│   │   ├── aft/                      Attention Free Transformer
│   │   ├── mlm/                      Masked Language Model
│   │   ├── mlp_mixer/                MLP-Mixer
│   │   ├── gmlp/                     gMLP
│   │   ├── vit/                      Vision Transformer
│   │   ├── primer_ez/                Primer EZ
│   │   ├── hour_glass/               Hourglass Transformer
│   │   ├── flash/                    Triton Flash Attention
│   │   └── jax_transformer/          JAX transformer implementation
│   ├── diffusion/                    Diffusion models
│   │   ├── ddpm/                     DDPM
│   │   └── stable_diffusion/         Stable Diffusion + DDIM + Latent Diffusion
│   │       ├── model/                UNet, autoencoder, CLIP
│   │       ├── sampler/              DDPM and DDIM samplers
│   │       └── scripts/              Generation scripts
│   ├── gan/                          GAN implementations
│   │   ├── original/                 Original GAN
│   │   ├── dcgan/                    DCGAN
│   │   ├── cycle_gan/                CycleGAN
│   │   ├── wasserstein/              WGAN + gradient penalty
│   │   └── stylegan/                 StyleGAN 2
│   ├── optimizers/                   Optimizer implementations
│   │   ├── __init__.py               GenericAdaptiveOptimizer base + WeightDecay
│   │   ├── adam.py                   Adam
│   │   ├── amsgrad.py                AMSGrad
│   │   ├── adam_warmup.py            Adam with warmup
│   │   ├── adam_warmup_cosine_decay.py  Adam warmup + cosine decay
│   │   ├── adam_fp16.py              Mixed-precision Adam
│   │   ├── noam.py                   Noam scheduler
│   │   ├── radam.py                  Rectified Adam
│   │   ├── ada_belief.py             AdaBelief
│   │   ├── sophia.py                 Sophia-G
│   │   ├── configs.py                OptimizerConfigs
│   │   └── mnist_experiment.py       MNIST training example
│   ├── normalization/                Normalization layers (7 types)
│   ├── sampling/                     Sampling techniques (4 types)
│   ├── rl/                           Reinforcement learning (PPO, DQN)
│   ├── neox/                         GPT-NeoX (model, checkpoint, evaluation)
│   ├── lora/                         LoRA fine-tuning
│   ├── rwkv/                         RWKV language model
│   ├── graphs/                       Graph neural networks (GAT, GATv2)
│   ├── cfr/                          Counterfactual Regret Minimization
│   ├── resnet/                       ResNet
│   ├── conv_mixer/                   ConvMixer
│   ├── unet/                         U-Net
│   ├── capsule_networks/             Capsule Networks
│   ├── lstm/                         LSTM
│   ├── recurrent_highway_networks/   Recurrent Highway Networks
│   ├── hypernetworks/                HyperLSTM
│   ├── sketch_rnn/                   Sketch RNN
│   ├── distillation/                 Knowledge distillation
│   ├── adaptive_computation/         PonderNet
│   ├── uncertainty/                  Evidential deep learning
│   ├── activations/                  Fuzzy Tiling Activations, Swish
│   ├── scaling/                      Zero3 memory optimization
│   ├── helpers/                      Shared infrastructure
│   │   ├── datasets.py               MNIST, CIFAR-10, text dataset loaders
│   │   ├── trainer.py                Training loops, TrainValidConfigs
│   │   ├── metrics.py                Accuracy tracking
│   │   ├── device.py                 Device management
│   │   ├── optimizer.py              Legacy optimizer configs
│   │   └── schedule.py               Learning rate schedules
│   ├── experiments/                  Experiment base configs
│   │   ├── cifar10.py                CIFAR-10 experiment configs
│   │   ├── mnist.py                  MNIST experiment configs
│   │   ├── nlp_autoregression.py     Text generation base config
│   │   ├── nlp_classification.py     Text classification base config
│   │   └── arithmetic_dataset.py     Arithmetic problem dataset
│   └── utils/                        Utilities
│       ├── __init__.py               MapStyleDataset
│       └── tokenizer.py              TokenizerConfigs
├── docs/                             Pre-rendered HTML documentation (2000+ files)
│   ├── *.html                        English documentation
│   ├── zh/                           Chinese translations
│   ├── ja/                           Japanese translations
│   └── si/                           Sinhala translations
├── translate_cache/                  Translation cache files
├── utils/                            Build utilities
│   ├── diagrams.py                   SVG diagram processor
│   ├── papers_list.py                Paper title collector
│   └── sitemap.py                    Sitemap generator
└── papers/                           Paper metadata
```

### Class Hierarchy

The optimizer hierarchy demonstrates the project's inheritance pattern:

```text
torch.optim.Optimizer
└── GenericAdaptiveOptimizer
    ├── Adam
    │   ├── AdamFP16
    │   └── AMSGrad
    │       ├── AdamWarmup
    │       ├── AdamWarmupCosineDecay
    │       ├── Noam
    │       └── RAdam
    │           └── AdaBelief
    └── Sophia
```

The GPT-NeoX module hierarchy:

```text
nn.Module
├── NeoXModule (custom load_state protocol)
│   ├── Embedding
│   ├── TransformerLayer
│   ├── FinalNorm
│   └── ReadoutLayer
├── RoPE
├── AttentionLayer
└── FFNLayer
```

The training infrastructure hierarchy:

```text
BaseConfigs
└── TrainingLoopConfigs
    └── TrainValidConfigs
        ├── SimpleTrainValidConfigs
        ├── NLPAutoRegressionConfigs
        │   └── ArithmeticAutoregression
        └── NLPClassificationConfigs
```

### Key Design Patterns

**1. Literate Programming Style**

Every Python file is simultaneously executable code and documentation. Docstrings use a custom format:

```python
"""
---
title: Multi-Head Attention
summary: >
  Implementation of multi-head attention mechanism.
---

# Multi-Head Attention

This implements multi-headed attention from
[Attention Is All You Need](https://arxiv.org/abs/1706.03762).
"""
```

The YAML frontmatter (`title`, `summary`) is consumed by the documentation renderer. The markdown body becomes the side-by-side explanation.

**2. Configuration via BaseConfigs**

Instead of command-line arguments or YAML files, experiments use Python class-based configuration:

```python
class Configs(CIFAR10Configs):
    model: nn.Module
    learning_rate: float = 0.001

@option(Configs.model, 'resnet')
def _model(c: Configs):
    return ResNetBase(...)
```

The `@option` decorator registers named alternatives. The `aggregate()` function creates preset bundles.

> [!NOTE]
> The configuration system uses Python's descriptor protocol under the hood. When you access `c.model`, it lazily evaluates the registered `@option` function. This means configuration objects are not plain data — they are evaluation graphs.

**3. Self-Contained Modules**

Each paper implementation is designed to be understood in isolation. There is minimal cross-referencing between implementations. The `__init__.py` of each subpackage acts as a catalog pointing to its children.

**4. Experiment + Module Separation**

Most subpackages separate the model (`__init__.py` or named module) from the experiment (`experiment.py`). The model defines `nn.Module` subclasses; the experiment wires them into a training pipeline via `Configs` subclasses.

### Module Reference

#### Core Transformer Modules (`labml_nn/transformers/`)

| Class | File | Methods | Purpose |
|---|---|---|---|
| `MultiHeadAttention` | `mha.py:69` | `__init__`, `get_scores`, `prepare_mask`, `forward` | Multi-head scaled dot-product attention |
| `PrepareForMultiHeadAttention` | `mha.py:33` | `__init__`, `forward` | Linear projection for Q/K/V |
| `FeedForward` | `feed_forward.py:43` | `__init__`, `forward` | Position-wise FFN with optional gating |
| `TransformerLayer` | `models.py:60` | `__init__`, `forward` | Single transformer layer (self-attn + optional cross-attn + FFN) |
| `Encoder` | `models.py:131` | `__init__`, `forward` | Stack of N transformer encoder layers |
| `Decoder` | `models.py:153` | `__init__`, `forward` | Stack of N transformer decoder layers |
| `EncoderDecoder` | `models.py:193` | `__init__`, `forward`, `encode`, `decode` | Full encoder-decoder transformer |
| `Generator` | `models.py:175` | `__init__`, `forward` | Linear + log-softmax output layer |
| `PositionalEncoding` | `positional_encoding.py:30` | `__init__`, `forward` | Sinusoidal positional encoding |
| `EmbeddingsWithPositionalEncoding` | `models.py:24` | `__init__`, `forward` | Token embedding + sinusoidal PE |
| `EmbeddingsWithLearnedPositionalEncoding` | `models.py:42` | `__init__`, `forward` | Token embedding + learned PE |
| `LabelSmoothingLoss` | `label_smoothing_loss.py:18` | `__init__`, `forward` | KL-divergence based label smoothing |
| `TransformerConfigs` | `configs.py:143` | — | Declarative transformer assembly |
| `FeedForwardConfigs` | `configs.py:20` | — | Declarative FFN assembly |

#### Optimizer Modules (`labml_nn/optimizers/`)

| Class | File | Parents | Key Methods |
|---|---|---|---|
| `GenericAdaptiveOptimizer` | `__init__.py:70` | `Optimizer` | `init_state`, `step_param`, `step` |
| `WeightDecay` | `__init__.py:167` | — | `__init__`, `defaults`, `__call__` |
| `Adam` | `adam.py:50` | `GenericAdaptiveOptimizer` | `init_state`, `get_mv`, `get_lr`, `adam_update`, `step_param` |
| `AMSGrad` | `amsgrad.py:27` | `Adam` | `init_state`, `get_mv` |
| `RAdam` | `radam.py:148` | `AMSGrad` | `step_param`, `calc_rectification_term`, `r_adam_update` |
| `AdaBelief` | `ada_belief.py:45` | `RAdam` | `init_state`, `get_ms`, `step_param` |
| `AdamWarmup` | `adam_warmup.py:18` | `AMSGrad` | `get_lr` |
| `Noam` | `noam.py:20` | `AMSGrad` | `get_lr` |
| `Sophia` | `sophia.py:62` | `GenericAdaptiveOptimizer` | `init_state`, `update_hessian`, `step_param` |
| `AdamFP16` | `adam_fp16.py:22` | `Adam` | `init_state`, `step_param` |

The optimizer base class validates hyperparameters at construction:

- `lr` must be ≥ 0
- `eps` must be ≥ 0
- `betas[0]` and `betas[1]` must be in [0, 1)
- `weight_decay` must be ≥ 0
- Sparse gradients are not supported (raises `RuntimeError`)

#### GPT-NeoX Modules (`labml_nn/neox/`)

| Class | File | Purpose |
|---|---|---|
| `NeoXModule` | `model.py:30` | Base class with `load_state` protocol |
| `Embedding` | `model.py:35` | Token embeddings |
| `RoPE` | `model.py:65` | Rotary positional encoding |
| `AttentionLayer` | `model.py:168` | Multi-head attention (flash + standard paths) |
| `FFNLayer` | `model.py:346` | Feed-forward network |
| `TransformerLayer` | `model.py:378` | Complete transformer block |
| `FinalNorm` | `model.py:447` | Final layer normalization |
| `ReadoutLayer` | `model.py:475` | LM head (vocabulary projection) |
| `LayerGenerator` | `model.py:503` | Lazy layer loading for memory efficiency |
| `Cache` | `utils/cache.py:18` | Key-value cache for autoregressive generation |
| `FineTuner` | `utils/finetune.py:9` | Parameter freezing abstraction |
| `FineTuneBiases` | `utils/finetune.py:43` | Bias-only fine-tuning |
| `NeoXDataset` | `utils/text_dataset.py:55` | Text dataset for NeoX training |
| `TrainerConf` | `utils/trainer.py:31` | Training configuration |
| `PipelineParallelTrainerConf` | `utils/trainer.py:178` | Multi-GPU pipeline parallel training |

#### Training Infrastructure (`labml_nn/helpers/`)

| Class | File | Purpose |
|---|---|---|
| `TrainValidConfigs` | `trainer.py:358` | Main training loop with train/validate phases |
| `SimpleTrainValidConfigs` | `trainer.py:448` | Simplified training (no validation) |
| `TrainingLoop` | `trainer.py:60` | Iterator-based training with SIGINT handling |
| `Trainer` | `trainer.py:254` | Epoch-based trainer with data loader |
| `BatchIndex` | `trainer.py:310` | Batch progress tracking |
| `Mode` | `trainer.py:237` | Context manager for train/eval mode |
| `Accuracy` | `metrics.py:43` | Accuracy metric tracking |
| `DeviceConfigs` | `device.py:34` | CPU/GPU device selection |
| `Schedule` | `schedule.py:4` | Learning rate schedule base |
| `Piecewise` | `schedule.py:34` | Piecewise linear schedule |
| `MNISTConfigs` | `datasets.py:24` | MNIST data loader configs |
| `CIFAR10Configs` | `datasets.py:106` | CIFAR-10 data loader configs |
| `TextDataset` | `datasets.py:180` | Character-level text dataset |
| `SequentialDataLoader` | `datasets.py:232` | Sequential text data loader |

### Error Messages Reference

| Error | Location | Trigger |
|---|---|---|
| `ValueError: Invalid learning rate: {lr}` | `optimizers/__init__.py:88` | `lr < 0` |
| `ValueError: Invalid epsilon value: {eps}` | `optimizers/__init__.py:90` | `eps < 0` |
| `ValueError: Invalid beta parameter at index 0: {betas[0]}` | `optimizers/__init__.py:92` | `beta1 < 0` or `beta1 >= 1` |
| `ValueError: Invalid beta parameter at index 1: {betas[1]}` | `optimizers/__init__.py:94` | `beta2 < 0` or `beta2 >= 1` |
| `ValueError: Invalid weight_decay value: {weight_decay}` | `optimizers/__init__.py:186` | `weight_decay < 0` |
| `RuntimeError: GenericAdaptiveOptimizer does not support sparse gradients` | `optimizers/__init__.py:150` | Sparse gradient tensor |
| `ValueError: Head size {d_k} too large for flash attention` | `neox/model.py:308` | Per-head dim exceeds flash attn limit |
| `ImportError: Please install bitsandbytes` | `neox/utils/llm_int8.py:37` | Missing bitsandbytes package |
| `FileNotFoundError` | `neox/utils/text_dataset.py:39` | Dataset file not found |
| `FileNotFoundError` | `helpers/datasets.py:300` | Text file not found |
| `StopIteration: SIGINT` | `helpers/trainer.py:119` | Ctrl+C during training |

### Code Patterns and Conventions

**Module Structure Pattern:**

```text
labml_nn/{topic}/
├── __init__.py          Module catalog (docstring) + core implementation
├── experiment.py        Training experiment using Configs
└── {variant}/           Sub-implementations
    ├── __init__.py      Variant implementation
    └── experiment.py    Variant experiment
```

**Naming Conventions:**

- Module classes: PascalCase matching the paper name (e.g., `MultiHeadAttention`, `ConvMixer`, `RWKV`)
- Config classes: Always named `Configs`, disambiguated by module path
- Experiment entry points: `main()` function in `experiment.py`
- Constants: UPPER_SNAKE_CASE (e.g., `CHECKPOINTS_URL`, `PREV_X_TIME`)

**Docstring Style:**

- YAML frontmatter for title/summary (consumed by doc renderer)
- Markdown with LaTeX math in `$...$` and `$$...$$`
- Inline comments explain every non-obvious line
- HTML anchor tags (`<a id="...">`) for cross-referencing

> [!NOTE]
> The codebase uses a mix of `labml.configs` options and direct class inheritance. Newer implementations tend toward simpler direct inheritance; older ones use the full `@option`/`aggregate` system. Both patterns are valid.

### Testing

No dedicated test suite is included in the repository. The primary verification is that each experiment script runs successfully end-to-end. The Capsule Networks implementation includes a Jupyter notebook that serves as an integration test.

> [!WARNING]
> There is no automated test suite. When modifying implementations, verify correctness by running the corresponding experiment and checking that loss curves match expected behavior.

### Known Issues (from Community)

| Issue | Status | Impact |
|---|---|---|
| [#107](https://github.com/labmlai/annotated_deep_learning_paper_implementations/issues/107) StyleGAN2 ignores latent vector | Open | StyleGAN output is not properly conditioned |
| [#135](https://github.com/labmlai/annotated_deep_learning_paper_implementations/issues/135) RETRO tensor size mismatch | Open | Batch processing crashes with variable-length inputs |
| [#165](https://github.com/labmlai/annotated_deep_learning_paper_implementations/issues/165) Missing activation in DDPM ResidualBlock | Open | Potential accuracy impact on diffusion models |
| [#177](https://github.com/labmlai/annotated_deep_learning_paper_implementations/issues/177) Classifier-free guidance error | Open | Incorrect diffusion sampling behavior |
| [#185](https://github.com/labmlai/annotated_deep_learning_paper_implementations/issues/185) Transformer-XL shift method bug | Open | Incorrect relative position encoding |
| [#216](https://github.com/labmlai/annotated_deep_learning_paper_implementations/issues/216) Potential errors in MHA | Open | Attention computation may be incorrect |
| [#238](https://github.com/labmlai/annotated_deep_learning_paper_implementations/issues/238) DDPM interpolation formula error | Open | Noise schedule calculation error |
| [#245](https://github.com/labmlai/annotated_deep_learning_paper_implementations/issues/245) PPO running error | Open | Gym version incompatibility |
| [#255](https://github.com/labmlai/annotated_deep_learning_paper_implementations/issues/255) GAE formula bug | Closed | Fixed in current version |
| [#270](https://github.com/labmlai/annotated_deep_learning_paper_implementations/issues/270) UNet error in DDPM | Open | Shape mismatch in skip connections |
| [#285](https://github.com/labmlai/annotated_deep_learning_paper_implementations/issues/285) TransformerXLLayer issues | Open | Code snippet problems |
| [#293](https://github.com/labmlai/annotated_deep_learning_paper_implementations/issues/293) Nucleus sampling doc corrections | Open | Documentation inaccuracies |

### Potential Improvements

Based on source analysis:

1. **Add automated tests.** No test suite exists — each implementation relies on end-to-end experiment execution for verification. Unit tests for core modules (attention, optimizers, normalization) would catch regressions.

2. **Pin dependency versions more tightly.** `setup.py` specifies no version bounds for most dependencies, while `requirements.txt` has minimum bounds. The `fairscale` dependency has no version constraint at all.

3. **Address open bugs.** Multiple community-reported issues (StyleGAN latent vector, DDPM formula errors, Transformer-XL shift) remain open. These affect correctness of the educational implementations.

4. **Standardize experiment entry points.** Some experiments use `main()` functions, others use `if __name__ == '__main__'` blocks. A consistent CLI pattern would improve discoverability.

5. **Add type hints to remaining modules.** Many modules have partial type annotations. Full typing would improve IDE support and catch errors.

6. **Update deprecated gym API.** The RL experiments use the legacy `gym` API (pre-Gymnasium). Migrating to `gymnasium` would fix version compatibility issues ([#245](https://github.com/labmlai/annotated_deep_learning_paper_implementations/issues/245)).

### Contributing

The project follows a literate-programming style. When adding a new paper implementation:

1. Create a new subpackage under `labml_nn/`
2. Write the implementation with extensive inline docstrings containing:
   - YAML frontmatter (`title`, `summary`)
   - Markdown explanation with LaTeX math
   - Line-by-line comments for non-obvious operations
3. Create an `experiment.py` with a `Configs` class and `main()` entry point
4. Add the implementation to the catalog in `labml_nn/__init__.py`

See community discussion [#275](https://github.com/labmlai/annotated_deep_learning_paper_implementations/issues/275) for contribution guidelines.

---

## 9. Troubleshooting

### Training hangs or produces NaN loss

**Cause:** Learning rate too high, or numerical instability in FP16 training.

**Solution:**

1. Reduce learning rate
2. If using `AdamFP16`, ensure `GradScalerFP16` is properly scaling gradients
3. Check for division by zero in normalization layers with very small epsilon

> [!WARNING]
> The custom optimizer implementations do not include gradient clipping by default. For unstable training, add `torch.nn.utils.clip_grad_norm_()` before `optimizer.step()`.

### GPT-NeoX checkpoint download fails

**Cause:** Network issues downloading from `mystic.the-eye.eu`.

**Solution:**

1. Check network connectivity to `https://mystic.the-eye.eu/public/AI/models/GPT-NeoX-20B/slim_weights/`
2. Manually download checkpoint files to `lab.get_data_path() / 'neox' / 'slim_weights'`
3. Ensure sufficient disk space (~40 GB)

> [!TIP]
> If the default download path doesn't work, the checkpoint code checks for both `neox_fast/slim_weights` and `neox/slim_weights` directories under `lab.get_data_path()`.

### labml experiment tracking not working

**Cause:** Missing or misconfigured labml setup.

**Solution:**

1. Ensure `labml` is installed: `pip install labml>=0.4.147`
2. Run `labml init` to create a `.labml.yaml` configuration file
3. Check that `lab.get_data_path()` points to a writable directory

---

_Generated May 2026 from labmlai/annotated_deep_learning_paper_implementations at 33ab022. AI-generated documentation. The library is checked by sampling (three repositories so far); this guide was not individually reviewed. Verify against the project before relying on it. Unofficial; not affiliated with or reviewed by the project._