TriCL

Triangular Contrastive Learning on Molecular Graphs

MinGyu Choi, Wonseok Shin, Yijingxiu Lu, and Sun Kim - ArXiv Link (2205.13279)


The markdown contents of this post is mainly refer to Gratus907’s blog, one of the co-author of the paper.

TriCL is a tri-modal contrastive learning framework with triangular loss function that learns the angular geometry of the embedding space through contrasting the area of positive and negative triplets.

We achieved SOTA performance on MoleculeNet classification tasks, the molecular property prediction dataset.


Alignment and Uniformity

Alignment and Uniformity suggested by Wang and Isola is a assessment metric of contrastive learning embedding space.

  • Alignment : Similar inputs should be mapped closely (augmented samples in unimodal, and different encoders in multimodal)
  • Uniformity : Dissimilar inputs should be mapped away from each other, so that the embedding could use the whole space.

Perfect alignment - all representations at one point - is useless, since the model cannot distinguish distinct features among samples.
Perfect uniformity - all representations fall apart - is also useless, since the model cannot learn useful correlationship from training data.

Therefore a careful control between alignment and uniformity is the key for effective pre-training.


Triangular Contrastive Learning Framework

Figure 1. TriCL Framework.

TriCL utilizes both intra-modal and inter-modal contrastive losses to learn the geometry of embedding space.

ANY neural network encoders could be adopted under the TriCL framework.

For intra-modal contrastive loss, we used NT-Xent loss function.

Note that NT-Xent loss also could be interpreted as a summation of alignment and uniformity.

NT-Xent Loss, decomposed to alignment and uniformity.

We proposed Triangular Area Loss, which is designed to pull positive triplets and push negative triplets.

Specifically, positive triplet means three representation vectors from three independent encoders; negative triplet means otherwise.

Since the geometric angle information among three vectors could be regarded through TAL, the model could better learn the geometry among independent encoders.

Triangular Area Loss, decomposed to alignment and uniformity.

Results

We achieved SOTA performance on MoleculeNet classification tasks, after pre-training five layers of GIN benchmark model with GEOM 50,000 molecules.

Performances on MoleculeNet Classification Tasks.

For more description, please refer to our paper.