Compatible with PyTorch Inspired by Pyro

Coverage_2 Coverage Docs

License GitHub Contributors Issues GitHub Pull Requests

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

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

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.

• 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 i️ 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

uv pip install relaxit

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

pip install -r requirements.txt

pip install -r requirements-dev.txt

Open In Colab

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)

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.

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

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.

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

• 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