Densitymodel adtype slots - #62
Conversation
Per review: instead of reading pushforward_performance directly, follow DI's own hvp_mode composition — ForwardOverAnything → ForwardOnGrad, everything else → ReverseOnGrad. Only the outer direction matters since we differentiate the already-built gradlogp.
… densitymodel-adtype-slots
|
@penelopeysm can i get your eyes on this? It is basically your ideas in #52 |
when our robot overlords revolt, I'm taking out github copilot first |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #62 +/- ##
==========================================
+ Coverage 89.99% 90.86% +0.87%
==========================================
Files 7 7
Lines 1069 1150 +81
==========================================
+ Hits 962 1045 +83
+ Misses 107 105 -2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
penelopeysm
left a comment
penelopeysm
left a comment
There was a problem hiding this comment.
The code all looks sensible to me & in line with what I had imagined, although I didn't look tooooo hard
|
|
||
| Backends are turned into prepared [DifferentiationInterface](https://github.com/JuliaDiff/DifferentiationInterface.jl) callables when sampling starts, and that preparation is reused for the rest of the run. Hand-written and AD-derived slots mix, so an analytical gradient with `hvp=AutoForwardDiff()` is fine. | ||
|
|
||
| A backend in `hvp` differentiates whatever the gradient slot holds; it is not a second derivative of `logdensity`. Over an AD-derived gradient that composition is second-order AD, and over a hand-written one it is a single AD pass across your own code. The same goes for the batched pair, and a `logdensity_batch` supplied without a `grad_logdensity_batch` has the batched gradient derived for it — one gradient of `sum(logdensity_batch(X))`, which is the stacked per-column gradients only because columns are independent, so `logdensity_batch` must not couple them. |
There was a problem hiding this comment.
it is not a second derivative of
logdensity. Over an AD-derived gradient that composition is second-order AD
This is contradictory or at best unclear. I think the second half of the first sentence should be cut. But honestly maybe it's just easiest to have a table liks this:
grad hvp result
adtype adtype hvp calculated via second order AD
...
There was a problem hiding this comment.
Oh i see you already did this in #63 haha, so you can ignore me
|
okay well that's the first time i've ever accidentally merged a PR....tried to merge in the docs refresh pr but i guess i wasn't looking |
|
bit confusing as i have branch protections too. I am able to bypass but i guess since i did this from #63 it just didn't even flag? odd / not optimal |
| that a wrapped backend still reaches its own normalization: dispatch happens | ||
| on what comes out of `DI.outer`, not on the `SecondOrder` around it. | ||
| =# | ||
| _hvp_forward_backend(backend::AbstractADType) = DI.outer(backend) |
There was a problem hiding this comment.
Technically that's only true if the backend is forward itself, or forward-over-reverse
| model-specific Hessian-vector product used by DEER to avoid AD through | ||
| problematic kernels when available. If omitted, the sampler differentiates | ||
| `logdensity` directly via DifferentiationInterface. | ||
| - `grad_logdensity` — callable `x -> AbstractVector`, or a backend applied to |
There was a problem hiding this comment.
"applied" sounds weird here
| holds; it does not take a second derivative of `logdensity`. Over an | ||
| AD-derived gradient that composition is second-order AD, and over a | ||
| hand-written gradient it is a single AD pass across your own code. Only the | ||
| outer half of a `DifferentiationInterface.SecondOrder` is used, the gradient |
There was a problem hiding this comment.
So do you interpret every SecondOrder as forward-over-reverse, even if it's not?
| made for; anything else (e.g. the `Dual`s an outer AD pass pushes through | ||
| this gradient when forming an HVP) goes through unprepared `DI.gradient`. |
There was a problem hiding this comment.
This could be an alternative mode for DI, between strict and lax: use preparation when the signatures match, otherwise go from scratch
Before this, you had to write a gradient
by hand for every model, and pass an AD backend to the sampler even when nothing
ever called it. Now you can hand
DensityModela log-density plus a backend andit works out the derivatives it needs.
You can build a model from just a log-density:
DensityModel(logp, AutoForwardDiff(), dim).Any of the four derivative slots (
grad_logdensity,hvp,grad_logdensity_batch,hvp_batch) takes an AD backend in place of a function you wrote yourself, and you canmix the two freely.
backend=onParallelMALASampleris now optional. A model that brings its ownhvp/hvp_batchdoes not need it. It used to be required by the API even when itwas never invoked.
AD gets prepared once per chain instead of once per step. The prepared model travels
in the sampler state, and a state handed to a different model via
initial_stateisre-prepared, so the model you pass to
sampleis always the one sampled.When your gradient is itself AD-derived, HVPs now go through DI's second-order
operator instead of one AD call nested inside another. Same numbers, roughly 10x
fewer allocations, and it is what asking for second-order AD should have meant.
You can pass a
SecondOrderand both halves get used. It previously kept the outerpass and threw the inner one away without saying so. This is also the only way to get
an AD HVP for a Turing or LogDensityProblems model, whose gradient arrives already
prepared and cannot be differentiated a second time.
Turing and LogDensityProblems models can reach the batched DEER path now, since both
extension constructors forward the batched slots. Neither package gives you a batched
log-density, so
logdensity_batchstill has to be written by hand.Slots that cannot be used now say so instead of being quietly dropped. Both batched
derivative slots require
logdensity_batch, and anhvp_batchwith no batchedgradient to pair with raises at sampling time.
This will resolve Allow supplying only the log posterior #40 and resolve Retool AD specification #52. Sort of addresses Use second order AD directly #37 (GPU still flakey, will be addressed hopefully soon after this)