Skip to content

Preserve higher derivatives in fallback bilinear autograd - #155

Open
M-Colley wants to merge 1 commit into
cornellius-gp:mainfrom
M-Colley:codex/fix-bilinear-fallback-higher-derivatives
Open

M-Colley wants to merge 1 commit into
cornellius-gp:mainfrom
M-Colley:codex/fix-bilinear-fallback-higher-derivatives

Conversation

@M-Colley

Copy link
Copy Markdown

The default LinearOperator._bilinear_derivative detaches every representation tensor and calls autograd.grad without create_graph. Operators that inherit this fallback consequently lose representation-input higher derivatives. For example, kron(Diag(ones(4)), Root(a) + Diag(0.06 * ones(2))).to_dense().sin().sum() has the correct first derivative with respect to a, but that derivative is detached despite create_graph=True. The equivalent dense torch.kron expression has a nonzero Hessian. Correlated multitask likelihood covariance materialization is one affected use case.

This repair retains the existing detached ordinary-backward path. When the caller requests a derivative graph, it builds independent view nodes for differentiable representation tensors and requests graph construction from the nested autograd call. The views are important: simply reusing the original tensors would include unwanted paths through dependent left/right vectors or through another representation tensor when computing a partial derivative. The resulting partial derivatives remain connected to all original tensors for subsequent backward passes.

Tests include:

  • A nonlinear fallback operator with dependent representations and vectors, checked against explicit analytic partials and mixed higher derivatives.
  • Kronecker Root-plus-Diagonal covariance compositions against independent dense algebra through third derivatives, including batches/broadcasting and CPU/CUDA.
  • Gradcheck/gradgradcheck and an ordinary-backward control.

The unmodified source fails eight new CPU/CUDA subcases. The repaired representative operator suites (Kronecker, concatenation, triangular, added-diagonal, sum and the new fallback regressions) pass 968 tests and ten subcases on Torch 2.13 CPU/CUDA, and 968 tests and five subcases on Torch 2.14.1 CPU. Formatting, Flake8 and git diff --check pass.

Related analytic specializations are proposed separately in #152 and #153. This change fixes the remaining general fallback and does not depend on either PR. Prepared with assistance from OpenAI Codex; the source reproductions and validation described above were executed locally.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant