Foundation Models of Scientific Knowledge for Chemistry: Opportunities, Challenges and Lessons Learned (2022)

Abstract: Foundation models pre-trained on large corpora demonstrate significant gains across many natural language processing tasks and domains e.g., law, healthcare, education, etc. However, only limited efforts have investigated the opportunities and limitations of applying these powerful models to science and security applications. In this work, we develop foundation models of scientific knowledge for chemistry to augment scientists with the advanced ability to perceive and reason at scale previously unimagined. Specifically, we build large-scale (1.47B parameter) general-purpose models for chemistry that can be effectively used to perform a wide range of in-domain and out-of-domain tasks. Evaluating these models in a zero-shot setting, we analyze the effect of model and data scaling, knowledge depth, and temporality on model performance in context of model training efficiency. Our novel findings demonstrate that (1) model size significantly contributes to the task performance when evaluated in a zero-shot setting; (2) data quality (aka diversity) affects model performance more than data quantity; (3) similarly, unlike previous work, temporal order of the documents in the corpus boosts model performance only for specific tasks, e.g., SciQ; and (4) models pre-trained from scratch perform better on in-domain tasks than those tuned from general-purpose models like Open AI’s GPT-2.

## introduction the emergence of foundation models @xcite such as large-scale autoencoding models (e.g., bert @xcite , roberta @xcite ) and autoregressive language models (e.g., gpt-2 @xcite , gpt-3 @xcite , megatron-turing @xcite and @xcite ) as well as multimodal vision and language models, such as flava @xcite and perceiver @xcite , established a paradigm shift in artificial intelligence (ai). these foundation models, also called neural platforms, are built using self-supervised pretraining at scale. they are then able to be easily adapted to a wide range of downstream tasks via transfer learning @xcite and finetuning @xcite . the wide community adoption of foundation models can be explained by their key properties, two of which are emergent behavior and homogenization -which also make foundation models appealing for adaption across science and security domains. emergence, or emergent behavior, reflect new behaviors that a model introduces or is capable of that it was not explicitly trained to perform. homogenization is the consolidation of methods for building machine learning systems across a wide range of tasks. another key advantage of scaling language models is that they perform competitively on language tasks using in-context learning without fine-tuning or gradient updates. thus, in-context learning allows foundation models to be effectively used across new downstream tasks with only simple instructions and a few optional examples. in this work we focus on a science domain (chemistry) and demonstrate the value and limitations of large-scale language models evaluated across a wide range of in-domain (science-focused) and out-of-domain tasks. unlike the majority of work on foundation models that focuses on pretraining these models on book corpora, web pages, wikipedia and mixed sources, e.g., the pile @xcite , we pretrain our models on scientific literature. using scientific literature presents unique opportunities and challenges. opportunities include the scale and diversity of scientific literature, the explicit structure, and explicit alignment across different modalities in the papers, e.g., table and figure references. challenges include limited benchmarks that can be used to perform model evaluation, model prompting and interactions. there are three major contributions of this work: (1) we collect and release a 0.67tb dataset covering research publication data across 10+ sources for chemistry; (2) we release 28 auto-regressive foundation models for chemistry that have been pretrained from scratch; and (3) we present a rigorous evaluation of model performance on 15+ indomain and out-of-domain tasks that investigates the effects of model and data scaling, knowledge depth (aka diversity), and temporal order on performance as described in research questions below. ## related work in this section we summarize previous efforts in two categories: mixed-domain continual pretraining that continues pretraining of a base model on domain data and in-domain pretraining from scratch that pretrains a from scratch on domain data. we present a model summary in ## model pretraining unlike the majority of related models that rely on a base bert (or variant) model, we adapt the open-ai's gpt-2 transformer decoder architecture @xcite to train autoregressive language models for chemistry. to understand the impact of model size (rq2), we experiment with four different transformer sizes: small (s), medium (m), large (l), and extra-large (xl). these models differ in the number of decoder layers, hidden size of the model, and the number of attention heads in transformer blocks as shown in our experiments leverage the gpt-neox python library @xcite developed with megatron @xcite and deep-speed @xcite . we optimize the autoregressive log-likelihood (i.e., cross-entropy loss) averaged over a 2048-token context. we set the micro batch size per gpu as 4, and the learning rate to 2 × 10 -4 , and rely on the cosine decay. we use an adam optimizer with β 1 = 0.9, β 2 = 0.99, and σ = 10 -8 and clip the gradient norm at 1.0. in addition, zero optimizer @xcite was used to reduce memory footprint by distributing optimizer states across several processes. to reduce memory and increase training throughput, we use mixed-precision training @xcite and the parallel attention and feed-forward implementations available in gpt-neox @xcite . we also use the rotary positional embeddings @xcite instead of the learned positional embeddings used in the gpt-2 model (radford et al., 2019) because they offer performance advantages in tasks with longer texts by capturing relative position dependency in self-attention. our models are pretrained across multiple workers with data parallelism. as the largest model in our experiments fit on a single gpu, we didn't use the model (tensor) or pipeline parallelism. models are pretrained from scratch for a total of 320k steps. the original gpt-2 models are fine-tuned for 150k steps. we perform experiments in a single dgx-a100 machine with 8 80gb gpus. ## data collection and processing we collected a large corpus of 53.45 million chemistry-focused scientific articles and abstracts, resulting in 670gb of text data. as shown in tokenization as used in gpt-2 model, we use a byte pair encoding (bpe) tokenizer. we train bpe tokenizers for each data sample with a vocabulary size of 64k as preliminary experiments varying vocabulary sizes from 64k to 256k for smaller scale model pretraining did not show significant differences in performance. we compare the gpt-2 vocabulary generated from the webtext and the in-domain vocabularies generated from our corpora and find that the in-domain vocabulary breaks chemical entities into fewer tokens. for example, dimethylnitroxide was tokenized into #dimethyl, #nitr, #oxide using the in-domain vocabulary and #dim, #ethyl, #nit, #rox, #ide using the gpt-2 vocabulary. ## analysis and results this section presents the analysis of 28 pretrained models evaluated on 15+ in-domain and out-ofdomain downstream tasks (rq1, section 5.1). we investigate the effects of model and data scaling (rq2, section 5.2), knowledge diversity (rq3, section 5.3), and temporal order (rq4, section 5.4) on the downstream performance. we also compare the results from continual vs. from scratch pretraining (section 5.5) and present the analysis of large-scale training efficiency (section 5.6). baseline models as we use a similar model architecture, we identify open ai's gpt-2 @xcite as a baseline comparison model. we compare our performance with four variants of the original gpt-2 models, corresponding to small (s), medium (m), large (l), and extra-large (xl) sized transformer architectures shown in our models we pretrained models with individual datasets (aminer, core, mag, pubmed, s2orc, wos) and combined abstracts and fulltexts. our goal is to systematically study data biases in the model performance when pretraining models with individual datasets. for example, pubmed publications cover mostly bio-medicinal terms @xcite , while the majority of s2orc publications are from medicine, biology, physics, and mathematics @xcite . we only use 4 gpus for the models pretrained with individual datasets and 8 gpus for the rest. this is to control the number of tokens seen during model pretraining (320,000 steps * 4 gpus * 4 micro batch size * 2,048 context @xmath0. we also trained one xl (4x) model with 4x larger batch size than what used in xl model to evaluate the impact of the number of training tokens. ## conclusions in this paper, we collected and released 0.67tb of research publication data collected across 10+ sources for chemistry. we pretrained and released 25+ foundation models for chemistry. we rigorously analyzed model performance on 15+ indomain and out-of-domain tasks.