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

Đang mở Phù hợp với người mới
  • Có gì không ổn?
#2,036 0 bình luận 0 reaction 0 người được giao Xem trên GitHub

@Andreas-Hum đang làm issue này rồi.

Từ ngày 1/11/2026.

  • #2037 của @Andreas-Hum — đang mở

Đánh giá

Độ khó
2/5
Thời gian dự kiến
1-3 giờ
Mức phù hợp với người mới
72/100
Loại issue
Lỗi
Độ rõ ràng
Đặc tả rõ ràng
Mức độ hoạt động
Sôi nổi
Công nghệ
python
Lĩnh vực
machine-learning

Hướng nghiên cứu

Bắt đầu từ tensorflow_probability/python/distributions/transformed_distribution.py tại _kl_transformed_transformed, nơi nhánh a.bijector == b.bijector trả về KL của phân phối cơ sở. Kiểm tra cờ _is_injective của bijector và ném NotImplementedError với các bijector không đơn ánh như AbsoluteValue, giữ nguyên các trường hợp đơn ánh. Hoàn thành khi các cặp folded-normal trong bản tái hiện đều ném lỗi và có một bài kiểm thử bao phủ đường đi này.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

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.

Ngôn ngữ chính
Jupyter Notebook
Star
4.4k
Fork
1.1k
Không có pull request nào được merge trong 30 ngày

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của tensorflow/probability

Tất cả issue của tensorflow/probability

Issue tương tự

Thêm issue về Machine Learning

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.