Just Relax It

Discrete Variables Relaxation

"Just Relax It" is a cutting-edge Python library designed to streamline the optimization of discrete probability distributions in neural networks, offering a suite of advanced relaxation techniques compatible with PyTorch.

📬 Assets

1. Technical Meeting 1 - Presentation
2. Technical Meeting 2 - Jupyter Notebook
3. Technical Meeting 3 - Jupyter Notebook
4. Blog Post 1
5. Blog Post 2
6. Documentation
7. Tests
8. Technical Report

💡 Motivation

For lots of mathematical problems we need an ability to sample discrete random variables. The problem is that due to continuous nature of deep learning optimization, the usage of truly discrete random variables is infeasible. Thus we use different relaxation methods. One of them, Concrete distribution or Gumbel-Softmax (this is one distribution proposed in parallel by two research groups) is implemented in different DL packages. In this project we implement different alternatives to it.

🗃 Algorithms

• Relaxed Bernoulli, also see 📝 paper
• Correlated relaxed Bernoulli, also see 📝 paper
• Gumbel-Softmax TOP-K, also see 📝 paper
• Straight-Through Bernoulli, also see 📝 paper
• Stochastic Times Smooth, also see 📝 paper
• Invertible Gaussian with KL implemented, also see 📝 paper
• Hard Concrete, also see 📝 paper
• Logistic-Normal and Laplace-form approximation of Dirichlet, also see ℹ️ wiki and 💻 stackexchange
• Generalized Gumbel-Softmax, also see 📝 paper
• REBAR, also see 📝 paper
• Decoupled Straight-Through Gumbel-Softmax, also see 📝 paper
• RELAX, also see 📝 paper

🛠️ Install Using uv (Recommended)

For Production

uv pip install relaxit

For Development

git clone https://github.com/intsystems/relaxit
cd relaxit
uv venv                    # create venv
source .venv/bin/activate  # activate venv
uv sync                    # install all the dependencies
uv pip install -e .        # make the relaxit package editable

To run tests:

uv run pytest tests/

To run Python scripts:

uv run python demo/vae_hard_concrete.py

To run notebooks:

uv run jupyter lab

⚒️ Install Using pip

For Production

pip install -r requirements.txt

For Development

pip install -r requirements-dev.txt

🚀 Quickstart

import torch
from relaxit.distributions import InvertibleGaussian

# initialize distribution parameters
loc = torch.zeros(3, 4, 5, requires_grad=True)
scale = torch.ones(3, 4, 5, requires_grad=True)
temperature = torch.tensor([1e-0])

# initialize distribution
distribution = InvertibleGaussian(loc, scale, temperature)

# sample with reparameterization
sample = distribution.rsample()
print('sample.shape:', sample.shape)
print('sample.requires_grad:', sample.requires_grad)

🎮 Demo

Laplace Bridge	REINFORCE in Acrobot environment	VAE with discrete latents

For demonstration purposes, we divide our algorithms in three[^*] different groups. Each group relates to the particular demo code:

• Laplace bridge between Dirichlet and LogisticNormal distributions
• REINFORCE
• RELAX
• Other relaxation methods

We describe our demo experiments here.

[^*]: We also implement REINFORCE algorithm as a score function estimator alternative for our relaxation methods that are inherently pathwise derivative estimators. This one is implemented only for demo experiments and is not included into the source code of package.

📚 Stack

Some of the alternatives for GS were implemented in pyro, so we base our library on their codebase.

🧩 Some details

To make to library consistent, we integrate imports of distributions from pyro and torch into the library, so that all the categorical distributions can be imported from one entrypoint.

👥 Contributors

• Daniil Dorin (Basic code writing, Final demo, Algorithms)
• Igor Ignashin (Project wrapping, Documentation writing, Algorithms)
• Nikita Kiselev (Project planning, Blog post, Algorithms)
• Andrey Veprikov (Tests writing, Documentation writing, Algorithms)
• Vladislav Minashkin (Project planning, Visualizations, Algorithms)
• Papay Ivan (Documentation writing, Blog post, Algorithms)
• Meshkov Vlad (Benchmarking, Demo, Algorithms)
• Stepanov Ilya (Tech. report, Code writing, Algorithms)
• You are welcome to contribute to our project!

🔗 Useful links

• About top-k GS
• VAE implementation with different latent distributions
• KL divergence between Dirichlet and Logistic-Normal implemented in R
• About score function (SF) and pathwise derivate (PD) estimators, VAE and REINFORCE