diff --git a/distrax/_src/distributions/mvn_from_bijector.py b/distrax/_src/distributions/mvn_from_bijector.py index eb70a3e..96e29ff 100644 --- a/distrax/_src/distributions/mvn_from_bijector.py +++ b/distrax/_src/distributions/mvn_from_bijector.py @@ -29,6 +29,7 @@ import jax import jax.numpy as jnp +import jax.scipy as jsp from tensorflow_probability.substrates import jax as tfp @@ -150,6 +151,56 @@ def stddev(self) -> Array: result = jnp.sqrt(self.variance()) return jnp.broadcast_to(result, self.batch_shape + self.event_shape) + def symmetrized_kl_divergence(self, other: 'MultivariateNormalLike') -> Array: + """Computes KL(self || other) + KL(other || self) without subtractive cancellation.""" + m, _ = _cholesky_cross_operator_from_dist(self, other) + cov_term = 0.5 * jnp.sum(jnp.square(m), axis=(-2, -1)) + + # pyrefly: ignore[bad-index,unsupported-operation] + delta_mu = (self.mean() - other.mean())[..., None] + v1 = _inv_scale_operator(self)(delta_mu) + v2 = _inv_scale_operator(other)(delta_mu) + mean_term = 0.5 * jnp.sum(jnp.square(v1) + jnp.square(v2), axis=(-2, -1)) + + return cov_term + mean_term + + def bhattacharyya_distance(self, other: 'MultivariateNormalLike') -> Array: + """Computes the Bhattacharyya distance without subtractive cancellation.""" + _, k = _cholesky_cross_operator_from_dist(self, other) + cov_term = 0.25 * _logdet_one_plus_k(k) + + # pyrefly: ignore[bad-index,unsupported-operation] + delta_mu = (self.mean() - other.mean())[..., None] + sigma_avg = 0.5 * (self.covariance() + other.covariance()) + l_avg = jnp.linalg.cholesky(sigma_avg) + v_avg = jsp.linalg.solve_triangular(l_avg, delta_mu, lower=True) + mean_term = 0.125 * jnp.sum(jnp.square(v_avg), axis=(-2, -1)) + + return cov_term + mean_term + + def geometric_jensen_shannon_divergence( + self, other: 'MultivariateNormalLike' + ) -> Array: + """Computes the Geometric Jensen-Shannon Divergence without subtractive cancellation.""" + _, k = _cholesky_cross_operator_from_dist(self, other) + cov_term = 0.25 * _trace_two_k_minus_log_one_plus_k(k) + + # pyrefly: ignore[bad-index,unsupported-operation] + delta_mu = (self.mean() - other.mean())[..., None] + v1 = _inv_scale_operator(self)(delta_mu) + v2 = _inv_scale_operator(other)(delta_mu) + + sigma_avg = 0.5 * (self.covariance() + other.covariance()) + l_avg = jnp.linalg.cholesky(sigma_avg) + v_avg = jsp.linalg.solve_triangular(l_avg, delta_mu, lower=True) + + mean_term = 0.125 * jnp.sum( + jnp.square(v1) + jnp.square(v2), axis=(-2, -1) + ) - 0.125 * jnp.sum(jnp.square(v_avg), axis=(-2, -1)) + mean_term = jnp.maximum(mean_term, 0.0) + + return cov_term + mean_term + MultivariateNormalLike = Union[ MultivariateNormalFromBijector, tfd.MultivariateNormalLinearOperator] @@ -211,6 +262,73 @@ def _has_diagonal_scale(d: MultivariateNormalLike) -> bool: return False +def _cholesky_cross_operator_from_dist(dist1, dist2): + """Computes cross-operator M and normalized Gramian K directly from bijectors.""" + x = _inv_scale_operator(dist2)(_scale_matrix(dist1)) + y = _inv_scale_operator(dist1)(_scale_matrix(dist2)) + m = x - jnp.swapaxes(y, -1, -2) + k = 0.25 * jnp.matmul(m, jnp.swapaxes(m, -1, -2)) + # Ensure exact symmetry + k = 0.5 * (k + jnp.swapaxes(k, -1, -2)) + return m, k + + +def _derive_taylor_threshold(dtype: jnp.dtype, max_order: int = 4) -> float: + eps = float(jnp.finfo(dtype).eps) + return float((eps * (max_order + 1)) ** (1.0 / (max_order + 1))) + + +def _logdet_one_plus_k(k: jax.Array) -> jax.Array: + """Computes ln det(I + K) for PSD K, stable for small ||K||.""" + tr_k2 = jnp.sum(jnp.square(k), axis=(-2, -1)) + norm_k = jnp.sqrt(tr_k2) + threshold = _derive_taylor_threshold(k.dtype) + + # Since K and K^2 are symmetric, tr(K^3) = _F and + # tr(K^4) = ||K^2||_F^2, requiring only a single matrix multiplication + # (K^2 = K @ K). + k2 = jnp.matmul(k, k) + tr_k = jnp.trace(k, axis1=-2, axis2=-1) + tr_k3 = jnp.sum(k * k2, axis=(-2, -1)) + tr_k4 = jnp.sum(jnp.square(k2), axis=(-2, -1)) + logdet_taylor = ( + tr_k - 0.5 * tr_k2 + (1.0 / 3.0) * tr_k3 - 0.25 * tr_k4 + ) + + # Large-norm Cholesky: 2 * sum(log(diag(chol(I + K)))) + eye = jnp.eye(k.shape[-1], dtype=k.dtype) + l_k = jnp.linalg.cholesky(eye + k) + logdet_chol = 2.0 * jnp.sum( + jnp.log(jnp.diagonal(l_k, axis1=-2, axis2=-1)), axis=-1 + ) + return jnp.where(norm_k < threshold, logdet_taylor, logdet_chol) + + +def _trace_two_k_minus_log_one_plus_k(k: jax.Array) -> jax.Array: + """Computes tr(2K - ln(I + K)) for PSD K without subtractive cancellation.""" + tr_k2 = jnp.sum(jnp.square(k), axis=(-2, -1)) + norm_k = jnp.sqrt(tr_k2) + threshold = _derive_taylor_threshold(k.dtype) + tr_k = jnp.trace(k, axis1=-2, axis2=-1) + + # Small-norm Taylor expansion: tr(K + K^2/2 - K^3/3 + K^4/4) + k2 = jnp.matmul(k, k) + tr_k3 = jnp.sum(k * k2, axis=(-2, -1)) + tr_k4 = jnp.sum(jnp.square(k2), axis=(-2, -1)) + val_taylor = ( + tr_k + 0.5 * tr_k2 - (1.0 / 3.0) * tr_k3 + 0.25 * tr_k4 + ) + + # Large-norm Cholesky: 2 tr(K) - ln det(I + K) + eye = jnp.eye(k.shape[-1], dtype=k.dtype) + l_k = jnp.linalg.cholesky(eye + k) + logdet_chol = 2.0 * jnp.sum( + jnp.log(jnp.diagonal(l_k, axis1=-2, axis2=-1)), axis=-1 + ) + val_chol = jnp.maximum(2.0 * tr_k - logdet_chol, 0.0) + return jnp.where(norm_k < threshold, val_taylor, val_chol) + + def _kl_divergence_mvn_mvn( dist1: MultivariateNormalLike, dist2: MultivariateNormalLike, diff --git a/distrax/_src/distributions/mvn_from_bijector_test.py b/distrax/_src/distributions/mvn_from_bijector_test.py index b499687..db7318e 100644 --- a/distrax/_src/distributions/mvn_from_bijector_test.py +++ b/distrax/_src/distributions/mvn_from_bijector_test.py @@ -297,6 +297,58 @@ def test_kl_divergence_raises_on_incompatible_distributions(self): with self.assertRaises(ValueError): dist1.kl_divergence(dist2) + @chex.all_variants(with_pmap=False) + @parameterized.named_parameters( + ('float32', jnp.float32), + ('float64', jnp.float64), + ) + def test_divergences_cancellation_dtypes(self, dtype): + prng = hk.PRNGSequence(jax.random.PRNGKey(42)) + eps = 1e-4 # Using 1e-4 perturbation + dim = 5 + + # Create two nearby distributions + loc1 = jax.random.normal(next(prng), (dim,), dtype=dtype) + loc2 = loc1 + eps * jax.random.normal(next(prng), (dim,), dtype=dtype) + + scale_tril1 = jnp.tril( + jax.random.normal(next(prng), (dim, dim), dtype=dtype) + ) + # Ensure it is well conditioned by adding to diagonal + scale_tril1 = scale_tril1 + 3.0 * jnp.eye(dim, dtype=dtype) + + scale_tril2 = scale_tril1 + eps * jax.random.normal( + next(prng), (dim, dim), dtype=dtype + ) + scale_tril2 = jnp.tril(scale_tril2) + + dist1 = MultivariateNormalFromBijector( + loc=loc1, + scale=TriangularLinear(matrix=scale_tril1), + ) + dist2 = MultivariateNormalFromBijector( + loc=loc2, + scale=TriangularLinear(matrix=scale_tril2), + ) + + # New methods calls + # pyrefly: ignore[missing-attribute] + sym_kl = self.variant(dist1.symmetrized_kl_divergence)(dist2) + # pyrefly: ignore[missing-attribute] + bhattacharyya = self.variant(dist1.bhattacharyya_distance)(dist2) + # pyrefly: ignore[missing-attribute] + gjsd = self.variant(dist1.geometric_jensen_shannon_divergence)(dist2) + + # We expect these values to be > 0 and computed accurately. + self.assertTrue(jnp.all(sym_kl > 0.0)) + self.assertTrue(jnp.all(bhattacharyya > 0.0)) + self.assertTrue(jnp.all(gjsd > 0.0)) + + # At small perturbation (eps = 1e-4), D_B and G-JSD asymptotically approach + # 1/8 D_J (since sym_kl = KL(1||2) + KL(2||1) = 2 * D_sym_half). + np.testing.assert_allclose(bhattacharyya, 0.125 * sym_kl, rtol=0.01) + np.testing.assert_allclose(gjsd, 0.125 * sym_kl, rtol=0.01) + if __name__ == '__main__': jax.config.update('jax_threefry_partitionable', False)