Don't carry the tangent of a clipped step into the next steps - #773
Open
etiferrier wants to merge 2 commits into
Open
etiferrier wants to merge 2 commits into
etiferrier wants to merge 2 commits into
Conversation
With traced jump_ts, a step clipped to just before a jump is often very short while the tangent of its length, d(jump) - d(t0), is O(1). The inner controller proposes the next step as a multiple of that length, so its tangent scaled the tangents of all the following step times by (step size / clipped step size), jump after jump: the derivative with respect to several close jump times was wrong by many orders of magnitude, in forward and reverse mode (and worse when step_ts are also present). After a step that lands on a jump_ts or step_ts time, keep the proposed step's value but give it the tangent of the new t0, so the following step times move rigidly with the time just landed on. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Without the step_ts branch of the fix, the derivative is still off by 1e2-3e10 when step_ts are present; the new parametrization catches it. The solver tolerance is tightened to 1e-10 so that the step_ts variant separates clearly (with the fix: <1e-6; without: >1e2). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
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
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.
Fixes #772.
Problem. With traced
jump_ts, a step clipped to just before a jump is often tiny, but the tangent of its length,d(jump) - d(t0), is O(1).PIDControllerproposes the next step asprev_dt * factor, so that tangent is scaled up into every later step time, and it compounds jump after jump. With several close jump times, the derivative with respect to them ends up wrong by many orders of magnitude, in forward and reverse mode.Fix. In
ClipStepSizeController.adapt_step_size, right after a step lands on ajump_tsorstep_tstime, the proposed next step keeps its value but takes the tangent of the newt0(_drop_step_tangent). The following step times then move rigidly with the time just landed on.stop_gradient(next_t1) + (next_t0 - stop_gradient(next_t0)).test_grad_of_discontinuous_forcingrelies on. That's why this isn'tprev_dt = lax.stop_gradient(t1 - t0)inPIDController: that also fixes Gradients with respect to traced jump_ts blow up: the tangent of a clipped step grows through the next steps #772, but breaks that test.Test.
test_grad_wrt_close_jump_tshas 7 pulses 0.04 apart, with all 14 edges traced, and runs with and without astep_tsgrid. It checks against central finite differences, forRecursiveCheckpointAdjoint,DirectAdjointandForwardMode.step_tsstep_tsmainThe rest of the test suite behaves as on
main.Downstream check. I ran dynamiqs, which passes pulse edges as
jump_ts, on 8 production-like quantum-device workloads: