Reprogramming Language Models for Molecular Representation Learning

Adversarial attacks on ML models in discrete input spaces have been found to cause models to misclassify input samples by manipulating tokens in sequence data [7]. Adversarial Reprogramming (AR) is a technique [3] that is able to repurpose a model (without significant changes to the architecture and parameters) by training an adversary to optimally transform input data such that we can choose the output of the model. This method has been proven to be successful in both, white-box and black-box settings [8]. The depth and the size of large general language models make them ideal candidates to be adversarially reprogrammed, as modifying the internal architecture or finetuning over 1 billion parameters is infeasible for a traditional transfer learning approach. The work in [5] assumes an exposed model, where they use the learned parameters of the pre-trained source model to generate context-based vocabulary map. However, the hypothesis of comparable vocabulary size does not hold when the source data are English vocabularies (on the order of $\sim$ 10 million) and the target data are amino acids (on the order of $\sim$ 20). To that end, we introduce dictionary learning to train our AP.

Typical amino acid/nucleotide representations of biological molecules or widely used SMILES representation of chemical molecules have fewer distinct tokens than that of language data. The significant reduction in dimensionality of the embedded space of English data to molecule data makes finding a mapping non-trivial. Results in [9] demonstrate that that representation learning algorithms have an advantage in transfer learning methods as they capture features relevant to alternate tasks. We require a mapping with high representational capacity. Work in [10] demonstrates that a sparse encoding is distributed and highly expressive: we can represent $O(2^{k})$ input regions with only $O(N)$ parameters. Distributed representations can be clustered to extract relevant features where component extraction algorithms can find the optimal representation (a dictionary). We use a k-SVD approximation algorithm [11] to mitigate computational expenses.

We are given a pretrained classifier, $X$, a source-task dataset $\{V_{S}\}^{n}_{i=1}$ and target-task dataset $\{V_{T}\}^{m}_{j=1}$. The embedded matrices are $V_{T}$ and $V_{S}$ respectively. We can encode an output label mapping function $h:l_{V_{S}}\longmapsto l_{V_{T}}$. We then aim to train an adversarial program (AP) $f_{\theta}$ that finds the optimal coefficients $\theta$ of our atoms in $V_{T}$ to represent a sparse encoding of the dictionary, $V_{S}$ such that $V_{T}=V_{S}\theta$. The AP $f_{\theta}$ is used to perform the target task through the transformation $h(X(f_{\theta}(t_{i}))$ where $t_{i}$ is a molecule data sample. While we do not make any modification the parameters or architecture of $X$, we assume access to the gradient $\nabla_{\theta}X$ for loss evaluation during training.

To reprogram the pretrained classifier, we use a similar structure as introduced in [5], $f_{\theta}:t_{i}\longmapsto s_{i}$ where $t_{i}\in V_{T}$ and $s_{i}\in V_{S}$. Dimension of the input space of $V_{S}$ and $V_{T}$ is $|V_{S}|$, and $|V_{T}|$ respectively, where $|V_{T}|\ll|V_{S}|$. The AP $f_{\theta}$ is parametrized by $\theta\in\mathbf{R}^{|V_{S}|\times|V_{T}|}$, which represents the coefficients of the atoms in $V_{T}$ such that $V_{T}=V_{S}\theta$ and $V_{S}$ is our dictionary. The observation of $|V_{T}|\ll|V_{S}|$ requires that our mapping has high representational capacity, so we encode a sparse representation of $V_{S}$, to extract relevant features from the embeddings of the source-model vocabulary $\{V_{S}\}^{n}_{i=1}$ for the alternate task. To that end, approximate the dictionary, we use a k-SVD solver to optimize over the cross entropy loss for updates to $\theta$.

We define $\theta\in\mathbf{R}^{|V_{S}|\times|V_{T}|}$ and use $\theta_{t}$ to denote its $t$-th column, where a signal (embedding vector) $v_{t}$, can be represented as a sparse linear combination of the signal atoms of columns of $V_{S}$, $v_{t}=V_{S}\theta_{t}$. $v_{t}$ is the representation of the AMP input sample in the dictionary space and satisfies $||v_{t}-V_{s}\theta_{t}||_{p}\leq\epsilon$. An exact solution $\theta_{t}^{*}$ such that $v_{t}=V_{S}\theta_{t}^{*}$ is computationally expensive to find, and is subject to various convergence traps, so for the purpose of our efficient AP approach we approximate $v_{t}\approx V_{s}\theta_{t}$ by limiting the number of iterations to converge to a solution for $v_{t}$. Our optimization problem for finding the optimal sparse representation of input data samples in $V_{T}$ is then minimize $||\theta||_{0}$ subject to $||v_{t}-V_{s}\theta_{t}||_{1}\leq\epsilon$ to enforce a sparse solution [11], where $||\cdot||_{0}$ denotes the $l_{0}$ norm. While algorithms exist to choose an optimal dictionary (an exact solution to k-SVD) that can be continually updated [11], we penalize computational expense over performance for the purpose of maintaining an efficient solution (at the cost of statistically insignificant improvements in accuracy) by using a predetermined number of steps for an approximate solution, $V_{S}$ that encodes the atoms of $V_{T}$. $\theta$ is then used to evaluate the cross entropy loss on $h(X(f_{\theta}(V_{T}))$, which will be updated in the AP $f_{\theta}$.

To benchmark the performance of R2DL we compare it with the current established benchmark of a trained classifier using the same AMP training data set [12].

To further investigate the efficacy of the transfer learning approach, we compare the performance of R2DL versus the model trained from scratch with AMP data, with a restricted training data set. The test accuracies indicate that R2DL performs better when fewer labeled training data samples are available. Tables 2 and 3 show that when trained in a setting with fewer labeled data samples, R2DL outperforms the train from scratch method after the threshold of approximately 5000 samples. Below 5000 samples, both methods approximate random prediction, with R2DL not successfully transferring any learned representations ¹¹15000 AMP training samples and below were excluded from the results below as they showed statistically insignificant test accuracy..

BERT and GPT-3 are 2 of the most common powerful language models, that generalize well across various NLP tasks, dealing with language (sequence) data. In this experiment, we use BERT, a bidirectional transformer, tuned for the sentiment classification task on the IMDB movie review dataset. As R2DL assumed access to the gradients of the source model ($\nabla X$), this approach works in a semi-black box setting given that we access but do not modify the internal architecture. In this task, we use sentiment classification as the source task, in which there are 2 output classes (positive, negative), and AMP toxicity classification as the target task (toxic, non-toxic). The output-label mapping $h$ is then a simple 1-1 correspondence between (positive, toxic) and (negative, non-toxic). The input data of the source model (BERT) is tokenized on a word-level which form the atoms for our dictionary representation of $V_{S}$. The input data to the target task, AMP sequences, are tokenized on a character level with only 7 distinct tokens. The embeddings of a source vocabulary token, $v_{s}$, is then represented as a weighted combination of the atoms of the AMP sequence tokens. Using the $l^{0}$ norm in our objective function, 100 k-SVD iterations and $\epsilon=0.045$, we are able to achieve accuracy on the order of the benchmark when trained from scratch in table 1. Table 4 shows accuracies for sentiment classification source models and the target task, and Table 5 shows accuracies as we increase the number of k-SVD iterations for a convergent solution.

From Table 5, we see that R2DL acheives the performance (89.34 %) of the baseline (93.7.0 %) when the pretrained classifier is a bidirectional transformer. Intuitively, we can understand why more expressive models have an increased representational capacity that can be leveraged to transfer relevant features from the pretrained classifier. Additionally, more precise k-SVD solutions beyond 100 iterations are not correlated with increased performance, and only increase computational cost, making this approach less efficient.