Forced Gauss-Newton step for last iterations of truncated backward.

[luisenp]

Motivation and Context

Our current derivation of implicit/truncated backward modes hols only for Gauss-Newton. However, regardless of the optimization method used, we can first find the fixed point, then force the optimizer to apply a GN step. Preliminary experiments with Levenberg-Marquardt show that this improves the error in gradient computation, compared with numerical derivatives.

Marking as bugfix, since this means that previous version of LM was incorrect.

How Has This Been Tested

Types of changes

• Docs change / refactoring / dependency upgrade
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)

Checklist

• My code follows the code style of this project.
• My change requires a change to the documentation.
• I have updated the documentation accordingly.
• I have read the CONTRIBUTING document.
• I have completed my CLA (see CONTRIBUTING)
• I have added tests to cover my changes.
• All new and existing tests passed.

[mhmukadam]

@luisenp, should we merge this now (after comments are resolved) given that we are not going with the full Hessian route?

[luisenp]

@mhmukadam Yes, but I'm planning to make a few changes before merging. In particular, I'm not sure we actually need this for truncated grad, in the sense that (AFAIK) there is no theory specifying that the grad enabled loops need to have some specific form. @bamos ?