Implement PaiNN (Polarizable Atomistic Interaction Neural Network)

[louischereau]

🚀 The feature, motivation and pitch

PaiNN (Polarizable Atomistic Interaction Neural Network) is a staple for molecular property prediction and MD simulations, offering a superior speed-accuracy trade-off compared to SchNet or DimeNet++. TI is often used as an indsutry standard for benchmarking. While implementations exist in specialist libraries (TorchMD-NET, SchNetPack), a native PyG implementation would allow users to leverage PyG's native batching, torch.compile optimizations, and broader GNN utility suite.

Read paper here: https://arxiv.org/pdf/2102.03150

Proposed Implementation:

I have a functional implementation and would like to PR it to torch_geometric.nn.models. Key technical details I plan to include:

• Equivariant Message Passing: Refactored to inherit from MessagePassing for scalar and vector feature streams.

• Performance: Optimized radial basis computations and distance-dependent filters to minimize redundant calculations during inference.

• Verification: Includes unit tests for rotation equivariance and translation invariance.

• Compatibility: Designed to be compatible with torch.compile and JIT for high-throughput production environments.

Are there specific naming conventions for vector-valued features or preferences for handling directional unit vectors ($r_{ij} / ||r_{ij}||$) that I should align with?

Alternatives

No response

Additional context

No response

[Ankita7033]

Hi @louischereau and @rusty1s! 👋

About me: I'm a CS student with experience in Graph Neural Networks and PyTorch implementations, and I'm looking to make my first substantial contribution to PyG by helping implement PaiNN.

Why I'm interested: PaiNN is a critical architecture for molecular property prediction—the paper shows significant improvements over SchNet and DimeNet on benchmarks. Having a native PyG implementation would unlock torch.compile optimizations and broader GNN utility.

Proposed Technical Approach:

1. Core Architecture (PaiNN class):

   • Inherit from torch.nn.Module (standard PyG pattern)
   • Implement equivariant message passing by subclassing MessagePassing
   • Key insight: PaiNN maintains both scalar (s) and vector (v) features, unlike SchNet

2. Message Passing Layer (PaiNNConv):

   • Scalar-to-scalar updates via MLP (standard)
   • Scalar-to-vector updates via radial basis functions (RBF) + directional components
   • Vector-to-scalar via norm pooling (preserves equivariance)
   • Verify: All operations must be rotation/translation equivariant

3. Radial Basis Function (RBF) Module:

   • Implement Gaussian RBF with learnable centers/widths (as in paper)
   • Optimize: Pre-compute basis evaluations to avoid redundant calculations
   • Cache distance matrices in forward pass

4. Implementation Structure:

   torch_geometric/nn/models/painn.py
   ├── PaiNN (main model)
   ├── PaiNNConv (equivariant convolution)
   ├── RadialBasis (RBF implementation)
   └── VectorFeatureNorm (vector normalization for equivariance)
   

5. Testing Strategy:

   • Unit tests for rotation equivariance (rotate input → rotated output)
   • Unit tests for translation invariance (shift coords → same result)
   • Benchmark against paper's published numbers on QM9/MD17
   • Integration test with PyG's data loaders

6. Design Questions for maintainers:

   • Should vector features use naming like x_v / x_s, or a dict structure like {'scalar': ..., 'vector': ...}?
   • For torch.compile compatibility, should we avoid dynamic shapes? Any constraints?
   • Where should we store the RBF module—as a separate layer or embed in PaiNNConv?

Next Steps (in order):

1. Get maintainer feedback on design above
2. Implement PaiNNConv with rotation/translation checks
3. Implement full PaiNN model + RBF module
4. Add comprehensive tests
5. Benchmark and optimize
6. Open PR for review

@louischereau—I see you have a functional implementation ready. Would you be open to collaborating? I'm happy to help refactor your code into PyG conventions, add tests, or handle optimization. Or if you prefer to lead this, I'm excited to review and contribute ideas.

Looking forward to building this! 👍