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.

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

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

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

Data Flow

The typical workflow for using labml_nn:

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

Required Dependencies

From setup.py:

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

From requirements.txt (pinned versions for development):

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:

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

TransformerConfigs Reference

FeedForwardConfigs Reference

MNISTConfigs Reference

CIFAR10Configs Reference

4. Deployment & Scaling

PyPI Package

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

Package metadata from setup.py:

Documentation Website

The documentation site at 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:

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:

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:

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:

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

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:

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

The GPT-NeoX module hierarchy:

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

The training infrastructure hierarchy:

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:

"""
---
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:

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/)

Optimizer Modules (labml_nn/optimizers/)

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/)

Training Infrastructure (labml_nn/helpers/)

Error Messages Reference

Code Patterns and Conventions

Module Structure Pattern:

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)

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).

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 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