freegaussianizer 0.1.0

Matricial Free Energy loss for training Gaussianizing autoencoders.

FreeGaussianizer

A Python package implementing the Matricial Free Energy loss function for training Gaussianizing autoencoders, based on the paper:

Sonthalia & Nadakuditi, "Matricial Free Energy as a Gaussianizing Regularizer: Enhancing Autoencoders for Gaussian Code Generation", arXiv:2510.17120

The key idea: minimizing the Free Loss over mini-batches of encoder outputs drives the latent code distribution toward a Gaussian — without any explicit generative model.

This is the Python (PyTorch) port of FreeGaussianizer.jl. The loss is implemented from the singular values of the code matrix (verified against the paper's reference value of −34.69; see Testing).

Installation

pip install freegaussianizer

Or from source:

git clone https://github.com/Thiruloksundar/FreeGaussianizer
cd FreeGaussianizer
pip install -e .

Quick start

import torch
from freegaussianizer import (build_encoder, generate_data, train,
                              ks_statistic, delta_ot, rel_err_free_loss)

# 1. Generate synthetic data (Section 3 of the paper)
X_train, _ = generate_data(1280, p=2, mu=[5.0, 5.0], seed=42)
X_test,  _ = generate_data(1280, p=2, mu=[5.0, 5.0], seed=123)

# 2. Build the paper's MLP encoder (p=2 input -> d=32 embedding)
model = build_encoder(2, 32)

# 3. Train with mini-batched Adam, minimizing the Free Loss
result = train(model, X_train, epochs=2000, lr=1e-2, batch_size=256, X_test=X_test)

# 4. Evaluate Gaussianity of the learned embeddings
with torch.no_grad():
    Y = model(torch.as_tensor(X_train, dtype=torch.float64))  # (n x 32) code matrix
print("KS statistic :", ks_statistic(Y))
print("Delta_OT     :", delta_ot(Y))
print("Rel. err.    :", rel_err_free_loss(Y))

Conventions

The code matrix Y is batch-first: shape (b, d) = (batch size × embedding dimension) with d < b. This matches PyTorch's standard (batch, features) layout, so the encoder is a plain nn.Sequential — no transposes needed. (The Julia package uses the transposed d × b convention; the mathematics is identical.)

API

Loss

Function	Description
free_loss(Y)	Free Loss L_free(Ỹ) = −Φ̄_c(Ỹ) — minimize this to Gaussianize (Eq. 12)
matricial_free_energy(Y)	Matricial free energy Φ̄_c(Ỹ) (Eq. 9)

Y is a (b, d) torch.Tensor with d < b. Both functions are autograd-differentiable for use inside a training loop.

Data

Function	Description
generate_data(n_half, p, mu, seed)	Generate balanced two-class χ²₁ data (Section 3)

Model

Function	Description
build_encoder(p, d, hidden=32)	5-layer MLP encoder matching the paper's architecture

Training

Function	Description
train(model, X_train, ...)	Mini-batched Adam training loop; returns TrainResult(train_losses, test_losses, snapshots)

Evaluation metrics

Function	Description
ks_statistic(entries)	Kolmogorov–Smirnov statistic vs N(0,1)
delta_ot(Y, n_ref)	Relative excess optimal transport cost Δ_OT (Eq. 13)
rel_err_free_loss(Y, n_ref)	Relative deviation of free loss from Gaussian reference
ot_cost(A, B)	Raw optimal transport cost between two matrices

Theoretical references

Function	Description
theoretical_min_free_loss(d, b, ...)	Expected minimum free loss for a Gaussian (b, d) matrix
marchenko_pastur_pdf(x, c)	Marčenko–Pastur density at aspect ratio c = d/b

Examples

The examples/ directory reproduces figures from the paper:

Script	Reproduces
free_loss_verification.py	The three loss-verification checks (paper value, σ / mean minimizers)
figure1.py	Training data scatter + loss curves (Fig. 1)
figure6_7.py	Histogram / Q–Q / eigenvalue snapshots across training (Figs. 6–7)
figure8.py	Ridgeline of code-entry distributions over training (Fig. 8)
figure9_10.py	KS / Δ_OT / RelErr scaling experiments (Figs. 9–10)

Run any example from the repo root (requires pip install -e ".[examples]"):

python examples/figure1.py

Testing

Install the test extras and run the suite with pytest:

pip install -e ".[test]"
pytest -v

The test suite covers:

Test file	What it checks
test_exports.py	All public functions are exported and accessible
test_loss.py	free_loss / matricial_free_energy values, autograd, and the paper's three verification checks (−34.69, σ = 1, mean = 0)
test_data.py	generate_data shape, label balance, reproducibility
test_model.py	build_encoder architecture, forward-pass shape, float64 weights
test_metrics.py	ks_statistic, ot_cost, delta_ot correctness
test_theoretical.py	theoretical_min_free_loss sign, marchenko_pastur_pdf integral

Citation

If you use this package, please cite the original paper:

@misc{sonthalia2025matricial,
  title   = {Matricial Free Energy as a Gaussianizing Regularizer:
             Enhancing Autoencoders for Gaussian Code Generation},
  author  = {Sonthalia, Rishi and Nadakuditi, Raj Rao},
  year    = {2025},
  eprint  = {2510.17120},
  archivePrefix = {arXiv},
}

License

MIT — see LICENSE.