Stable logsumexp softmax gradients without rewrites - #2313
Open
ricardoV94 wants to merge 3 commits into
Open
Conversation
Baking an OpFromGraph inner graph disabled recognition rewrites wholesale, to stop a SymbolicOp such as Softmax from re-creating itself inside its own body. That is too coarse: it also stripped recognition from ordinary OpFromGraphs, so a naive exp(x) / sum(exp(x)) that evaluates correctly at the top level returned nan once wrapped in one. Only SymbolicOps can recurse this way, so exclude recognition just for them.
pt.logsumexp built a bare log(sum(exp(x))) and left the max subtraction to local_log_sum_exp. That rewrite runs after grad, so the gradient was built from the raw form as exp(x) / sum(exp(x)), which overflows past log(float64 max) and is 0/0 for very negative x. Issue pymc-devs#2288 only addressed the forward value. Add a LogSumExp SymbolicOp that subtracts the maximum as built, and route logsumexp and logaddexp through it. Differentiating the stabilized body is already stable, so no pullback override is needed. Inlined at specialize by late_inline_OpFromGraph, as for XLogY. local_log_sum_exp and local_log_add_exp now emit that Op rather than expanding the stable form inline, so hand-written and helper-built forms share one representation. They are tagged symbolic_op_recognition, which keeps them from firing inside LogSumExp's own body, and registered at stabilize only: at specialize they would re-wrap what late_inline_OpFromGraph had just inlined.
ricardoV94
force-pushed
the
stable_logsumexp_softmax_grads
branch
from
July 29, 2026 11:23
bbd2bd4 to
26f8fa7
Compare
ricardoV94
force-pushed
the
stable_logsumexp_softmax_grads
branch
from
July 29, 2026 11:25
26f8fa7 to
e88dd42
Compare
ricardoV94
changed the title
Collaborator
|
Is this type of rework similar in philosophy (but different mechanism) to my preliminary Taylor series remainder stuff? There we need to take care because differentiating naturally constructs things that catastrophically cancel, and does so in a sort of worst-case way. So it's a regime where you effectively need to guard against symbolic differentiation in order to protect numerical stability. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Related to pymc-devs/pymc-extras#720
Related to #2289
Gotta bite the bullet and make logsumexp/logaddexp eager SymbolicRVOp so their gradients are stable.
Also figure out a more clean way to avoid recursion in rewrites but still handle OFG(naive_graph)