Stable logsumexp softmax gradients without rewrites - #2313
Conversation
26f8fa7 to
e88dd42
Compare
ricardoV94
changed the title
|
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. |
No, this is just applying the commonly known subtract by max trick consistently. So if user calls |
4283614 to
6e332f1
Compare
|
Failing tests are the known llvm<22 bugs, hopefully they will disappear after #2273 |
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.
6e332f1 to
0b9a773
Compare
2 participants
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)