NNPOps.neighbors.getNeighborPairs does not seem to allow double backward

[JMorado]

Hi,

This is more of a discussion than an issue specifically relatd to smee, but it appears that when using NNPOps.neighbors.getNeighborPairs with torch.autograd.grad(..., create_graph=True), the resulting gradient tensor has grad_fn=None, preventing second-order gradients from being computed. This is a blocker for anyone aiming to fit forces in periodic systems. The following code should reproduce the issue I have encountered:

import NNPOps.neighbors
import torch

epsilon = torch.tensor([0.3], requires_grad=True, device="cuda")
sigma = torch.tensor([3.0], requires_grad=True, device="cuda")

# Two particles
coords = torch.tensor(
    [[0.0, 0.0, 0.0], [3.1, 0.0, 0.0]], requires_grad=True, device="cuda"
)
box = torch.eye(3, device="cuda") * 10.0

# Get distance via NNPOps
_, _, distances, _ = NNPOps.neighbors.getNeighborPairs(coords, 5.0, -1, box)

# Compute LJ energy
sig_r = sigma / distances
energy = 4.0 * epsilon * (sig_r**12 - sig_r**6)

print(f"Distance: {distances[0].item():.4f}")
print(f"LJ Energy: {energy.item():.4f}")

# Compute forces: F = -dE/dcoords
forces = -torch.autograd.grad(energy, coords, create_graph=True, retain_graph=True)[0]

print(f"Forces: {forces}")
print(f"forces.grad_fn: {forces.grad_fn}")

# This requires d(forces)/d(epsilon) and d(forces)/d(sigma)
force_loss = (forces**2).sum()
grad_epsilon = torch.autograd.grad(force_loss, epsilon, create_graph=True)[0]
print(f"d(force_loss)/d(epsilon): {grad_epsilon}")

I think this has already been mentioned in relation to #143 and #141, but since getNeighborPairs does not seem to allow double backward, would it make sense to rewrite it using PyTorch operations to enable higher-order gradients?

Many thanks in advance for any ideas on this.

[JMorado]

Interestingly enough, this seems to be device dependent, i.e. the above snippets work when device="cpu" but not when device="cuda".

[lilyminium]

Thanks for raising this @JMorado! Yeah I agree that re-implementing that in Pytorch would be useful. It looks like the cpu version implements the function entirely with torch operations, which might be why it works on cpu.

[JMorado]

Happy to make a PR -- I've got a working branch that implements a PyTorch-native function that mimics NNPOps.neighbors.getNeighborPairs behaviour.

[JMorado]

I was benchmarking my implementation, and unfortunately a PyTorch-native implementation seems to be ca. 3 times slower. However, doing something like the following is equally performant:

def _compute_pairwise_periodic(
    conformer: torch.Tensor, box_vectors: torch.Tensor, cutoff: torch.Tensor
) -> PairwiseDistances:
    import NNPOps.neighbors

    assert box_vectors is not None, "box vectors must be specified for PBC."
    assert len(conformer.shape) == 2, "the conformer must not have a batch dimension."

    (
        pair_idxs,
        _,
        _,
        _,
    ) = NNPOps.neighbors.getNeighborPairs(conformer, cutoff.item(), -1, box_vectors)

    are_interacting = ~torch.isnan(distances)
    pair_idxs = pair_idxs[:, are_interacting]

    # ensure i < j
    pair_idxs, _ = pair_idxs.sort(dim=0)

    deltas = conformer[pair_idxs[0]] - conformer[pair_idxs[1]]
    shifts = torch.round(deltas @ torch.linalg.inv(box_vectors))
    deltas = deltas - shifts @ box_vectors

    distances = torch.linalg.norm(deltas, dim=-1)

    return PairwiseDistances(pair_idxs.T.contiguous(), deltas, distances, cutoff)

Essentially we grab the pair_idxs inside the cutoff and only compute the deltas and distances for those.

EDIT: This comment illustrates the good scaling behaviour of NNPOps.neighbors.getNeighborPairs mir-group/nequip#288 (comment)

[lilyminium]

Hmm -- so recompute deltas and distances ourselves? (by the way, I think are_interacting depends on the distances returned by getNeighborPairs so that will need to be retained still). I might need to think about this a bit -- would you be up for opening a PR and seeing if anything breaks in CI with this change?

[lilyminium]

Thanks for looking into this so carefully!

[JMorado]

Just for the record, this is the performance I get on CUDA for the three implementations:

I'll open a PR with the NNPOps + PyTorch distances one, which corresponds to the implementation I pasted above. Thanks!

[JMorado]

Closed via #151.