A convolutional Beta-VAE built from scratch in JAX. No autograd, no neural network libraries. Every forward pass, backward pass, gradient, and parameter update is written manually using jax.numpy and jax.lax.
Trained on the full CelebA dataset on GPU.
Full writeup here: Convolutional Beta VAE from Scratch in JAX
Evaluated on 1,000 held out images.
| Metric | Value |
|---|---|
| SSIM | 0.9364 |
| PSNR | 22.49 dB |
| Final Reconstruction Loss | 95.38 |
| Final KL Divergence | 110.14 |
| Channel Balance (R / G / B) | 0.56 / 0.49 / 0.44 |
Top row is the original input, bottom row is the model's reconstruction.
Faces generated by sampling z from a standard normal and running it through the decoder.
Linear interpolation between the latent representations of two different identities.
ELBO, reconstruction loss, and KL divergence over 100 epochs.
Everything below was written manually with explicit gradient derivations. No autograd anywhere.
- Conv2D forward and backward (Encoder.py)
- TransposedConv2D forward and backward (Decoder.py)
- Dense layers with manual gradients (Dense.py)
- ReLU and Sigmoid activations (Activation.py)
- Reparameterization trick (Reparameterize.py)
- ELBO loss with reconstruction and KL terms (ELBO.py)
- AdamW optimizer with decoupled weight decay (Adam.py)
- Xavier/Glorot weight initialization
- Gradient clipping
The encoder is three conv layers that downsample the input from 64x64 down to 8x8, followed by a flatten and two dense heads that output μ and log σ². The latent space is 128-dimensional. Sampling uses the reparameterization trick, z = μ + σ ⊙ ε where ε is drawn from a standard normal.
The decoder mirrors the encoder. A dense layer maps the latent vector back to 8192 dimensions, reshapes it, and three transposed conv layers upsample back to 64x64. Final activation is sigmoid.
| Encoder | Channels | Kernel | Stride | Activation | Output |
|---|---|---|---|---|---|
| Conv2D | 3 to 32 | 4 | 2 | ReLU | 32 x 32 x 32 |
| Conv2D | 32 to 64 | 4 | 2 | ReLU | 64 x 16 x 16 |
| Conv2D | 64 to 128 | 4 | 2 | ReLU | 128 x 8 x 8 |
| Flatten + Dense | 8192 to 128 | μ, log σ² |
| Decoder | Channels | Kernel | Stride | Activation | Output |
|---|---|---|---|---|---|
| Dense + Reshape | 128 to 8192 | 128 x 8 x 8 | |||
| TransposedConv2D | 128 to 64 | 4 | 2 | ReLU | 64 x 16 x 16 |
| TransposedConv2D | 64 to 32 | 4 | 2 | ReLU | 32 x 32 x 32 |
| TransposedConv2D | 32 to 3 | 4 | 2 | Sigmoid | 3 x 64 x 64 |
The objective is ELBO = Reconstruction Loss + β × KL Divergence, with β = 0.5.
CelebA, all 202,599 images, resized to 64x64 and stored as CHW tensors. The data loader preloads everything into memory for fast GPU training.
git clone https://github.com/ojayballer/bvaex.git
cd bvaex
pip install jax jaxlib numpy matplotlib pillow
python train.py| Setting | Value |
|---|---|
| GPU | NVIDIA Tesla P100, 16GB |
| Epochs | 100 |
| Batch Size | 512 |
| Optimizer | AdamW, lr=0.001, β₁=0.9, β₂=0.999, weight decay=0.01 |
| Latent Dimension | 128 |
| KL Weight β | 0.5 |
| Training Time | ~2 hours |
| Platform | Kaggle |
JAX compiles everything through XLA, so even though the layers and gradients are all manually implemented, they still run efficiently on GPU.
Trained weights are in weights/epoch_100/.
python evaluate.py
python metrics.py
python interpolate.pyOutputs go to results/. The reconstructions are slightly blurry because of MSE loss. More on that in the blog post.
bvaex/
├── model/
│ ├── __init__.py
│ ├── Encoder.py
│ ├── Decoder.py
│ ├── Dense.py
│ ├── ELBO.py
│ ├── Adam.py
│ ├── Activation.py
│ ├── Reshape.py
│ ├── Reparameterize.py
│ └── model.py
├── train.py
├── evaluate.py
├── metrics.py
├── interpolate.py
├── plots.py
├── load_data.py
├── results/
│ ├── reconstruction_grid.png
│ ├── generated_faces.png
│ ├── latent_interpolation.png
│ └── loss_curves.png
└── weights/
└── epoch_100/
- Kingma and Welling, Auto-Encoding Variational Bayes, ICLR 2014
- Kingma and Welling, An Introduction to Variational Autoencoders, Foundations and Trends in ML
- Kingma and Ba, Adam: A Method for Stochastic Optimization, ICLR 2015
- Higgins et al., β-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework, ICLR 2017
MIT