jejjohnson / jejjohnson/pipekit

pipekit-train: KL sample-based fallback (v0.1.1)

Open
#20 0 comments 0 reactions 0 assignees View on GitHub
type:feature
Dominant language
Python
Stars
0
Forks
0
PR merge metrics
No merged PRs in 30d

Description

## Problem / Request

The `KL` loss currently supports only the **analytic** path (`predicted.kl_divergence(target)`). The design committed to a sample-based fallback `-(target.log_prob(samples) - predicted.log_prob(samples)).mean()` with samples drawn from `predicted`. Deferred from v0.1 because it needs the per-batch RNG; landed under v0.1.x.

## Motivation

- Many distributions don't expose `.kl_divergence` (custom densities, mixture distributions, normalising flows). The sample-based fallback unblocks them.
- Small, contained change: just needs an RNG threaded through `Loss.__call__`.

## Proposed API

Today's `Loss.__call__(predicted, target)` extends to `Loss.__call__(predicted, target, key=None)`:

```python
class KL(Operator):
n_samples: int = 32 # for the sample-based fallback

def _apply(self, predicted, target, key=None):
if hasattr(predicted, "kl_divergence"):
return predicted.kl_divergence(target).mean(), {"kl": ...}
if key is None:
raise TypeError(
"Sample-based KL requires a PRNG key. The Equinox "
"adapter supplies one via task.loss_fn; for ad-hoc "
"use, pass key=jax.random.key(seed) explicitly."
)
samples = predicted.sample(seed=key, sample_shape=(self.n_samples,))
log_q = predicted.log_prob(samples)
log_p = target.log_prob(samples)
return (log_q - log_p).mean(), {"kl": ...}
```

`_SynthesisedTask.loss_fn` in the Equinox adapter already has `key`; thread it through.

## Implementation Steps

- [ ] Extend `Loss.__call__` signature to accept optional `key` (backwards compatible — default `None`).
- [ ] `KL._apply` implements the sample-based fallback.
- [ ] `_SynthesisedTask.loss_fn` passes `key` to the loss call.
- [ ] Test: KL between two simple distributions (e.g. mocked Normal + Normal) via the analytic path; same pair via the sample-based path; assert the two agree within tolerance.

## Definition of Done

- [ ] Sample-based KL works under the Equinox adapter.
- [ ] Existing analytic path unchanged.
- [ ] Docstring updated; `KL` no longer raises `TypeError` for non-analytic distributions.

## Relationships

- Parent: pipekit-train v0.1.x (small follow-up)

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.