Annotated Deep Learning Paper Implementations — User Guide

1. Quick Start

labml_nn is a collection of simple, annotated PyTorch implementations of neural networks and related algorithms from landmark deep learning research papers. Each implementation is documented with inline explanations, and the companion website (nn.labml.ai) renders the source code as side-by-side formatted notes — code on one side, explanation on the other.

Key features:

• 60+ paper implementations covering transformers, diffusion models, GANs, optimizers, normalization layers, reinforcement learning, and more
• Every implementation is a working PyTorch module you can import and use directly
• Side-by-side annotated documentation at nn.labml.ai
• Experiment scripts included for training on standard datasets (MNIST, CIFAR-10, Tiny Shakespeare, Atari Breakout)
• Translations available in Chinese, Japanese, and Sinhala
• MIT licensed

Prerequisites: Python 3.7+, PyTorch 1.10+

Install:

pip install labml-nn

Minimal example — use a multi-head attention layer:

from labml_nn.transformers import MultiHeadAttention
import torch

# 8-head attention with 512-dimensional model
mha = MultiHeadAttention(heads=8, d_model=512)
q = k = v = torch.randn(32, 10, 512)  # batch=32, seq=10, dim=512
output = mha(query=q, key=k, value=v)
print(output.shape)

torch.Size([32, 10, 512])

Tip

Browse the full implementation catalog at nn.labml.ai to find the paper you want, then import the corresponding module from labml_nn.

What to try next:

• Browse the implementation catalog below to find specific papers
• Run one of the included experiment scripts to train a model
• Read the Admin Guide for architecture details and module internals

2. Installation

System Requirements

Install Methods

pip (recommended)

pip install labml-nn

From source

git clone https://github.com/labmlai/annotated_deep_learning_paper_implementations.git
cd annotated_deep_learning_paper_implementations
pip install -r requirements.txt
pip install -e .

Important

The package name on PyPI is labml-nn (with a hyphen), but the Python import is labml_nn (with an underscore).

Dependencies

Optional dependencies (required only for specific implementations):

Post-Install Verification

python -c "import labml_nn; print('labml_nn installed successfully')"

labml_nn installed successfully

Warning

If you see ModuleNotFoundError: No module named 'labml', ensure you installed labml-nn (which pulls in labml as a dependency), not just labml_nn from source without installing dependencies.

Install Troubleshooting

Error: pip install labml-nn generated errors

Cause: Version conflicts between PyTorch, torchvision, and torchtext. Community issue #247.

Solution:

1. Install PyTorch first following pytorch.org/get-started for your platform
2. Then install labml-nn:

pip install labml-nn

Prevention: Always install PyTorch before labml-nn to ensure compatible versions.

Error: No module named 'labml_helpers'

Cause: The labml-helpers package was not installed alongside labml-nn.

Solution:

pip install labml-helpers>=0.4.84

3. Implementation Catalog

The library organizes implementations by research area. Every module listed below is a working PyTorch implementation you can import and use directly.

Transformers

The largest collection, covering the original transformer and 20+ variants.

Diffusion Models

Generative Adversarial Networks (GANs)

Optimizers

Normalization Layers

Sampling Techniques

Reinforcement Learning

Other Implementations

4. Core Usage

Importing and Using Modules

Every implementation is a standard PyTorch nn.Module. Import the class, instantiate it with the required parameters, and use it in your models.

Using a transformer encoder:

from labml_nn.transformers import Encoder, MultiHeadAttention
from labml_nn.transformers.feed_forward import FeedForward
from labml_nn.transformers.models import TransformerLayer, EmbeddingsWithPositionalEncoding
import torch.nn as nn

# Build a 6-layer transformer encoder
d_model = 512
heads = 8
d_ff = 2048
n_layers = 6

self_attn = MultiHeadAttention(heads=heads, d_model=d_model)
feed_forward = FeedForward(d_model=d_model, d_ff=d_ff)
layer = TransformerLayer(d_model=d_model, self_attn=self_attn,
                         src_attn=None, feed_forward=feed_forward)
encoder = Encoder(layer=layer, n_layers=n_layers)

Note

Each module is self-contained. You can use individual layers (like MultiHeadAttention) in your own architectures without importing the full transformer stack.

Using a normalization layer:

from labml_nn.normalization.layer_norm import LayerNorm
import torch

ln = LayerNorm(normalized_shape=[512])
x = torch.randn(32, 10, 512)
output = ln(x)

Using an optimizer:

from labml_nn.optimizers.adam import Adam
from labml_nn.optimizers import WeightDecay

optimizer = Adam(
    model.parameters(),
    lr=1e-3,
    betas=(0.9, 0.999),
    eps=1e-8,
    weight_decay=WeightDecay(0.01)
)

Tip

The custom optimizer implementations include a WeightDecay class that applies decoupled weight decay (AdamW-style) separately from the gradient update. Pass WeightDecay(decay_rate) to any optimizer's weight_decay parameter.

Using a sampling technique with a language model:

from labml_nn.sampling.nucleus import NucleusSampler
from labml_nn.sampling.temperature import TemperatureSampler

# Nucleus sampling with p=0.95 and temperature=0.8
sampler = NucleusSampler(p=0.95, sampler=TemperatureSampler(temperature=0.8))
next_token = sampler(logits)

All sampler classes accept a logits tensor of shape [..., n_tokens] and return sampled token indices. You can chain them — for example, TopKSampler wraps an inner TemperatureSampler.

Using the Configuration System

Many experiments use the labml configuration system (BaseConfigs). This provides declarative, composable experiment setups.

from labml_nn.transformers.configs import TransformerConfigs

# Create a configured transformer
conf = TransformerConfigs()
conf.d_model = 256
conf.heads = 4
conf.d_ff = 1024
conf.n_layers = 4

The TransformerConfigs class exposes these configurable parameters:

The FeedForwardConfigs class supports GLU variants and gating:

Reading the Annotated Source

The primary value of this project is its annotated source code. Each Python file contains extensive docstrings written as tutorial prose, with mathematical notation in LaTeX. The website at nn.labml.ai renders these as side-by-side views — code on the right, explanation on the left.

To browse locally, the docs/ directory contains pre-rendered HTML files for each implementation.

Note

The documentation website is available in three languages: English, Chinese, and Japanese.

5. Common Tasks

Train a ResNet on CIFAR-10

Goal: Run the included ResNet experiment to train on CIFAR-10.

Prerequisites: labml-nn installed, GPU recommended.

Steps:

1. Run the experiment script:

cd annotated_deep_learning_paper_implementations
python -m labml_nn.resnet.experiment

2. The experiment uses labml for tracking. Monitor progress in your terminal or via the labml dashboard.

Verify: Training loss should decrease over epochs. The experiment uses the CIFAR10Configs configuration class.

Train a Capsule Network on MNIST

Goal: Train a capsule network classifier on MNIST.

Steps:

1. Run the experiment:

python -m labml_nn.capsule_networks.mnist

Or use the provided Colab notebook:

Open In Colab

Experiment with Language Model Sampling

Goal: Compare greedy, temperature, top-k, and nucleus sampling using GPT-2.

Prerequisites: transformers package installed (pip install transformers).

Steps:

1. Run the sampling experiment:

python -m labml_nn.sampling.experiment

2. This loads HuggingFace's GPT-2 model and generates text using each sampling strategy with the prompt "I saw an interesting dream last night."

Tip

Modify the prompt variable in labml_nn/sampling/experiment.py to test with your own prompts. Adjust sampler parameters (temperature, k, p) to see how they affect output diversity.

Generate Text with GPT-NeoX

Goal: Run inference with the 20-billion parameter GPT-NeoX model.

Prerequisites: A GPU with 48 GB+ VRAM. The checkpoint download requires ~40 GB of disk space.

Steps:

1. Checkpoints are automatically downloaded on first run:

python -m labml_nn.neox.samples.generate

2. The script downloads slim checkpoint weights from the Eleuther AI repository and loads them into the model.

Warning

GPT-NeoX requires significant GPU memory. For reduced memory usage, use the LLM.int8() quantized version at labml_nn.neox.samples.llm_int8, which requires the bitsandbytes package (pip install bitsandbytes -U).

Fine-Tune GPT-2 with LoRA

Goal: Fine-tune GPT-2 on Tiny Shakespeare using Low-Rank Adaptation.

Steps:

1. Run the LoRA experiment:

python -m labml_nn.lora.experiment

2. This freezes the pre-trained GPT-2 weights and trains only the low-rank decomposition matrices (A and B), dramatically reducing trainable parameters. The scaling factor is controlled by alpha / r, where r is the rank of the decomposition.

Train a U-Net for Image Segmentation

Goal: Train U-Net on the Carvana car segmentation dataset.

Prerequisites: Download the Carvana dataset, opencv-python installed.

Steps:

1. Run the experiment:

python -m labml_nn.unet.experiment

Run PPO on Atari Breakout

Goal: Train a PPO agent to play Atari Breakout.

Prerequisites: gym[atari] and opencv-python installed.

Steps:

1. Run the PPO experiment:

python -m labml_nn.rl.ppo.experiment

Note

The RL experiments use a custom Game wrapper that applies the same action for 4 frames, converts observations to grayscale at 84×84 resolution, and stacks 4 consecutive frames. Each life is treated as a separate episode. The environment is BreakoutNoFrameskip-v4.

Run Knowledge Distillation on CIFAR-10

Goal: Train a small model to mimic a larger model's predictions.

Steps:

1. First train the large (teacher) model:

python -m labml_nn.distillation.large

2. Then train the small (student) model:

python -m labml_nn.distillation.small

3. Finally, run the distillation experiment that transfers knowledge:

python -m labml_nn.distillation

Train RWKV Language Model

Goal: Train an RWKV model, a recurrence-based alternative to transformers.

Steps:

1. Run the RWKV experiment:

python -m labml_nn.rwkv.experiment

Tip

RWKV uses linear-time inference (O(1) per token) compared to transformers' quadratic attention, making it more efficient for long sequences.

6. Troubleshooting

Error: RuntimeError: stack expects each tensor to be equal size, but got [2, 32] at entry 0 and [1, 32] at entry 29

Cause: Mismatched tensor sizes in the RETRO implementation when batch elements have different chunk counts. Community issue #135.

Solution: Ensure all inputs in the batch have the same sequence length. Pad shorter sequences to match the longest sequence in the batch before passing to the RETRO model.

Prevention: Use a collation function that pads sequences to uniform length.

Error: PPO code running error

Cause: Version incompatibilities with the gym package. Community issue #245.

Solution:

1. Install a compatible version of gym:

pip install gym[atari]==0.21.0

2. Ensure opencv-python is installed for frame processing.

Error: U-Net error in DDPM

Cause: Shape mismatch in the U-Net skip connections used within the DDPM implementation. Community issue #270.

Solution: Ensure input image dimensions are divisible by the number of downsampling stages. For the default DDPM U-Net configuration, use image sizes that are multiples of 16.

Warning

The DDPM implementation uses center-cropping for skip connections (community issue #281 discusses why center crop is used instead of padding). If input dimensions are not cleanly divisible, spatial dimension mismatches will occur during the concatenation step.

Error: ValueError: Invalid learning rate: -0.001

Cause: Negative learning rate passed to an optimizer. The GenericAdaptiveOptimizer base class validates that lr >= 0, eps >= 0, and 0 <= beta < 1 for both beta parameters.

Solution: Ensure all optimizer hyperparameters are non-negative:

from labml_nn.optimizers.adam import Adam

# Correct: all positive values
optimizer = Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.999), eps=1e-8)

Error: ImportError: Please install bitsandbytes

Cause: The bitsandbytes package is required for LLM.int8() quantization but not installed.

Solution:

pip install bitsandbytes -U

Error: Head size {d_k} too large for flash attention

Cause: The GPT-NeoX attention layer raises this when the per-head dimension exceeds what the flash attention kernel supports.

Solution: Use the standard attention path instead of flash attention, or reduce the model's head dimension.

Error: Transformer-XL shift method bug

Cause: Potential issue in the relative position encoding shift operation. Community issue #185.

Solution: Check the labml_nn.transformers.xl.relative_mha module for the latest fixes. Update to the latest labml-nn.

Caution

Several implementations include known community-reported issues. Always check the GitHub Issues page for the latest status before relying on an implementation in production.

Error: RuntimeError: GenericAdaptiveOptimizer does not support sparse gradients

Cause: Passing sparse gradients (e.g., from nn.Embedding with sparse=True) to the custom optimizers.

Solution: Use PyTorch's built-in SparseAdam for layers with sparse gradients, or set sparse=False on the embedding layer.

Error: Missing FloatDynamicHyperParam / IntDynamicHyperParam in labml.configs

Cause: These classes may not be available in your installed version of labml. Community issue #291.

Solution:

pip install labml --upgrade

Error: DDPM interpolation formula error

Cause: Potential mathematical error in the DDPM noise schedule interpolation. Community issue #238.

Solution: Cross-reference the implementation against the original DDPM paper equations. Check the GitHub issue for discussion and proposed fixes.

Error: GAE formula bug in PPO

Cause: Incorrect Generalized Advantage Estimation computation. Community issue #255 (fixed/closed).

Solution: Update to the latest version of labml-nn where this has been addressed.

7. FAQ

Q: Is this library meant for production use or for learning?

The primary purpose is educational — each implementation prioritizes clarity and readability over optimization. However, every module is a functional PyTorch nn.Module that can be used in real projects. For production workloads, consider using official or optimized implementations (e.g., HuggingFace Transformers) alongside these for understanding.

Q: How do I use a specific paper's implementation in my own project?

Import the relevant module and use it as a standard PyTorch layer:

from labml_nn.normalization.layer_norm import LayerNorm
from labml_nn.transformers.mha import MultiHeadAttention

Each module is self-contained and can be composed with your own code.

Q: What datasets do the experiments use?

Most experiments use standard benchmarks: MNIST, CIFAR-10, Tiny Shakespeare (text), Carvana (segmentation), and Atari games (RL). Datasets are downloaded automatically via torchvision or labml utilities on first run.

Q: Do I need a GPU?

Not for importing modules or small-scale experimentation. However, training experiments (especially GPT-NeoX, Stable Diffusion, or RL agents) require CUDA-capable GPUs. GPT-NeoX inference requires 48 GB+ VRAM (or use the LLM.int8() variant for reduced memory).

Q: How do I contribute a new paper implementation?

Follow the project's annotation style — each Python file should contain literate-programming-style docstrings that explain the algorithm alongside the code. See the contributing discussion on GitHub.

Q: The StyleGAN 2 implementation seems to ignore the latent vector. Is this a bug?

Yes, this is a known issue (#107). Check the GitHub issue for the latest status and community workarounds.

Q: Is the classifier-free guidance in the diffusion implementation correct?

Community issue #177 reports a potential error. Verify against the original paper if using in production.

Q: What is the labml configuration system?

The labml library provides a declarative configuration system where experiment parameters are defined as class attributes on BaseConfigs subclasses, with options registered via @option decorators. This allows composable experiment definitions — see the Admin Guide for full details.

Q: Where are the annotated docs rendered?

The website nn.labml.ai renders the annotated source as side-by-side formatted notes. The docs/ directory in the repository contains pre-rendered HTML pages. Translations are available in Chinese (/zh/), Japanese (/ja/), and Sinhala (/si/).

Q: Can I use this with JAX instead of PyTorch?

The library is PyTorch-native. However, there is a JAX implementation of the basic transformer at labml_nn.transformers.jax_transformer. All other implementations are PyTorch only.

Q: How do I train with multiple GPUs?

The GPT-NeoX implementation includes pipeline-parallel training support via PipelineParallelTrainerConf. The LoRA fine-tuning example trains on Tiny Shakespeare. For other models, integrate with PyTorch's DistributedDataParallel or use fairscale (included as a dependency).