Skip to content

Raise in kl_divergence of TransformedDistributions with a non-injective bijector - #2037

Open
Andreas-Hum wants to merge 2 commits into
tensorflow:mainfrom
Andreas-Hum:kl-transformed-non-injective
Open

Andreas-Hum wants to merge 2 commits into
tensorflow:mainfrom
Andreas-Hum:kl-transformed-non-injective

Conversation

@Andreas-Hum

Copy link
Copy Markdown

Fixes #2036

What

For two TransformedDistributions with equal bijectors, kl_divergence returns the KL of the base distributions. The identity KL(g(P) || g(Q)) = KL(P || Q) holds only for injective g. For a non-injective bijector such as tfb.AbsoluteValue, 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. For example, KL(|N(-1, 1)| || |N(1, 1)|) is returned as 2.0 although the two distributions are identical.

TransformedDistribution._entropy already raises NotImplementedError for non-injective bijectors. This PR makes _kl_transformed_transformed raise as well, with a message naming the bijector. Injective bijectors are unchanged.

Test

  • New testTransformedKLNonInjectiveBijectorFails in transformed_distribution_test.py. It fails on current main (in all three execution regimes) and passes with this change.
  • With this change, all of transformed_distribution_test.py passes: 144 passed, 11 skipped (TensorFlow 2.21, CPU, Python 3.13).

…ve bijector

KL(g(P) || g(Q)) equals KL(P || Q) only for injective g. For a shared
non-injective bijector such as AbsoluteValue, the base KL overstates the
divergence: KL(|N(-1, 1)| || |N(1, 1)|) was returned as 2.0 although the
two distributions are identical. Raise NotImplementedError instead, as
TransformedDistribution._entropy already does for such bijectors.

Fixes tensorflow#2036
@google-cla

google-cla Bot commented Nov 1, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant