Superior Molecular Representations from Intermediate Encoder Layers

Luis Pinto

Keywords: Property Prediction, Molecular Encoders, Pretrained Encoders

TL;DR: Information flow analysis shows intermediate layers in molecular encoders are more general than the specialized final-layer. Empirical studies confirm superior results for both frozen and finetuned intermediate embeddings vs. the default final-layer.

Abstract: Pretrained molecular encoders have become indispensable in computational chemistry for tasks such as property prediction and molecular generation. However, the standard practice of relying solely on final-layer embeddings for downstream tasks may discard valuable information. In this work, we first analyze the information flow in five diverse molecular encoders and find that intermediate layers retain more general-purpose features, whereas the final-layer specializes and compresses information. We then perform an empirical layer-wise evaluation across 22 property prediction tasks. We find that using frozen embeddings from optimal intermediate layers improves downstream performance by an average of 5.4%, up to 28.6%, compared to the final-layer. Furthermore, finetuning encoders truncated at intermediate depths achieves even greater average improvements of 8.5%, with increases as high as 40.8%, obtaining new state-of-the-art results on several benchmarks. These findings highlight the importance of exploring the full representational depth of molecular encoders to achieve substantial performance improvements and computational efficiency.

Submission Track: Paper Track (Full Paper)

Submission Category: AI-Guided Design + Automated Material Characterization

Institution Location: Montreal, Canada

AI4Mat RLSF: Yes

Submission Number: 13