Skip to content

Preserve concatenated operator higher derivatives with analytic bilinear VJP - #153

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

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

Conversation

@M-Colley

Copy link
Copy Markdown

Concatenated operators currently inherit the generic _bilinear_derivative, which reconstructs detached arguments and evaluates autograd.grad without building a graph. When CatLinearOperator is used inside another operator's matmul (for example BlockDiagLinearOperator), first derivatives can be numerically correct while covariance second derivatives are detached or incomplete.

This appears in independent multitask Gaussian covariances: the squared dense covariance loss for two simple blocks yields correct first gradients with create_graph=True, but those gradients have requires_grad=False, so the next backward raises. In a two-output exact-GP posterior example the resulting mixed Hessian-vector error is 0.213, although values and all first derivatives agree with dense Gaussian conditioning to about 8e-15. Central differences of the unmodified first gradients converge to the dense result. This Cat-specific analytic derivative reduces that full model discrepancy to 2.3e-12.

Delegate the bilinear VJP to each child with the corresponding row, column, or batch vector slice. Broadcast batches and move vectors to each child's device, retaining the incoming vector graphs without materializing the concatenated matrix. The forward implementation is unchanged. This is a separate missing analytic delegation from the triangular operator repair in #152; applying that repair alone does not fix this example.

Validation against main 1c5e24b:

  • New regressions fail in 36 CPU/CUDA subcases before the repair.
  • Torch 2.13 CPU/CUDA: 309 tests and 36 subcases pass, including the existing Cat operator suite.
  • Torch 2.14.1 CPU: 309 tests and 15 subcases pass.
  • Coverage includes row/column/batch concatenation, extra broadcast batches, mixed first/second/third derivatives, gradcheck/gradgradcheck, unbatched vector operands, independent covariance assembly, nested child operators, and mixed CPU/CUDA child/output devices.
  • ufmt check, flake8, and git diff --check pass for the changed files.

AI assistance: implemented and tested with OpenAI Codex. The numerical reproductions, regression runs, and formatting checks 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