Tensor Channel Equivariant Graph Neural Networks for Molecular Polarizability Prediction

2026-05-16

AI Summary

This work addresses the challenge of efficiently and accurately predicting molecular polarizability tensors, with a particular focus on their anisotropic components. The authors propose an equivariant graph neural network based on the PaiNN architecture that innovatively incorporates explicit second-order tensor channels throughout the message-passing process. By integrating geometry-driven tensor bases with traceless parameterization, the model achieves end-to-end alignment with the target tensor structure. In contrast to existing approaches that construct tensors only at the readout stage, the proposed method significantly reduces prediction errors for both the full polarizability tensor and its anisotropic part on the QM7-X dataset, while also offering faster inference. Notably, this performance gain stems from architectural design rather than increased model size.

Abstract

We introduce a tensor-channel equivariant graph neural network for direct prediction of molecular polarizability tensors. Building on the efficient PaiNN architecture, we augment the hidden representation with explicit symmetric rank-2 tensor channels aligned with the decomposition of polarizability into isotropic and anisotropic components. In contrast to approaches that construct tensor outputs only at readout, our model propagates tensor structure throughout message passing using geometrically motivated tensor bases. This yields a target-aligned architecture for tensor-valued molecular prediction. On optimized QM7-X geometries, the proposed model achieves lower full-tensor and anisotropic error than both a PaiNN-style readout baseline and a dielectric MACE baseline under matched training conditions and at nearly identical parameter count. In this controlled setting, it also outperforms MACE while remaining substantially faster at inference. Ablation studies show that the gain does not arise from increased capacity alone, but from the combination of explicit tensor propagation and a traceless target parameterization matched to the anisotropic part of the polarizability tensor. Among the tensor bases considered, the strongest results are obtained from interactions between learned directional features, indicating that these are particularly effective for modeling molecular polarizability. Rotational equivariance tests further confirm that all compared models are numerically equivariant, so the observed improvements are attributable to better learning of the target tensor itself. Overall, our results show that for structured tensor-valued targets, propagating target-aligned tensor features can outperform both readout-only tensor construction and a more general higher-order equivariant model in the present training setting.

Problem

• molecular polarizability
• tensor prediction
• graph neural networks
• equivariance
• anisotropic components

Innovation

• tensor-channel equivariance
• molecular polarizability prediction
• equivariant graph neural networks
• tensor propagation
• anisotropic tensor modeling