Machine Learning's Ultimate Learning Archive
A clean, educational implementation of Machine Learning's most influential works using JAX.
Mula is a curated collection of seminal machine learning papers and algorithms, reimplemented from scratch in JAX. The goal is to provide:
- Clear, readable implementations that prioritize understanding over performance
- Educational focus with detailed comments and documentation
- Modern tooling leveraging JAX's functional approach and automatic differentiation
- Reproducible results matching original papers where possible
- Deep Q-Network (DQN) - Mnih et al., 2015
- Proximal Policy Optimization (PPO) - Schulman et al., 2017
- Deep Deterministic Policy Gradient (DDPG) - Lillicrap et al., 2015
- Soft Actor-Critic (SAC) - Haarnoja et al., 2018
- Convolutional Neural Networks - LeCun et al., 1989
- ResNet - He et al., 2015
- Vision Transformer (ViT) - Dosovitskiy et al., 2020
- YOLO - Redmon et al., 2016
- Attention is All You Need (Transformer) - Vaswani et al., 2017
- BERT - Devlin et al., 2018
- GPT - Radford et al., 2018
- Qwen - Bai et al., 2023
- Backpropagation - Rumelhart et al., 1986
- Adam Optimizer - Kingma & Ba, 2014
- Batch Normalization - Ioffe & Szegedy, 2015
- Dropout - Srivastava et al., 2014
- Variational Autoencoders (VAE) - Kingma & Welling, 2013
- Generative Adversarial Networks (GAN) - Goodfellow et al., 2014
- Diffusion Models (DDPM) - Ho et al., 2020 (paper)
- Graph Convolutional Networks (GCN) - Kipf & Welling, 2017 (paper)
- Graph Attention Networks (GAT) - Veličković et al., 2018 (paper)
- Python 3.12+
- uv (recommended) or pip
Using uv (recommended):
# Clone the repository
git clone https://github.com/yourusername/mula.git
cd mula
# Install dependencies
uv syncUsing pip:
pip install -e .mula/
├── docs/ # Documentation and paper summaries
├── mula/
│ ├── rl/ # Reinforcement Learning implementations
│ ├── cv/ # Computer Vision implementations
│ ├── nlp/ # NLP implementations
│ ├── generative/ # Generative models
│ └── commons/ # Shared utilities and helpers
├── examples/ # Example scripts and notebooks
├── tests/ # Unit tests
└── main.py # Entry point
Each implementation includes:
- Links to the original paper
- Mathematical background and intuition
- Code walkthrough with inline comments
- Training tips and hyperparameters
- Functional Programming: Clean, composable code
- Auto-differentiation: Native gradient computation
- JIT Compilation: Performance when needed
- Hardware Acceleration: Seamless GPU/TPU support
Apache 2.0 License - see LICENSE for details
This project stands on the shoulders of giants. We acknowledge all the researchers whose groundbreaking work made modern machine learning possible.
Built with ❤️ and JAX