Iterative Target Augmentation for Effective Conditional Generation

1 INTRODUCTION . Deep architectures are becoming increasingly adept at generating complex objects such as images , text , molecules , or programs . Many useful generation problems can be seen as translation tasks , where the goal is to take a source ( precursor ) object such as a molecule and turn it into a target satisfying given design characteristics . Indeed , molecular optimization of this kind is a key step in drug development , though the adoption of automated tools remains limited due to accuracy concerns . We propose here a simple , broadly applicable meta-algorithm to improve translation quality . Translation is a challenging task for many reasons . Objects are complex and the available training data pairs do not fully exemplify the intricate ways in which valid targets can be created from the precursors . Moreover , precursors provided at test time may differ substantially from those available during training — a scenario common in drug development . While data augmentation and semisupervised methods have been used to address some of these challenges , the focus has been on either simple prediction tasks ( e.g. , classification ) or augmenting data primarily on the source side . We show , in contrast , that iteratively augmenting translation targets significantly improves performance on complex generation tasks in which each precursor corresponds to multiple possible outputs . Our iterative target augmentation approach builds on the idea that it is easier to evaluate candidate objects than to generate them . Thus a learned predictor of target object quality ( a filter ) can be used to effectively guide the generation process . To this end , we construct an external filter and apply it to the complex generative model ’ s sampled translations of training set precursors . Candidate translations that pass the filter criteria become part of the training data for the next training epoch . The translation model is therefore iteratively guided to generate candidates that pass the filter . The generative model can be viewed as an adaptively tuned prior distribution over complex objects , with the filter as the likelihood . For this reason , it is helpful to apply the filter at test time as well , or to use the approach transductively1 to adapt the generation process to novel test cases . The approach is reminiscent of self-training or reranking approaches employed with some success for parsing ( McClosky et al. , 2006 ; Charniak et al. , 2016 ) . However , in our case , it is the candidate generator that is complex while the filter is relatively simple and remains fixed during the iterative process . 1Allowing the model to access test set precursors ( but not targets ) during training . We demonstrate that our meta-algorithm is quite effective and consistent in its ability to improve translation quality in the supervised setting . On a program synthesis task ( Bunel et al. , 2018 ) , under the same neural architecture , our augmented model outperforms their MLE baseline by 8 % and their RL model by 3 % in top-1 generalization accuracy ( in absolute measure ) . On molecular optimization ( Jin et al. , 2019a ) , their sequence to sequence translation baseline , when combined with our target data augmentation , achieves a new state-of-the-art result and outperforms their graph based approach by over 10 % in success rate . Their graph based methods are also improved by iterative target augmentation with more than 10 % absolute gain . The results reflect the difficulty of generation in comparison to evaluation ; indeed , the gains persist even if the filter quality is reduced somewhat . Source side augmentation with unlabeled precursors ( the semi-supervised setting ) can further improve results , but only when combined with the filter in the target data augmentation framework . We provide ablation experiments to empirically highlight the effect of our method and also offer some theoretical insights for why it is effective . 2 RELATED WORK . Molecular Optimization The goal of molecular optimization is to learn to modify compounds so as to improve their chemical properties . Jaques et al . ( 2017 ) ; You et al . ( 2018 ) ; Popova et al . ( 2018 ) used reinforcement learning approaches , while Jin et al . ( 2019a ; b ) formulated this problem as graphto-graph translation and significantly outperformed previous methods . However , their performance remains imperfect due to the limited size of given training sets . Our work uses property prediction models to check whether generated molecules have desired chemical properties . Recent advances in graph convolutional networks ( Duvenaud et al. , 2015 ; Gilmer et al. , 2017 ) have provided effective solutions to predict those properties in silico . In this work , we use an off-the-shelf property prediction model ( Yang et al. , 2019 ) to filter proposed translation pairs during data augmentation . Program Synthesis Program synthesis is the task of generating a program ( using domain-specific language ) based on given input-output specifications ( Bunel et al. , 2018 ; Gulwani , 2011 ; Devlin et al. , 2017 ) . One can check a generated program ’ s correctness by simply executing it on each input and verifying its output . Indeed , Zhang et al . ( 2018 ) ; Chen et al . ( 2019 ) leverage this idea in their respective decoding procedures , while also using structural constraints on valid programs . Semi-supervised Learning Our method is related to various approaches in semi-supervised learning . In image and text classification , data augmentation and label guessing ( Berthelot et al. , 2019 ; Xie et al. , 2019 ) are commonly applied to obtain artificial labels for unlabeled data . In machine translation , Norouzi et al . ( 2016 ) sample new targets from a stationary distribution in order to match the model distribution to the exponentiated payoff distribution centered at a single target sentence . Back-translation ( Sennrich et al. , 2015 ; Edunov et al. , 2018 ) creates extra translation pairs by using a backward translation system to translate unlabeled sentences from a target language into a source language . In contrast , our method works in the forward direction because many translation tasks are not symmetric . Moreover , our data augmentation is carried out over multiple iterations , in which we use the augmented model to generate new data for the next iteration . In syntactic parsing , our method is closely related to self-training ( McClosky et al. , 2006 ) . They generate new parse trees from unlabeled sentences by applying an existing parser followed by a reranker , and then treat the resulting parse trees as new training targets . However , their method is not iterative , and their reranker is explicitly trained to operate over the top k outputs of the parser ; in contrast , our filter is independent of the generative model . In addition we show that our approach , which can be viewed as iteratively combining reranking and self-training , is theoretically motivated and can improve the performance of highly complex neural models in multiple domains . Co-training ( Blum & Mitchell , 1998 ) and tri-training ( Zhou & Li , 2005 ; Charniak et al. , 2016 ) also augment a parsing dataset by adding targets on which multiple baseline models agree . Instead of using multiple learners , our method uses task-specific constraints to select correct outputs . 3 ITERATIVE TARGET AUGMENTATION . Our iterative target augmentation framework can be applied to any conditional generation task with task-specific constraints . For example , molecular optimization ( Jin et al. , 2019a ; b ) is the task of transforming a given molecule X into another compound Y with improved chemical properties , Algorithm 1 Augmentation by iterative target augmentation Input : Original training set D = [ ( X1 , Y1 ) , . . . , ( Xn , Yn ) ] 1 : procedure AUGMENTDATASET ( D , Mt ) 2 : Dt+1 = D . Initialize augmented dataset . 3 : for ( Xi , Yi ) in D do 4 : for attempt in 1 , . . . , C do 5 : Apply modelMt to Xi to sample candidate Y ′ 6 : if Y ′ passes external filter then 7 : Add ( Xi , Y ′ ) to Dt+1 8 : if K successful translations added then 9 : break from loop 10 : return augmented dataset Dt+1 11 : procedure TRAIN ( D ) 12 : for epoch in 1 , . . . , n1 do . Regular training 13 : Train model on D. 14 : for epoch in 1 , . . . , n2 do . Iterative target augmentation 15 : Dt+1 = AUGMENTDATASET ( D , Mt ) 16 : Mt+1 ← Train modelMt on Dt+1 . while constraining Y to remain similar to X . Program synthesis ( Bunel et al. , 2018 ; Chen et al. , 2019 ) is the task of generating a program Y satisfying input specification X ; for example , X may be a set of input-output test cases which Y must pass . Without loss of generality , we formulate the generation task as a translation problem . For a given input X , the model learns to generate an output Y satisfying the constraint c. The proposed augmentation framework can be applied to any translation model M trained on an existing dataset D = { ( Xi , Yi ) } . As illustrated in Figure 1 , our method is an iterative procedure in which each iteration consists of the following two steps : • Augmentation Step : Let Dt be the training set at iteration t. To construct each next training set Dt+1 , we feed each input Xi ∈ D ( the original training set , not Dt ) into the translation model up toC times to sampleC candidate translations Y 1i . . . Y C i . 2 We take the firstK distinct translations for each Xi satisfying the constraint c and add them to Dt+1 . When we do not find K distinct valid translations , we simply add the original translation Yi to Dt+1 . • Training Step : We continue to train the modelMt over the new training set Dt+1 for one epoch . The above training procedure is summarized in Algorithm 1 . As the constraint c is known a priori , we can construct an external filter to remove generated outputs that violate c during the augmentation step . At test time , we also use this filter to screen predicted outputs . To propose the final translation of a given input X , we have the model generate up to L outputs until we find one satisfying the constraint c. If all L attempts fail for a particular input , we just output the first of the failed attempts . 2One could augment Dt instead of D and continuously expand the dataset , but the empirical effect is small ( see Appendix B.3 ) . We note this augmentation step can be trivially parallelized if speed is a concern . Finally , as an additional improvement , we observe that the augmentation step can be carried out for unlabeled inputs X that have no corresponding Y . Thus we can further augment our training dataset in the transductive setting by including test set inputs during the augmentation step , or in the semi-supervised setting by simply including unlabeled inputs . 4 MOTIVATION FOR ITERATIVE TARGET AUGMENTATION . We provide here some theoretical motivation for our iterative target augmentation framework . For simplicity , we consider an external filter cX , Y that is a binary indicator function representing whether output Y satisfies the desired constraint in relation to inputX . In other words , we would like to generate Y such that Y ∈ B ( X ) = { Y ′|cX , Y ′ = 1 } . If the initial translation model P ( 0 ) ( Y |X ) serves as a reasonable prior distribution over outputs , we could simply “ invert ” the filter and use P ( ∗ ) ( Y |X ) ∝ P ( 0 ) ( Y |X ) · cX , Y ( 1 ) as the ideal translation model . While this posterior calculation is typically not feasible but could be approximated through samples , it relies heavily on the appropriateness of the prior ( model prior to augmentation ) . Instead , we go a step further and iteratively optimize our parametrically defined prior translation model Pθ ( Y |X ) . Note that the resulting prior can become much more concentrated around acceptable translations . We maximize the log-likelihood that candidate translations satisfy the constraints implicitly encoded in the filter EX [ logPθ ( cX , Y = 1 | X ) ] ( 2 ) In many cases there are multiple viable outputs for any given input X . The training data may provide only one ( or none ) of them . Therefore , we treat the output structure Y as a latent variable , and expand the inner term of Eq . ( 2 ) as logPθ ( cX , Y = 1 | X ) = log ∑ Y Pθ ( Y , cX , Y = 1 | X ) ( 3 ) = log ∑ Y P ( cX , Y = 1 | Y , X ) Pθ ( Y |X ) ( 4 ) = log ∑ Y cX , Y · Pθ ( Y |X ) ( 5 ) Since the above objective involves discrete latent variables Y , we propose to maximize Eq . ( 5 ) using the standard EM algorithm ( Dempster et al. , 1977 ) , especially its incremental , approximate variant . The target augmentation step in our approach is a sampled version of the E-step where the posterior samples are drawn with rejection sampling guided by the filter . The number of samples K controls the quality of approximation to the posterior.3 The additional training step based on the augmented targets corresponds to a generalized M-step . More precisely , let P ( t ) θ ( Y |X ) be the current translation model after t epochs of augmentation training . In epoch t + 1 , the augmentation step first samples C different candidates for each input X using the old model P ( t ) parameterized by θ ( t ) , and then removes those which violate the constraint c , interpretable as samples from the current posterior Q ( t ) ( Y |X ) ∝ Pθ ( t ) ( Y |X ) cX , Y . As a result , the training step maximizes the EM auxiliary objective via stochastic gradient descent : J ( θ | θ ( t ) ) = EX [ ∑ Y Q ( t ) ( Y |X ) logPθ ( Y |X ) ] ( 6 ) We train the model with multiple iterations and show empirically that model performance indeed keeps improving as we add more iterations . The EM approach is likely to converge to a different and better-performing translation model than the initial posterior calculation discussed above .

Summary: This paper proposes a training scheme to enhance the optimization process where the outputs are required to meet certain constraints. The authors propose to insert an additional target augmentation phase after the regular training. For each datapoint, the algorithm samples candidate outputs until it find a valid output according the an external filter. The model is further fine-tuned on the augmented dataset. In experiments, the authors evaluated the proposed training scheme on two datasets: molecular optimization and program synthesis. Results show that the proposed training algorithm improves the success rate and the diversity among generated targets.