optimal-latent-quantization

Extends classical optimal quantization theory (Lloyd-Max, entropy-constrained vector quantization) to neural latent spaces with learned, potentially non-Gaussian distributions.

Overview

Neural network latent spaces (VAEs, autoencoders, diffusion models) produce representations whose distributions are rarely Gaussian — they can be multimodal, heavy-tailed, or exhibit complex dependencies. Classical quantization theory provides optimal quantizers for uniform and Gaussian sources, but the optimal quantizer for an arbitrary learned distribution requires adapting to the empirical density.

This package implements:

• Lloyd-Max quantization for arbitrary distributions (not just uniform/Gaussian)
• Entropy-constrained vector quantization (ECVQ) for learned latent spaces
• Optimality conditions: centroid + nearest-neighbor + entropy constraint
• Rate-distortion bounds: R(D) for learned distributions
• Non-uniform quantization: optimal bin widths via companding
• Sensitivity analysis: rate penalty of suboptimal (uniform) quantization
• Numerical verification: Lloyd-Max vs uniform on various distributions

Key Mathematics

Lloyd-Max Conditions (Scalar Quantization)

For MSE distortion, the following two conditions are necessary and sufficient for local optimality:

1. Centroid condition: y_i = E[X | X in bin_i] — reconstruction levels are conditional means
2. Nearest-neighbor condition: b_i = (y_i + y_{i+1}) / 2 — boundaries are midpoints

Entropy-Constrained VQ

Minimizes the Lagrangian: L = E[d(X,Y)] + λ·H(Y)

• The Lagrange multiplier λ controls the rate-distortion tradeoff
• At the optimum: dD/dR = -λ (slope of the R-D curve)

Rate Penalty (Uniform vs Optimal)

For non-Gaussian latents with variance σ² and distortion D:

ΔR ≤ ½·log₂(σ²/D) − H_optimal

This quantifies how many extra bits a uniform quantizer wastes compared to the distribution-optimal quantizer. For multimodal latents, this penalty can exceed 1 bit/dimension.

Convergence

Lloyd's algorithm converges to a local optimum in O(log(1/ε)) iterations for an ε-optimal solution.

Installation

git clone git@github.com:drwjkirkpatrick-web/optimal-latent-quantization.git
cd optimal-latent-quantization
pip install -e .

Usage

Lloyd-Max quantization from a density function

from scipy.stats import norm
from optimal_latent_quantization import lloyd_max_quantizer

# Define a density (standard normal)
pdf = lambda x: norm.pdf(x, 0, 1)

# Compute the optimal 8-level quantizer
result = lloyd_max_quantizer(pdf, n_levels=8, x_range=(-6, 6))

print(f"Distortion: {result.distortion:.6f}")
print(f"Rate: {result.rate:.4f} bits")
print(f"Converged: {result.converged} in {result.iterations} iterations")
print(f"Levels: {result.quantizer.levels}")

Lloyd-Max from empirical samples (e.g., VAE latents)

import numpy as np
from optimal_latent_quantization import lloyd_max_from_samples

# Suppose these are latent samples from a trained model
latent_samples = np.random.randn(50000)  # replace with real latents

result = lloyd_max_from_samples(latent_samples, n_levels=16)
print(f"Optimal distortion: {result.distortion:.6f}")

Entropy-constrained VQ

from optimal_latent_quantization import ecvq

result = ecvq(latent_samples, n_levels=32, lambda_val=0.05)
print(f"Distortion: {result.distortion:.6f}")
print(f"Rate: {result.rate:.4f} bits")
print(f"Lagrangian: {result.lagrangian:.6f}")

Compare optimal vs uniform quantization

from optimal_latent_quantization import compare_quantizers

comparison = compare_quantizers(latent_samples, n_levels=16,
                                distribution_name="VAE latents")
print(f"Distortion ratio (uniform/optimal): {comparison.distortion_ratio:.3f}")
print(f"Rate difference: {comparison.rate_at_same_distortion['rate_difference']:.4f} bits")

Rate-distortion bounds

from optimal_latent_quantization import gaussian_rate_distortion, rate_distortion_lower_bound

# Gaussian R(D)
R = gaussian_rate_distortion(sigma_sq=4.0, distortion=1.0)
print(f"Gaussian R(D=1) with σ²=4: {R:.4f} bits")

# Shannon lower bound for arbitrary source
R_lb = rate_distortion_lower_bound(differential_entropy=1.5, distortion=0.5)
print(f"Shannon lower bound: {R_lb:.4f} bits")

Project Structure

optimal-latent-quantization/
├── optimal_latent_quantization/
│   ├── __init__.py          # Public API exports
│   └── core.py              # Core theory implementation
├── tests/
│   └── test_core.py         # 30+ tests (all theory verified)
├── proofs/
│   ├── main_theorem.tex     # Formal LaTeX proofs (3 theorems)
│   └── main_theorem.pdf     # Compiled PDF (force-added)
├── pyproject.toml
├── requirements.txt
├── .gitignore
└── README.md

Formal Proofs

The proofs/ directory contains formal LaTeX proofs of three theorems:

1. Theorem 1 (Lloyd-Max Local Optimality): The centroid and nearest-neighbor conditions are necessary and sufficient for local MSE optimality.
2. Theorem 2 (ECVQ Optimality and Slope Condition): The entropy-constrained Lagrangian yields dD/dR = -λ at the optimum.
3. Theorem 3 (Rate Penalty Bound): ΔR ≤ ½·log₂(σ²/D) − H_optimal.

Compile with: cd proofs && pdflatex main_theorem.tex

Testing

pytest tests/ -v

All 30+ tests verify theoretical properties: optimality conditions, convergence, monotonicity, symmetry, and rate-distortion bounds.

Dependencies

• Python ≥ 3.9
• NumPy ≥ 1.20
• SciPy ≥ 1.7
• pytest (for testing)

License

MIT

Author

Walker Kirkpatrick