KL divergence between two `TransformedDistribution`s with a shared non-injective bijector (`AbsoluteValue`) returns the base KL

Abierto Apto para principiantes
  • ¿Algo no está bien?
#2,036 0 comentarios 0 reacciones 0 asignados Ver en GitHub

@Andreas-Hum ya está trabajando en esto.

Desde el 1/11/2026.

  • #2037 de @Andreas-Hum — abierto

Evaluación

Dificultad
2/5
Tiempo estimado
1-3 horas
Aptitud para principiantes
72/100
Tipo de issue
Error
Claridad
Bien especificado
Estado de actividad
Activo
Stack tecnológico
python
Área
machine-learning

Línea de trabajo

Empieza en tensorflow_probability/python/distributions/transformed_distribution.py en _kl_transformed_transformed, donde la rama a.bijector == b.bijector devuelve el KL de la distribución base. Comprueba la marca _is_injective del bijector y lanza NotImplementedError para bijectors no inyectivos como AbsoluteValue, dejando los casos inyectivos sin cambios. Se considera terminado cuando los pares folded-normal de la reproducción lanzan un error y una prueba cubre esa ruta.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

Description

tfd.kl_divergence(a, b) for two TransformedDistributions returns kl_divergence(a.distribution, b.distribution) whenever a.bijector == b.bijector. That identity, KL(g(P) || g(Q)) = KL(P || Q), holds only for injective g. For a non-injective bijector the KL of the pushforwards can be strictly smaller, because information is lost: by the chain rule of KL, the difference is E_{Y ~ g(P)} KL(P(. | g = Y) || Q(. | g = Y)) >= 0.

tfb.AbsoluteValue is documented as "This non-injective bijector allows for transformations of scalar distributions with the absolute value function". TransformedDistribution.log_prob handles it correctly by summing over both preimages, but the KL rule does not check injectivity. For example KL(|N(-1, 1)| || |N(1, 1)|) returns 2.0, although the two distributions are identical. TransformedDistribution._entropy already refuses non-injective bijectors (L570-L572).

Reproduction

NumPy substrate (tensorflow_probability.substrates.numpy); the KL rule is the same code for TF and JAX:

import numpy as np
from scipy.integrate import quad
import tensorflow_probability.substrates.numpy as tfp

tfd, tfb = tfp.distributions, tfp.bijectors
print("tensorflow_probability", tfp.__version__, "(NumPy substrate)")


def folded(m, s=1.0):  # |X|, X ~ N(m, s^2)
    return tfd.TransformedDistribution(tfd.Normal(np.float64(m), np.float64(s)), tfb.AbsoluteValue())


def true_kl(p, q):  # KL(p || q) by quadrature, using TFP's own (correct, two-branch) log_prob
    return quad(lambda y: np.exp(p.log_prob(y)) * (p.log_prob(y) - q.log_prob(y)), 0, 40, limit=200)[0]


p0 = folded(-1.0)
print(f"log_prob of |N(-1,1)| integrates to {quad(lambda y: np.exp(p0.log_prob(y)), 0, 40)[0]:.6f}")
for name, p, q in [("|N(-1,1)| vs |N(1,1)| (same law)", folded(-1.0), folded(1.0)),
                   ("|N(3,1)| vs |N(0,1)|", folded(3.0), folded(0.0)),
                   ("|N(0.5,1)| vs |N(0,1)| (= HalfNormal(1))", folded(0.5), folded(0.0))]:
    kl, kl_true = float(tfd.kl_divergence(p, q)), true_kl(p, q)
    print(f"{name:42s} kl_divergence = {kl:.4f}   true KL = {kl_true:.4f}")
    assert kl > kl_true + 0.1
print("Bug reproduced: the base-distribution KL is returned for a non-injective bijector")

Output:

tensorflow_probability 0.25.0 (NumPy substrate)
log_prob of |N(-1,1)| integrates to 1.000000
|N(-1,1)| vs |N(1,1)| (same law)           kl_divergence = 2.0000   true KL = 0.0000
|N(3,1)| vs |N(0,1)|                       kl_divergence = 4.5000   true KL = 3.8107
|N(0.5,1)| vs |N(0,1)| (= HalfNormal(1))   kl_divergence = 0.1250   true KL = 0.0136
Bug reproduced: the base-distribution KL is returned for a non-injective bijector
Expected behavior

Either the true KL or NotImplementedError (there is no closed form for the folded-normal KL). The base KL is right only when the conditional distribution of the sign given |X| is the same under both distributions, e.g. when both bases are centred at 0.

Root cause

distributions/transformed_distribution.py#L731-L754 (v0.25.0, 9709569; identical on main, 6fa068e):

  with tf.name_scope(name or 'kl_transformed_transformed'):
    if a.bijector == b.bijector:
      return kullback_leibler.kl_divergence(a.distribution, b.distribution)
Proposed fix
--- a/tensorflow_probability/python/distributions/transformed_distribution.py
+++ b/tensorflow_probability/python/distributions/transformed_distribution.py
@@ -747,6 +747,10 @@
   """
   with tf.name_scope(name or 'kl_transformed_transformed'):
     if a.bijector == b.bijector:
+      if not a.bijector._is_injective:  # pylint: disable=protected-access
+        raise NotImplementedError(
+            'Unable to calculate KL divergence between {} and {} because '
+            'their bijector {} is not injective.'.format(a, b, a.bijector))
       return kullback_leibler.kl_divergence(a.distribution, b.distribution)
   raise NotImplementedError(
       'Unable to calculate KL divergence between {} and {} because '
Verifying the fix

fix_check.py applies this change to the installed 0.25.0 function (NumPy substrate) in memory and re-registers it. It then checks that the folded pairs raise NotImplementedError, and that pairs with injective bijectors (Exp, Shift(Scale)) give exactly the same values as before:

tensorflow_probability 0.25.0 (NumPy substrate), patched _kl_transformed_transformed
|N(-1,1)| vs |N(1,1)|    -> NotImplementedError: ...value", batch_shape=[], min_event_ndims=0) is not injective.
|N(0.5,1)| vs |N(0,1)|   -> NotImplementedError: ...value", batch_shape=[], min_event_ndims=0) is not injective.
LogNormal via Exp            -> fixed 0.443147, released 0.443147
Shift(1)(Scale(2)) of Normal -> fixed 0.183243, released 0.183243
FIX OK: non-injective bijectors raise NotImplementedError; injective bijectors unchanged
Versions

tensorflow-probability 0.25.0 (NumPy substrate; the JAX substrate of 0.25.0 does not import with the installed jax 0.12), numpy 2.5.4, scipy 1.18.1 for the reference quadrature, Python 3.13, Linux, CPU only.

I would be happy to open a PR with this fix and a test.

Lenguaje dominante
Jupyter Notebook
Estrellas
4.4k
Forks
1.1k
Sin PR fusionados en 30 d

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de tensorflow/probability

Todos los issues de tensorflow/probability

Issues similares

Más issues de Machine Learning

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.