[論文レビュー] UNiTE: Unitary N-body Tensor Equivariant Network with Applications to Quantum Chemistry.

Zhuoran Qiao, Anders S. Christensen+4 more|arXiv (Cornell University)|May 31, 2021

Machine Learning in Materials Science参考文献 84被引用数 23

ひとこと要約

UNiTEは、対称テンソルを処理する際に線形時間計算量を達成するユニタリN体テンソル等長性ニューラルネットワークを導入し、量子化学におけるスケーラブルで対称性を保った学習を可能にした。ベンチマークにおいて最先端の手法を110%以上上回り、従来の数値的手法よりも3桁速い、堅牢なゼロショット一般化を実現した。

ABSTRACT

Equivariant neural networks have been successful in incorporating various types of symmetries, but are mostly limited to vector representations of geometric objects. Despite the prevalence of higher-order tensors in various application domains, e.g. in quantum chemistry, equivariant neural networks for general tensors remain unexplored. Previous strategies for learning equivariant functions on tensors mostly rely on expensive tensor factorization which is not scalable when the dimensionality of the problem becomes large. In this work, we propose unitary $N$-body tensor equivariant neural network (UNiTE), an architecture for a general class of symmetric tensors called $N$-body tensors. The proposed neural network is equivariant with respect to the actions of a unitary group, such as the group of 3D rotations. Furthermore, it has a linear time complexity with respect to the number of non-zero elements in the tensor. We also introduce a normalization method, viz., Equivariant Normalization, to improve generalization of the neural network while preserving symmetry. When applied to quantum chemistry, UNiTE outperforms all state-of-the-art machine learning methods of that domain with over 110% average improvements on multiple benchmarks. Finally, we show that UNiTE achieves a robust zero-shot generalization performance on diverse down stream chemistry tasks, while being three orders of magnitude faster than conventional numerical methods with competitive accuracy.

研究の動機と目的

• 量子化学などの分野における高次テンソルに対するスケーラブルな等長性ニューラルネットワークの欠如に対処すること。
• 既存のテンソル因子分解に基づく手法の等長性学習におけるスケーラビリティの制限を克服すること。
• 対称N体テンソルにおけるユニタリ群の等長性（例：3次元回転）を維持するニューラルネットワークアーキテクチャの開発。
• 対称性を保ったままモデルの一般化を向上させる、新しい対称性保存型正規化技術、等長性正規化を導入すること。
• 高い速度と正確性を兼ね備えた、多様な下流化学タスクにおけるゼロショット一般化を可能にすること。

提案手法

• 3次元回転などのユニタリ群作用における等長性を保つように設計された、対称N体テンソルに特化した新しいニューラルネットワークアーキテクチャUNiTEを提案する。
• 非ゼロテンソル要素の数に比例する線形時間計算量を達成する、パラメータ効率の良い設計を採用し、スケーラビリティを実現する。
• 対称性を保持しながらモデルの一般化を向上させる、等長性正規化と呼ばれる正規化手法を導入する。
• 因子分解に高コストがかかるのを避けるために、構造的にテンソル分解の原則を活用し、等長性を維持する。
• N体相互作用におけるメッセージパッシングフレームワークを採用し、特徴変換に学習可能な等長性層を用いる。
• 量子系における多粒子相互作用を符号化する、対称テンソル表現を採用する。

実験結果

リサーチクエスチョン

• RQ1 一般の対称N体テンソルに対して、ユニタリ群作用における等長性を維持するニューラルネットワークアーキテクチャを設計できるか？
• RQ2 非ゼロテンソル要素の数に対して線形時間計算量を達成できるか。これにより、高次元問題へのスケーラビリティが可能になるか？
• RQ3 等長性正規化は、対称性を保持したまま一般化を向上させることができるか？
• RQ4 UNiTEは、量子化学ベンチマークで最先端の性能を達成し、既存手法を顕著に上回る改善を実現できるか？
• RQ5 UNiTEは、多様な化学タスクにおいてゼロショット設定で堅牢に一般化できるか。また、従来の数値的ソルバーをはるかに高速に処理できるか？

主な発見

• UNiTEは、最先端の機械学習手法と比較して、複数の量子化学ベンチマークで平均110%以上の性能向上を達成した。
• モデルは多様な下流化学タスクにおいて、対称性からの強いインダクティブバイアスを示す、堅牢なゼロショット一般化を示した。
• UNiTEは、従来の数値的手法よりも3桁速く動作しながら、競争力のある正確性を維持した。
• アーキテクチャは、非ゼロテンソル要素の数に対して線形時間計算量を維持しており、高次元系へのスケーラビリティを実現した。
• 等長性正規化は、対称性を破らずにモデルの一般化を向上させ、未学習タスクにおけるより良いパフォーマンスに寄与した。
• UNiTEは、従来の等長性ニューラルネットワークが未到達としてきた、高次元テンソル上での等長性関数の学習に成功した。