Tensor Channel Equivariant Graph Neural Networks for Molecular Polarizability Prediction

Jean Philip Filling; Daniel Franzen; Michael Wand

cs.LGMay 16, 2026

Tensor Channel Equivariant Graph Neural Networks for Molecular Polarizability Prediction

Authors: Jean Philip Filling, Daniel Franzen, Michael Wand

Organizations: Institute for Computer Science, Johannes Gutenberg University Mainz, Germany

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.

Explore similar work

May 11, 2026cs.LG

QT-Net: Rethinking Evaluation of AI Models in Atomic Chemical Space

Atomic properties such as partial charges or multipoles encode chemically meaningful information that can inform downstream molecular property prediction, but their evaluation as machine learning targets has been complicated by the absence of a principled out-of-distribution evaluation protocol at the atomic level. In this work, we propose a held-out evaluation protocol that clusters atomic environments by SOAP descriptors and computes metrics accounting only for cluster labels unseen during training. Following this procedure, we use 5

\times

5 cross-validation and Tukey's HSD to run a statistically rigorous comparison of E(3)-equivariant against non-equivariant, rotationally augmented models for predicting electron populations and multipoles of H, C, N, and O atoms. Building on our results, we introduce the Quantum Topological Neural Network (QT-Net), a rotationally augmented, non-equivariant graph neural network. We show that QT-Net can be used to infer properties of atoms in molecules from QM9 outside our training set, and that these inferred properties can yield improvement when used as input features for downstream molecular property prediction. To further validate the framework, molecular dipole moments computed from QT-Net's per-atom outputs recover the ground-truth values reported in QM9. We release all code and data, including a JAX implementation of QT-Net, to support the broader use of learned QTA properties as inductive biases for atomic-scale molecular machine learning.

Pablo Martínez Crespo, Stefano Ribes, Martin Rahm +6

Molecular Property Prediction Vast Chemical Space Quantum Neural Networks Artificial Intelligence Models

Ask PDF

May 19, 2026cs.LG

Group-Algebraic Tensors: Provably-optimal Equivariant Learning and Physical Symmetry Discovery

We introduce the

\star_G

tensor algebra, in which any finite group

G

defines the multiplication rule, making equivariance an intrinsic algebraic property rather than an architectural constraint. The framework rests on three machine-verified theoretical pillars: (i)~an Eckart-Young optimality guarantee for the

\star_G

-SVD: the first such result for symmetry-preserving tensor approximation, exact and polynomial-time; (ii)~a Kronecker factorization that composes multiple symmetries by replacing

F_G

with

F_{G_1} \otimes F_{G_2}

with no architectural redesign; and (iii)~~a 600-line Lean~~4 formalization of the

\star_G

algebra. The framework provides capabilities that equivariant neural networks (ENNs) structurally cannot: a closed-form per-irreducible-representation decomposition of every prediction, and data-driven discovery of the symmetry group that best fits a dataset. As a non-trivial empirical demonstration, decomposing QM9 molecular geometry over the chiral octahedral subgroup of SO(3) recovers the Wigner--Eckart selection rules of angular momentum from data alone, with no quantum mechanical input: scalar properties are A

_1

-dominated, dipole components are T

_1

-dominated, the isotropic polarizability is uniquely insensitive to

l\!=\!1

as the rank-2-trace decomposition

l\!=\!0 \oplus l\!=\!2

requires, and the T

_1

_1

predictive-power ratio separates vector observables from scalar observables by a factor of five. On full QM9 (130{,}831 molecules),

\star_G

-SVD with ridge regression provides closed form predictions at

\sim50-90\times

fewer parameters than parameter-matched MLPs. Algebraic equivariance thus complements architectural equivariance not as a faster-better-cheaper alternative but as a different mathematical affordance: provably-optimal symmetry-preserving compression, per-irrep interpretability, and data-driven physical discovery.

Paulina Hoyos, Shashanka Ubaru, Dongsung Huh +5

Tensor Symmetry Exact Equivariance Singular Value Decomposition

Ask PDF

Aug 19, 2026cond-mat.mtrl-sci

The parity gap in crystal tensor prediction

Crystal symmetry dictates whether a physical response tensor must vanish, establishing a direct test for machine learning predictions independent of property calculations. We derive the parity gap, a group-theoretic metric quantifying the piezoelectric tensor freedom permitted by a crystal's proper rotation subgroup

SO(3)

but eliminated by inversion symmetry in

O(3)

. Across state-of-the-art equivariant neural network architectures, unconstrained

SO(3)

models systematically predict forbidden non-zero responses matching the parity gap of each centrosymmetric crystal class, while polar distortion paths dynamically map output responses to the loss of inversion symmetry. Regression controls confirm that enforcing full

O(3)

parity incurs no consistent accuracy cost across predictive tasks. Crucially, while training interventions using explicit zero labels reduce violation magnitudes, they leave residual forbidden outputs. Exact physical compliance instead requires structural enforcement through

O(3)

representation design or explicit output antisymmetrization. The parity gap thus provides a unified framework to distinguish empirical error reduction from exact structural compliance with physical law.

Can Polat, Mustafa Kurban, Erchin Serpedin +1

Symmetry Crystallographic Plane Indices Crystal Piezoelectric

Ask PDF