Deep Extrapolation for Attribute-Enhanced Generation

1 Introduction . Deep generative neural networks can generate realistic data across data-types , from sequences to images to time-series data , with applications in domains such as natural language processing ( NLP ) , computer vision , and speech . Beyond canonical domains , the scientific application of synthetic design of proteins , molecules , and materials can be cast as generative modeling of sequences , graphs , or images ( Anand & Huang , 2018 ; De Cao & Kipf , 2018 ; Madani et al. , 2020 , 2021 ) . Most often , the goal is to design or generate a sample that improves upon the attribute label of interest ( Fig . 1- ( left ) ) , which we term attribute-enhanced generation . Examples include generating a protein sequence with higher binding affinity or a nanomaterial structure with an energetically favorable state , as compared to all of the samples in the training distribution . In these scientific fields , traditional methods for synthetic object design with improved attributes are iterative and expensive , relying on labor- or compute-intensive methods ( Bepler & Berger , 2021 ; Wu et al. , 2021 ; Hie & Yang , 2021 ) . Hence , deep generative models that can design new proteins , molecules , and materials with improved attributes have the potential to dramatically accelerate design research . Beyond scientific applications , extrapolation in generation has potential applications in NLP , such as reducing toxicity or operating in low-resource settings . It is , however , a well-known challenge for deep neural networks to generate samples beyond the training distribution ( Arora et al. , 2017 ; Radford et al. , 2019 ; Xu et al. , 2020 ) . In this work , we develop a method for extrapolation , particularly for sequences . Our approach , called GENhance , is designed to generate an enhanced sequence using a learned latent space . GENhance consists of a generator ( sampler ) and a discriminator ( ranker ) that are jointly trained to minimize generation and discrimination losses , regularized by latent vector smoothing and a cycle-consistency loss . We evaluate GENhance in two data domains . First , we use the Stanford Sentiment Treebank ( SST ) , a natural language benchmark containing movie reviews with five discrete sentiment attributes ( Socher et al. , 2013 ) , to show that GENhance generates strongly positive reviews , after training with no * Equal Contribution Correspondence to amadani @ salesforce.com 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . positive examples . Second , we develop a protein stability dataset for the ACE2 protein ( Chan et al. , 2020 ) with a change in free energy ( ddG ) continuous attribute , and show that GENhance can generate protein sequences with higher stability than the training set ( Fig . 1- ( right ) ) . GENhance significantly outperforms baseline methods based on ( i ) a generator-discriminator model with rejection sampling and ( ii ) an algorithm using Metropolis-Hastings Markov chain Monte Carlo sampling with a trained discriminator . GENhance ’ s performance is further improved when provided access to a few examples with attribute scores beyond the training distribution . Our contributions are summarized below : • We formalize the task of extrapolation for deep generative models , focused on enhancing attributes in sequence generation , with important scientific applications in synthetic object design . • We introduce GENhance , a regularized encoder-decoder framework with a learned latent space , and demonstrate its superior performance with respect to rigorous baseline techniques . • We curate extrapolation benchmarks in NLP and proteins . We release the data , evaluation metrics , along with oracle models and scripts for automatic evaluation of generation quality . 2 Related Work . Generalization to Low Data Regimes : Previous approaches aim to generalize classification and regression to low data settings . Imbalanced classification methods upsample or downsample classes ( Chawla et al. , 2002 ; García & Herrera , 2009 ) or reweight the training cost function ( Huang et al. , 2016 ; Cao et al. , 2019 ; Cui et al. , 2019 ) . Yang et al . ( 2021 ) improve the generalization of regression models in extrapolation and interpolation of the continuous data domain by smoothing both the label and features of the training data . Unlike prior work in this area , GENhance aims to generate samples in low/no data settings . Methods that can better generalize discriminators to these regions are complimentary and orthogonal to our work . Data-Driven Design : Data-driven design aims to learn a distribution over a high-dimensional input space that is optimized for a fitness function corresponding to a desirable property . Design methods often iterate sampling from a generator , and then updating the generator to assign a higher probability to inputs that a discriminator predicts to have higher fitness ( Bedbrook et al. , 2019 ; Biswas et al. , 2021 ; Mansouri Tehrani et al. , 2018 ) . Auto-focused oracles ( Fannjiang & Listgarten , 2020 ) also adapt discriminators throughout this optimization process to using re-weighting of the training examples in the cost function to make them more reliable in the regions where the generator is more likely to generate . CbAS ( Brookes et al. , 2019 ) and DbAS ( Brookes & Listgarten , 2018 ) use a fixed discriminator/oracle model and iteratively learns the distribution of inputs conditioned on a desirable property using importance sampling . CbAS is an improved version of DbAS which also re-weights samples based on how close they are to the original training data . We view these techniques as complementary as GENhance proposes a model-specific architecture for optimizing attributes . Das et al . ( 2021 ) train VAE to learn latent space and use latent space classifiers to sample latent vectors through rejection sampling and decode them into sequences that would have the target attribute/label . Hawkins-Hooker et al . ( 2021 ) also decode generations of a VAE by conditioning on latent vectors that correspond to the target attribute/label . Hoffman et al . ( 2020 ) seek to optimize molecular designs by using zeroth-order optimization on query-based prediction of candidate molecules ’ properties . Gómez-Bombarelli et al . ( 2018 ) build a Gaussian Process ( GP ) regression model trained with latent vectors to predict their inputs ’ labels and use gradient-based optimization on the GP to find sequences with target attributes . Compared with these previous works , the core difference in our approach is the combination of cycle-consistency and contrastive discriminatory objective to train the generator and discriminator as one model . Controllable Text Generation : Our work is also related to controllable text generation , which aims to generate text that corresponds to a user-specified attribute ( Kikuchi et al. , 2016 ; Ficler & Goldberg , 2017 ) . CTRL ( Keskar et al. , 2019 ) generates controlled fluent texts through the use of control codes which are meta-data prepended to the text during generation . Krause et al . ( 2020 ) use a generative discriminator resulting from contrasting predictions from opposing control codes to guide generation . CoCon ( Chan et al. , 2021 ) performs zero-shot controllable text generation without attribute labels . ( Ziegler et al. , 2019 ) optimizes language generation for desirable attributes via human in-the-loop reinforcement learning . Similarly to GENhance , PPLM ( Dathathri et al. , 2020 ) applies a discriminator on top of the latent space of a generative model to guide generation , however , GENhance uses an autoencoder rather than a language model . Lastly , text style transfer methods have used autoencoders with disentangled style latent representations ( Shen et al. , 2017 ; Hu et al. , 2017 ; Yang et al. , 2018 ) . Unlike text style transfer and previous approaches toward controllable text generation , GENhance differs , aside from its model formulation , in that the goal is to optimize and extrapolate a particular attribute beyond the training distribution . 3 Methods . Our goal is to generate sequences with target attribute values that are better than the the training data . Formally , assume that there is a ground-truth oracle ( O ) that maps each sample ( x 2 Rd ) to the target attribute value ( y 2 R ) , i.e. , y = O ( x ) . Given a dataset of oracle labeled samples ( D ) , we aim to generate new sequences where its ground-truth attribute value is better than that of this dataset : ynew > y⌧ , 8 ( x , y ) 2 D : y < y⌧ ( 1 ) To generate samples that satisfy this criterion with high probability , we develop a sampling-ranking framework that consists of a sampler S that proposes a pool of candidate sequences and a ranker R model to infer the relative scores of these candidates . First , we describe two baseline generation techniques that are natural choices for this task and build on them to develop our GENhance model . 3.1 Generator-Discriminator Rejection Sampling . The first baseline , Gen-Disc , is a rejection sampling approach that uses a generator model as the sampler S and a separate discriminator model as the ranker R . The generator is trained to model the training distribution p ( x ) through a language modeling objective where it learns to auto-regressively ( Manning et al. , 1999 ; Bengio et al. , 2003 ) construct the training samples , p ( xt , . . . , xl|x1 , . . . , xt 1 ) = lY i=t p ( xi|x1 , . . . , xi 1 ) , x = xi , . . . , xl ( 2 ) where l is the length of the training sequence . The discriminator model is trained with an objective to predict the relative ranking of sequences from a pair of samples , based on their attribute . Given two training samples ( xa , ya ) and ( xb , yb ) , the pairwise contrastive loss is : Lcontrast = log 1 1 + exp ( ỹa ỹb ) , ỹa = fdisc ( xa ) , ya > yb ( 3 ) where fdisc denotes the discriminator which outputs a scalar score value for each input sequence . We employ the contrastive loss term for this objective since it can be applied to both continuousand discrete-labeled samples . After training , we sample candidate sequences from the generator model in an auto-regressive fashion and use the discriminator to rank the sequences according to the discriminator ’ s output score values . 3.2 Metropolis-Hastings Markov Chain Monte Carlo . Traditional methods for data-driven design rely on iterative optimization of candidates with better attributes . To mimic this process , we design a method that generates candidates using MetropolisHastings MCMC sampling from a population of better candidates . We start with an initial population of sequences , sampled from the training set . In the sampling step , new candidates are proposed by making edits to samples from this initial population , scored with the ranker R , and compared with the previous population of candidates . The probability that new generations are kept in the population depends on the score predicted by R. The cycle of sampling and ranking repeats until terminated . The ranker R takes the form of a neural network , identical to the discriminator model in the Gen-Disc setup . 3.3 GENhance . In the Gen-Disc framework , since the generator is trained only to model the training distribution , there is no direct way to steer the generation distribution towards a certain direction beyond the training data . In the MCMC framework , it might be challenging to find desirable candidates with stochastic mutation operations , since the search space can be large and high-dimensional for many design problems ( Kumar & Levine , 2019 ) . To overcome these limitations , we propose using a learned latent space ( Kingma & Welling , 2013 ) to control the attributes of generated sequences . Architecture : GENhance is an encoder-decoder framework with a latent space between its encoder ( ENC ) and decoder ( DEC ) modules ( Figure 2 ) . The latent vector ( z 2 Rdz ) of a particular input sequence ( x 2 Rdx ) is the output from the encoder module , i.e. , z = ENC ( x ) . In our experiments , z is the hidden state at the location of a < cls > token ( Devlin et al. , 2018 ) that is prepended to the input sequence . Within the latent vector z , representation relevant and irrelevant to the attribute of interest is stored in z|| and z ? respectively , i.e. , z = [ z|| ; z ? ] . To train the encoder to store information about the target attribute in z|| , we train it with a contrastive objective that aims to learn which of the two samples has the better value of the attribute : Lcontrast = log 1 1 + exp ( ỹa ỹb ) , ỹa = f|| ( za|| ) , [ za|| ; za ? ] = ENC ( xa ) , ya > yb ( 4 ) where ( xa , ya ) and ( xb , yb ) are a pair of training samples , each containing an input sequence x and its label y. f|| is an operation that maps z|| to a scalar value . Here we use z|| of dimension 1 and f|| is simply an identity operation . We train GENhance to generate sequences using an objective where the decoder will autoregressively reconstruct a sequence while conditioned on the latent vector z . For an input sequence x of length l , parameterizing the ENC with ✓ and DEC with , we get the reconstruction loss : Lrecon = lX i=t logp ( xi|z , { x1 , . . . , xi 1 } ) = lX i=t log p✓ , ( xi|x ) ( 5 ) To ensure that the perturbed latent vector z would result in plausible generations , we include a smoothing objective , the deterministic Wasserstein autoencoder-maximum mean discrepancy ( WAEMMD ) objective ( Tolstikhin et al. , 2017 ) , to train the latent space as it has shown to be effective for discrete sequences . The WAE-MMD term ( defined in the Supplement A.1 ) penalizes divergence of the latent vectors z from a target prior distribution Pz , which is a unit Gaussian in our case . Lsmooth = MMD ( Pz , z ) ( 6 ) To help learn a better latent space and stronger discriminator within GENhance , we propose a cycleconsistency learning objective ( Lcyc-con ) to train ENC to correctly predict the relative rank between two reconstructed inputs : Lcyc-con = log 1 1 + exp ( ŷa ŷb ) , ŷa = f|| ( ẑa|| ) , [ ẑa|| ; ẑa ? ] = ENC ( x̃a ) , x̃a = DEC ( ENC ( xa ) ) , ya > yb ( 7 ) The intuition behind this objective is two-fold . First , since the discriminator ( ENC ) is used to rank generated sequences during inference , we can improve its performance on these synthetic sequences by also training the discriminator ( ENC ) on generated sequences ( x̃ ) during the training phase . Secondly , by backpropagating the Lcyc-con term through GENhance , it could learn a latent space which generates sequences that are easy for the discriminator to rank accurately . Combining all the training objectives , we can optimize using stochastic gradient descent to approximate the optimal parameters for GENhance : ✓ ⇤ , ⇤ = argmin ✓ , ( contrastLcontrast + reconLrecon + smoothLsmooth + cyc-conLcyc-con ) ( 8 ) To the best of our knowledge , this is the first instance of using cycle-consistency with contrastive loss to train a generative model . Sampling & Ranking : After training , we can sample candidates from GENhance ’ s latent space and rank the generated samples with the scalar scores output by the GENhance ’ s ENC . First , we encode a training sample with ENC to sample a latent vector z . To obtain the latent vector encoding for a new candidate sequence with an improved attribute , we can make a perturbation ( z|| ) to the target attribute-aligned latent component z|| . At the final step , GENhance ’ s DEC conditions on this perturbed latent vector ( z0 ) to generate the improved candidate x̃0 : x̃0 = DEC ( z0 ) , z0 = ⇥ ( z|| + z|| ) ; z ? ⇤ , ⇥ z|| ; z ? ⇤ = ENC ( x ) ( 9 ) The perturbation z|| is determined as the direction that increases the f|| ’ s score output , i.e. , @ f|| ( z|| ) @ z|| . For a linear layer f|| , this term is the weight of the layer while for our case where f|| is an identity operator , z|| is simply a scalar . After generating a pool of candidates with GENhance , we can rank and filter out top-scoring candidates with the GENhance ENC ’ s predicted score : ŷ = f|| ( ENC ( x̃ 0 ) ) ( 10 )

Summary: The paper considers the problem of generating sequences that have label values higher than those observed in training. An approach is proposed to encode information about the attribute to be optimized into the latent vector of an autoencoder through joint training with a discriminator and a cycle consistency objective. The model is evaluated on two different datasets, one on sentiment in movie reviews and the other on protein stability.