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 describes GenEnhance, a method for generating sequences with attributes that do not exist in the training distribution. GenEnhance is based on a discriminator / generator architecture trained with a cycle consistency loss similar to existing methods. The main innovation of GenEnhance is that the discriminator is trained with a contrastive loss for learning to rank generated samples by their attribute.  The paper is well written and combining a contrastive loss with a cycle consistency loss for controllable generation is new as far as I know. However, it is unclear how the cycle-consistency loss performs relative to regression losses and how GenEnhance performs relative to existing methods for model-based optimization in the visible and latent space.