Adaptive Pseudo-labeling for Quantum Calculations

1 INTRODUCTION . Accurate quantum mechanical ( QM ) calculations of drug-like molecules at CCSDT ( coupled cluster single-double-triple ) or MP2 ( second order Møller-Plesset ) ( Watts et al. , 1992 ) level of theory , which is essential to characterize biomolecular interactions , continue to be prohibitively expensive , despite recent advances in hardware capabilities . Machine learning ( ML ) models have astonishing performance in approximating these calculations at a fraction of the computational cost ( von Lilienfeld & Burke , 2020 ) . Such speedups could accelerate the discovery of new therapeutics . In the absence of large-scale benchmark data sets reporting CCSDT or MP2 level calculations , most publications on this topic have relied on QM9 , a standard benchmark of Density Functional Theory ( DFT ) -level energy and properties , for training and evaluating QM/ML models . This data set contains molecular geometries and energies that were generated using DFT , which provides a faster ( compared to CCSDT , for example ) , but less accurate description of molecular energetics and properties . Previous publications have utilized computed geometries and properties from in-distribution molecules in QM9 to test and validate model performance , showing excellent performance of QM/ML models as demonstrated by low mean absolute error ( MAE ) of predicted energies of other properties relative to the gold-standard DFT calculations ( i.e. , gold label ) . However , it is unclear how the reported architectures ( e.g. , SchNet ( Schütt et al. , 2017 ) ) would perform in the regime of low but accurate data ( CCSDT or MP2 ) . Two challenges remain in the way of the realistic adoption of ML-aided QM calculations . Firstly , training molecules only cover part of the distribution , and real-world adoption requires out-ofdistribution generalization . Secondly , computing QM properties for datasets of the size of QM9 is costly . Widespread applicability of ML for QM ( e.g. , such as to CCSDT/MP2 ) thus requires models to generalize well given low-data constraints . Both challenges listed above can be attributed to the difficulty and cost of generating experimental data in chemistry and materials science . If we have a large and diverse set of labeled molecules , the ML model can generalize to larger chemical spaces and achieve better predictive performance . Present work . Our key observation is that state-of-the-art ML models can obtain a low-precision estimation of the properties even given a small set of labeled molecules . This suggests that the ML predicted label for any unlabelled molecule can be used as a low-precision estimation of the true label . We hypothesize that by considering these low-precision pseudo-labels as the training targets in the training process , we can largely increase the training dataset size and diversity and alleviate the fundamental label scarcity issue , along with the associated challenges in low-data and out-of-distribution generalization . In our study , we develop a simple , effective , and model-agnostic pseudo-labeling strategy called PSEUDσ . Notably , we solicit an extensive array of unlabeled molecules from the PC9 dataset ( Glavatskikh et al. , 2019 ) where an ML model assigns pseudo-labels for them . This generates a large set of ” labeled ” training data points , even when only a relatively small number of reference QM properties are available for training . As the unlabeled dataset is diverse and unseen in QM9 , it also helps generalization to unseen data . One crucial issue in pseudo-labeling is the introduced bias from low-quality pseudo-labels . To resolve this , we rely on a key observation that a data point with less evidence/higher model uncertainty is more likely to be of low-quality pseudo-label ( Section 6 ) . Thus , we use model-generated evidential uncertainty to quantify each unlabeled data and use it to adaptively lower the weight of bad pseudo-labels in the training loss to reduce the bias effect . In summary , we have made the following contributions : ( 1 ) Previous QM focus on in-distribution and label abundant setting while we point out that the realistic ML-aided QM challenges lie in low-data and out-of-distribution settings ; ( 2 ) Pivoting away from the status quo in improving the physics-based representation , we propose to look at data-centric approaches on learning from the vast array of unlabeled molecules ; ( 3 ) We propose an episodic pseudo-labeling scheduling strategy designed specifically for QM since many classic pseudo-labeling tricks do not work well in QM ( e.g. , data augmentation , model noise , student-training , re-initialization ) ; ( 4 ) To prevent bad pseudo-label in biasing the training , we devise an adaptive weighting scheme where the weights are generated using evidential uncertainty such that bad pseudo-labels are automatically filtered out ; ( 5 ) We derive a theoretical connection between the evidential loss and entropy minimization framework ; ( 6 ) Empirically , we show PSEUDσ can improve QM accuracy for any atomistic model across full-data , low-data , and out-of-distribution settings . 2 RELATED WORKS . ML-aided quantum calculations . Recently , many ML models have been proposed to improve quantum calculations . They mainly focus on improving the physics-based representation in the full QM9 dataset setting ( Schütt et al. , 2017 ; Unke & Meuwly , 2019 ; Anderson et al. , 2019 ; Lu et al. , 2019 ; Klicpera et al. , 2020 ; Liu et al. , 2021 ; Qiao et al. , 2021 ) . In contrast , our work proposes to shift the focus on improving training strategy instead of the model architecture . In addition , PSEUDσ is model-agnostic , and it can improve on any atomistic model . Additionally , we focus on realistic quantum calculations such as learning in the low-data regime and out-of-distribution inference . Pseudo-labeling . Pseudo-labeling/self-training generates pseudo-labels for unlabeled data . Numerous works exist for how to assign pseudo-labels , notably through trained ML model prediction ( Lee et al. , 2013 ) , label propagation ( Shi et al. , 2018 ; Iscen et al. , 2019 ) , and history cache ( Likhomanenko et al. , 2021 ; Higuchi et al. , 2021 ) . PSEUDσ is different as it focuses on detecting and preventing bad pseudo-label from affecting the model . Also , PSEUDσ adopts a novel episodic pseudo-label strategy with a re-initialized learning rate . ( Xie et al. , 2020 ) re-initialize the network as a student when a new pseudo-label set is generated along with noise per epoch . In contrast , PSEUDσ has no student , and no noise as both are shown to be ineffective for QM in Section 6 . In addition , small perturbational noise in 3D molecular geometry could easily lead to a drastic energy difference . Thus , the noise strategy does not work for QM tasks . More related is a concurrent work ( Rizve et al. , 2021 ) that develops an uncertainty-aware pseudo-labeling strategy , but they introduce additional hyperparameters to remove pseudo-labels at some uncertainties . In contrast , PSEUDσ uses an effective adaptive weighting scheme , along with an episodic pseudo-labeling training schedule . Additionally , PSEUDσ is the first method that studies pseudo-label in quantum calculations , which present unique challenges and motivations . Uncertainty . Model uncertainty is a well-studied subject ( Kendall & Gal , 2017 ; Lakshminarayanan et al. , 2017 ; Blundell et al. , 2015 ) . ( Amini et al. , 2020 ) use evidential uncertainty to add a prior over the gaussian parameters to search for higher-order patterns for regression tasks . PSEUDσ leverages evidential uncertainty as the uncertainty measure . Note that PSEUDσ is uncertainty measure-agnostic . We can easily switch to alternative uncertainty measures . Recently , ( Soleimany et al. , 2021 ) adapt evidential uncertainty and has shown it can successfully help guide property prediction . In contrast , we leverage evidential uncertainty as a proxy for pseudo-label quality to tackle low-data and out-ofdistribution challenges in realistic quantum calculations setup . 3 PROBLEM FORMULATION . Let X = { x1 , . . . , xN } denote N molecules , where each molecule xi is uniquely defined by 3D coordinates { ( aij , bij , cij ) } Nij=1 for Ni atoms with atom types { tj } Nij=1 in the corresponding molecule . We then denote Y = { y1 , . . . yN } a set of quantum mechanical properties for each molecule . The labeled dataset thus consists of a set of pairs of 3D coordinates and scalar labels D = { X , Y } . In addition to the labeled data , we solicit a large quantity of unlabeled data to generate pseudo-labels . We denote an unlabeled dataset U = { x1 , . . . , xM } , where M is the size of the unlabeled dataset . Given an atomistic model f ( · ) , we can generate pseudo-labels Ŷ = { ŷ1 , . . . , ŷM } , where ŷi = f ( xi ) for xi ∈ U . The problem is to find a machine learning-based atomistic model f : x 7→ y that can establish an accurate map from 3D coordinates to the quantum mechanical properties of the molecules , with the help of pseudo-labeled dataset U . 4 PSEUDσ : ADAPTIVE PSEUDO-LABELING FOR QUANTUM CALCULATIONS PSEUDσ ( Figure 1 ) is an approach for quantum chemical property prediction . Building on theoretical motivation from Section 5 , PSEUDσ solicits pseudo-labels on a vast array of an unlabeled dataset to increase the diversity of the training space via an episodic labeling strategy . Then , it adaptively weights the pseudo-labels using evidential uncertainty to allow a positive transfer . The overview is in Algorithm 1 . Episodic Pseudo-labeling . We devise a pseudo-labeling strategy that can ensure learning from the pseudo-labels to the fullest extent for QM . We have made two distinct modifications compared to existing works . First is the pseudo-label scheduling . In the standard pseudo-labeling ( Lee et al. , 2013 ) , pseudo-labels are updated in every update and the model is continuously trained . In contrast , we devise an episodic training strategy , where each episode consists of K epochs , and pseudo-labels are regenerated in every episode , while the model is continuously trained . This is important because we observe that fast updates on pseudo-labeling prevents the model to fully learn from all useful information in pseudo-labels . In contrast , our episodic approach gives the model more time to absorb useful information from a given set of pseudo-labels . Second modification is the model update . In self-training ( Rizve et al. , 2021 ; Xie et al. , 2020 ) , a set of pseudo-labels are regenerated after K epochs and the model is reinitialized . Instead , we train the same model throughout episodes . This new strategy allows the model to be exposed to a larger number of labels or training data points given the same time frame . For each episode , we also reinitialize the learning rate with a small step-wise decay strategy to allow the model a chance to jump out of the local optimum from the previous set of pseudo-labels . Formally , PSEUDσ mainly consists of three stages : in the first stage , regular training is conducted on labeled data D , and the output model is the initialized model f ( 1 ) . In the second stage , the updated model at episode k then conducts inference on the entire unlabeled data Ŷ = f ( k ) ( U ) to generate the pseudo-label set . The per-episode pseudo-label set is then combined with the gold-labeled data to form the training data for the next episode . In the third stage , the model is further trained using the combined dataset to get a new model f ( k+1 ) after N epochs . The second and third stages are then reiterated till the loss stops decreasing . Evidential uncertainty quantification . Pseudo-labels are noisy . Many are incorrect and can potentially lead to negative transfer . Thus , it is instrumental in detecting the quality of pseudo-labels . However , there is no auxiliary information in the dataset about the pseudo-labels . Thus , we need to quantify it through some proxies that can be assigned without auxiliary information . Our key observation is that low-quality pseudo-labels have high model uncertainty , and high-quality pseudo-labels have low model uncertainty . Another advantage of model uncertainty is that it can be estimated solely from x , if we make it model uncertainty-aware . Building on the theoretical motivation about the connection between evidential uncertainty and entropy minimization in Section 5 , we use evidential uncertainty as to the proxy for label quality . The evidential modeling of molecular property allows us to derive analytical solution of uncertainty , which could be directly used to weight the pseudo-labels . Formally , we can model the label probabilistically as it is drawn from ( y1 , · · · , yi ) ∼ N ( µ , σ2 ) , where the mean µ are variance σ2 are unknown . To estimate them , we pose a prior µ ∼ N ( γ , σ2v−1 ) , σ2 ∼ Γ−1 ( α , β ) , ( 1 ) where the parameters θ = ( µ , σ ) is an instantiation of the posterior p ( µ , σ2|γ , v , α , β ) . The choice of prior allows the factorization p ( µ , σ2 ) = p ( µ ) p ( σ2 ) ( Jordan , 2009 ) . The posterior then becomes a NormalInvGamma ( γ , v , α , β ) where the maximum likelihood estimation of θ can be analytically found as E [ µ ] = γ , E [ σ2 ] = β α− 1 . ( 2 ) Here , E [ σ2 ] plays the role of the aleatoric ( data ) uncertainty . The uncertainty of the model prediction can also be calculated , i.e . epistemic uncertainty : Var [ µ ] = E [ σ2 ] /v = β v ( α− 1 ) . ( 3 ) As the MLE is deterministic , the model can output the four prior parameters { γ , v , α , β } directly where the prediction and uncertainty can be derived from them analytically . The prior is optimized by evidential loss Levi ( Amini et al. , 2020 ) : Levii = −log St ( yi ; γ , β ( 1 + v ) vα , 2α ) + λ|yi − γ| ( 2v + α ) , ( 4 ) where the first term is to maximize the log-likelihood of the posterior predictive , which is derived as the Student ’ s t-distribution . The second term is a regularizer that encourages lower total evidence , which correlates positively to higher epistemic uncertainty , when the gap between ground truth and prediction is high ( i.e . when model prediction is with high error ) . Similarly , it encourages lower uncertainty when the model prediction is errorneous . This encourages the model to generate an accurate estimate of uncertainty or the degree of errors for the pseudo-labeled data points . The regularization is controlled by a hyperparameter λ. Adaptive weighting . The evidential uncertainty detects the low-quality pseudo-labels . The next step is to remove the noisy effect from the model training . Naive methods often use removal based on a threshold ( Rizve et al. , 2021 ) . However , it has two disadvantages : ( 1 ) it introduces a new hyperparameter - the threshold ; ( 2 ) it removes a portion of unlabeled noisy data , which can contain useful information . Instead , we propose an adaptive weighting mechanism that adapts the evidential loss given the inverse epistemic uncertainty . Intuitively , a higher uncertainty data point requires a lower effect in the loss function because it is more likely that the sample has a low pseudo-label quality , Algorithm 1 : PSEUDσ Algorithm . Input : Labeled data D = { ( x1 , y1 ) , · · · , ( xN , yN ) } , unlabled data U = { x1 , · · · , xNU } Û ← { } , Ŵ ← { } // Initialize with empty unlabeled data for k ∈ { 1 , · · · , K } // Outer-loop with K episodes do T ← D ∪ Û // Join updated pseudo-labels for ( xi , yi ) ∈ T // Inner-loop with N epochs do θi = ( γi , vi , αi , βi ) = f ( k−1 ) ( xi ) // Evidental parameters ŷi = E [ µ ] = γi // Posterior prediction L = L ( ŷi , yi , θi , Ŵi ) // Adaptive evidential loss via Eq . 5 f ( k−1 ) = Update ( f ( k−1 ) , L ) // Inner-loop update end f ( k ) ← f ( k−1 ) // Update teacher model for pseudo-labels for xi ∈ U do θ̂i = ( γ̂i , v̂i , α̂i , β̂i ) = f ( k ) ( xi ) ŷi = γ̂i // Infer a new set of pseudo-labels Ûi ← ( xi , ŷi ) // Update pseudo-labels Ŵi ← Var [ µ ] −1i = v̂i ∗ ( α̂i − 1 ) /β̂i // Update adaptive weights end end and we want to reduce its effect on the model . Conversely , if a pseudo-label has low uncertainty , the label quality is high enough to be used as a proxy of a gold-label . Thus , we should assign a higher score in the loss for it . The uncertainty is from the teacher model in the previous episode and is fixed throughout the current episode . Thus , the adaptive weight for each pseudo data point i becomes Ŵi = Var [ µ ] −1i . The final loss then becomes L = 1|D| ∑ i∈D Levii + ∑ i∈U Ŵi∑ i∈U Ŵi Levii , ( 5 ) where the first term corresponds to the labeled dataset D does not have weights unlike the second term corresponding to the unlabeled data . This adaptive loss solves two disadvantages since it has zero hyperparameters , and it removes the effect of bad pseudo-labels while retaining all training examples including the noisy ones to maximize the diversity of the training space . 5 PSEUDσ MOTIVATION : CONNECTION TO ENTROPY MINIMIZATION We derive motivation about why evidential uncertainty and the weighting mechanism could be beneficial to pseudo-labeling based on the entropy minimization framework for semi-supervised learning from Grandvalet & Bengio ( 2004 ; 2006 ) ; Lee et al . ( 2013 ) . Notably , our use of Bayesian modeling enables us to analytically derive a conditional entropy for pseudo-labeled data . We find that evidential loss strongly relates to conditional entropy , and minimizing evidential loss directly minimizes entropy . Secondly , we find the conditional entropy could be decomposed into the inverse epistemic uncertainty and the log-likelihood , which motivates our weighting mechanism . In the standard regression setting , one seeks to maximize the likelihood of the model pθ ( Y|X ) on the labeled data set D. To utilize the unlabeled data set , we need to extract some useful information on how the model behaves on the unlabeled dataset and inject this information to improve the model . To measure the utility , entropyH ( Y | U ) is introduced ( Grandvalet & Bengio , 2006 ) as a proxy to measure the amount of information in unlabeled data : H ( Y | U ) = ∑ xi∈U Ey∼pθ ( y |xi ) [ −log pθ ( y |xi ) ] . ( 6 ) Throughout the text , we are referring to entropy as Shannon entropy . Intuitively , high entropy is indicative of the poor model confidence on the unlabeled data point , while the low entropy can be interpreted as high confidence of the model predictions . High entropy , associated with random predictions , while low entropy is associated with non-random behavior . Hence force , we hypothesize that small entropy may be indication of a signal that our model can benefit from . Small entropy , as seen below , corresponds to high model confidence . and vice versa . Large entropy corresponds to high model uncertainty . Entropy minimization framework casts the regression as the following optimization problem : arg max θ [ log pθ ( Y|X ) − cH ( Y | U ) ] , ( 7 ) where c is the proportionality constant . Intuitively , here , the objective tends to maximize the loglikelihood on the labeled dataset while minimizing the entropy on the unlabeled data set at the same time to transfer knowledge from unlabeled data . In previous works ( Schütt et al. , 2017 ; Liu et al. , 2021 ) , molecule properties are not modeled probabilistically such that entropy calculation is infeasible . In contrast , PSEUDσ uses Bayesian modeling approaches that allow us to analytically calculate the entropy . For every molecule xi the machine learning model outputs four parameters f ( xi ) = ( αi , βi , γi , νi ) . Based on these parameters , the likelihood of label y given the input molecule xi is given by the Student ’ s t-distribution in the context of evidential regression pθ ( y|xi ) = St ( y ; γi , σ2st , i , 2αi ) ( 8 ) evaluated at location parameter γi , Student ’ s t-distribution scale parameter σ2st , i = βi ( 1+νi ) νiαi and 2αi degrees of freedom . The entropy of the Student ’ s t-distribution given in terms of evidential parameters is readily available ( Appendix A ) : H ( y |xi ) = 2αi + 1 2 ( Ψ ( 2αi + 1 2 ) −Ψ ( αi ) ) + log √ 2αi B ( αi , 1 2 ) + 1 2 log σ2st , i , ( 9 ) where Ψ is a digamma function and B ( · , · ) is a beta function . If we take model ( epistemic uncertainty ) ( Eq . 3 ) as our evidential uncertainty , we can show that minimizing entropy directly relates to minimizing epistemic uncertainty . We plot the relation between the entropy and epistemic uncertainty in the Figure 2 . As the next steps , we aim to uncover the dependence of the entropy on the model uncertainty of our pseudo-labeling approach . This can be done if we make two simplifications in entropy evaluation . Firstly , to introduce iterations as in pseudo-labeling , we replace entropy with the crossentropy between two probability distributions : the predictions y are generated from the probability distribution pθ ( t−1 ) ( y|xi ) at iteration step t − 1 and log-likelihood are evaluated with respect to probability distribution pθ ( t ) ( y|xi ) at iterative step t : H ( Y | U ) ≈ ∑ xi∈U Ey∼pθ ( t−1 ) ( y|xi ) [ − log pθ ( t ) ( y |xi ) ] . ( 10 ) Upon convergence , t → +∞ , the probability distributions at every iterative step pθ ( t−1 ) ( y|xi ) ≈ pθ ( t ) ( y|xi ) and are approximately the same and one can view introduced cross-entropy with respect to time step t as entropy . At the earlier stages of training cross-entropy acts as a regularizer encouraging network parameters θ ( t ) to match θ ( t− 1 ) . As a second approximation , to uncover model uncertainty in mathematical formulas , we approximate the probability distribution at time step t− 1 . We resort to empirical estimate of the entropy , as done in ( Grandvalet & Bengio , 2004 ) . We select labels y at the highest mode of probability distribution pθ ( t−1 ) ( y |xi ) , which corresponds to y = γt−1i . We obtain the following approximate for the entropy : H ( Y | U ) ≈ Hemp ( Y | U ) = − ∑ xi∈U Et−1i log pθ ( t ) ( γt−1i |xi ) , ( 11 ) where the log probabilities are weighted by empirical probabilities as weights Et−1i evaluated at iterative step t− 1 when plugged into Eq . 8 ( also see Appendix . Eq . 18 for exact formula of Student ’ s t-distribution ) Et−1i = St ( y = γi , σ2st , i , 2αi ) = 1√ 2αi σ2st , iB ( 1 2 , αi ) . ( 12 ) To establish a relationship between empirical weights Et−1i and aleatoric E [ σ2i ] / epistemic Var [ µi ] uncertainties we rewrite σ2st , i = αi − 1 αi ( Var [ µi ] + E [ σ2i ] ) ( 13 ) Et−1i = ( Var [ µi ] + E [ σ2i ] ) − 1 2√ 2 B ( 12 , αi ) √ αi − 1 . ( 14 ) Empirical coefficients depend on aleatoric , epistemic uncertainties and αi parameter , which can be interpreted as virtual observations in support of the variance estimation ( Jordan , 2009 ) . In the limiting case αi 1 one can approximate beta function via Stirling formula B ( 12 , αi ) ≈ √ π α − 12 i and empirical weights become Et−1i ≈ ( Var [ µi ] + E [ σ2i ] ) − 1 2 . ( 15 ) We can express the empirical coefficients depend both on aleatoric and epistemic uncertainties in a symmetric fashion . We selected adaptive pseudo-labeling coefficientsWi in our pseudo-labeling approach Eq . 5 to be inverse epistemic/model uncertainties . We can see , that those coefficients directly relate to empirical coefficients derived from entropy minimization approach Eq . 14 , as empirical coefficients also depend on model uncertainty in the inverse fashion . As the model uncertainty increases , the empirical coefficients Ei tend to decrease to minimize the entropy .

Summary: The paper is focused on improving quantum-mechanical properties of molecules in a setting similar to semi-supervised learning, where X is known for samples without Y. In an iterative fashion, the current model is used to predict pseudo-labels for those samples, then pseudo-labels together with labeled samples are used to improve the model. The paper uses existing approaches to quantify the uncertainty in individual pseudo-labels and then uses the uncertainty to weight their contribution to the loss.