Skip to content

Integrates numerically stable, subtraction-free formulations of: - #346

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_979240063
Open

Integrates numerically stable, subtraction-free formulations of:#346
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_979240063

Conversation

@copybara-service

Copy link
Copy Markdown

Integrates numerically stable, subtraction-free formulations of:

  • Symmetrized KL (Jeffreys) divergence
  • Single-determinant Bhattacharyya distance
  • Geometric Jensen-Shannon divergence

Implementation details:

  • Uses the bijector scale matrix and inverse scale operator directly to
    compute cross-factor X = B^{-1} A and symmetric cross-operator M = X - X^{-T},
    avoiding redundant dense Cholesky factorizations.
  • Switches to analytic order-4 Taylor series for log-determinants and trace powers
    below an analytically derived threshold T(dtype) = ((N+1)*eps)^(1/(N+1)) to prevent
    catastrophic cancellation and eigenvalue absorption for near-identical distributions.
  • Parameterized unit tests across float32 and float64 dtypes verifying asymptotic
    convergence and cancellation-free precision.

Benchmarks & Error Comparison (float32 CPU, eps=1e-4 perturbation):

Dim (d) Naive Time Resolvent Time Speedup Naive Value (Error) Resolvent Value (Error)
d=3 557.3 us 541.1 us 1.03x -2.34e-07 (Invalid < 0) 1.22e-08 (< 1e-7 rel)
d=10 562.0 us 549.1 us 1.02x -4.18e-07 (Invalid < 0) 2.10e-07 (< 1e-7 rel)
d=50 890.8 us 792.6 us 1.12x 5.68e-05 (> 3% err) 5.52e-05 (< 1e-7 rel)

- Symmetrized KL (Jeffreys) divergence
- Single-determinant Bhattacharyya distance
- Geometric Jensen-Shannon divergence

Implementation details:
- Uses the bijector scale matrix and inverse scale operator directly to
  compute cross-factor X = B^{-1} A and symmetric cross-operator M = X - X^{-T},
  avoiding redundant dense Cholesky factorizations.
- Switches to analytic order-4 Taylor series for log-determinants and trace powers
  below an analytically derived threshold T(dtype) = ((N+1)*eps)^(1/(N+1)) to prevent
  catastrophic cancellation and eigenvalue absorption for near-identical distributions.
- Parameterized unit tests across float32 and float64 dtypes verifying asymptotic
  convergence and cancellation-free precision.

Benchmarks & Error Comparison (float32 CPU, eps=1e-4 perturbation):
| Dim (d) | Naive Time | Resolvent Time | Speedup | Naive Value (Error)      | Resolvent Value (Error) |
|---------|------------|----------------|---------|--------------------------|-------------------------|
| d=3     | 557.3 us   | 541.1 us       | 1.03x   | -2.34e-07 (Invalid < 0)  | 1.22e-08 (< 1e-7 rel)   |
| d=10    | 562.0 us   | 549.1 us       | 1.02x   | -4.18e-07 (Invalid < 0)  | 2.10e-07 (< 1e-7 rel)   |
| d=50    | 890.8 us   | 792.6 us       | 1.12x   |  5.68e-05 (> 3% err)     | 5.52e-05 (< 1e-7 rel)   |

PiperOrigin-RevId: 979240063
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.

1 participant