Triplet Interaction Improves Graph Transformers: Accurate Molecular Graph Learning with Triplet Graph Transformers

7 Feb 2024  ·  Md Shamim Hussain, Mohammed J. Zaki, Dharmashankar Subramanian ·

Graph transformers typically lack third-order interactions, limiting their geometric understanding which is crucial for tasks like molecular geometry prediction. We propose the Triplet Graph Transformer (TGT) that enables direct communication between pairs within a 3-tuple of nodes via novel triplet attention and aggregation mechanisms. TGT is applied to molecular property prediction by first predicting interatomic distances from 2D graphs and then using these distances for downstream tasks. A novel three-stage training procedure and stochastic inference further improve training efficiency and model performance. Our model achieves new state-of-the-art (SOTA) results on open challenge benchmarks PCQM4Mv2 and OC20 IS2RE. We also obtain SOTA results on QM9, MOLPCBA, and LIT-PCBA molecular property prediction benchmarks via transfer learning. We also demonstrate the generality of TGT with SOTA results on the traveling salesman problem (TSP).

PDF Abstract

Results from the Paper


Task Dataset Model Metric Name Metric Value Global Rank Uses Extra
Training Data
Benchmark
Drug Discovery LIT-PCBA(ALDH1) EGT+TGT-At-DP AUC 0.806 # 1
Drug Discovery LIT-PCBA(KAT2A) EGT+TGT-At-DP AUC 0.746 # 1
Drug Discovery LIT-PCBA(MAPK1) EGT+TGT-At-DP AUC 0.743 # 1
Initial Structure to Relaxed Energy (IS2RE), Direct OC20 TGT-At (Mean Aggregation) Validation Mean Energy MAE 0.4022 # 1
Test Mean Energy MAE 0.4131 # 1
Test Mean EWT (%) 8.00 # 3
Initial Structure to Relaxed Energy (IS2RE), Direct OC20 TGT-At (Median Aggregation) Validation Mean Energy MAE 0.4030 # 2
Test Mean Energy MAE 0.4147 # 3
Test Mean EWT (%) 8.30 # 1
Graph Property Prediction ogbg-molpcba TGT-Ag+TGT-At-DP Test AP 0.3167 ± 0.0031 # 3
Number of params 47000000 # 4
Ext. data Yes # 1
Graph Regression PCQM4Mv2-LSC TGT-At Validation MAE 0.0671 # 1
Test MAE 0.0683 # 1
Link Prediction TSP/HCP Benchmark set TGT-Agx4 F1 0.871 # 1

Methods