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: Self-training, or pseudo-labeling, is very simple and popular approach in different domains like speech and image recognition.  For the first time, pseudo-labeling is designed for quantum mechanics (QM) calculations problem. This problem is challenging for pseudo-labeling as it cannot use data augmentation ("small perturbational noise in 3D molecular geometry could easily lead to a drastic energy difference") or model noise (like dropout or layer drop due to design of models to work with 3D molecular geometry) which were found to be crucial for pseudo-labeling in other domains. Moreover, the main real case scenario for QM is small amount of supervised data available (low-resource) which makes pseudo-labeling is even harder. In this paper authors propose learning a single model with time-to-time regenerating pseudo-labels and incorporating these pseudo-labeled data according to the model uncertainty. The uncertainty is modeled in the following way: labels are assumed to be from normal distribution and model is designed to predict parameters of this normal distribution, thus model estimates expectation of the label and its dispersion. This uncertainty detects and prevents bad pseudo-label from affecting the model training by adaptively weighting the data according to inverse of uncertainty value. Pseudo-labeling results show that direction of improving learning strategy instead of improving physics-based representation has potentials. With experiments it is demonstrated that proposed pseudo-labeling algorithm, pseud$\sigma$, improves results over supervised baseline, works across different backbone models and applicable for low-resource regime. Moreover it demonstrates improved results for out of distribution data.