Environment Diversification with Multi-head Neural Network for Invariant Learning
Anonymous Author s
Abstract
Neural networks are often trained with empirical risk minimization; however, ithas been shown that a shift between training and testing distributions can cause unpredictable performance degradation.
On this issue, a research direction, invariantlearning, has been proposed to extract causal features insensitive to the distributional changes.
This work proposes an invariant learning framework containing amulti-head neural network to absorb data biases.
We show that this framework doesnot require prior knowledge about the environment or strong assumptions aboutthe pre-train model.
We also reveal that the proposed algorithm has theoreticalconnections to recent studies discussing properties of variant and invariant features.
Finally, we demonstrate that empirically models trained with this framework aremore robust against distributional shift.
Introduction
Ensuring model performance on unseen data is a common yet challenging task in machine learning.
Awidely adopted solution would be empirical risk minimization (ERM), where training and testing dataare assumed to be independent and identically distributed.
However, data in real-world applicationscan come with undesired biases, causing a shift between training and testing distributions.
It hasbeen known that the distributional shifts can severely harm ERM model performance and evencause the trained model be worse than random predictions [9].
In this work, we focus on Invariant
Learning , which aims at learning causality expected to be robust against distributional changes.
Invariant risk minimization (IRM) [1] has been proposed as a popular solution for invariant learning.
Specifically, IRM is based on an assumption that training data can be partitioned into multiple subsetsor environments having distinct data distributions.
The learning objective is then designed as astandard ERM loss function with a penalty term constraining the trained model (e.g. classifier) to beoptimal in all the environments.
IRM has shown to be effective; however, we note that IRM and many invariant learning methods relyon strict assumptions which limit the practical impacts.
The limitations are summarized as follows.
Prior knowledge of environments: IRM assumes training data are partitioned by environments andthe environment labels (i.e. which data instance belongs to which environment) are given.
However,the environment labels are often unavailable.
Moreover, the definition of environments can be implicit,making human labeling more difficult and expensive.
To find environments without supervision,
Creager et al.
[6] propose EIIL, a min-max optimization framework training the model of interestand inferring the environment labels.
Another work, HRM [17] (or the extension, KerHRM [18]),parameterizes the environments and proposes clustering-based approaches to estimate the parameters.
A recent method, ZIN [15], also learns to label data.
However, it relies on carefully chosen featuressatisfying a series of theoretical constraints, and thus human efforts are still required.
Delicate initialization: EIIL is able to infer environments but requires an ERM model for initialization.
Crucially, the ERM model has to be underfitted and heavily depend on spurious correlations.
Asthe distributional shift is assumed to be unknown in the training stage, appropriate initialization isdifficult to guarantee.
Scalability issue: HRM and KerHRM, though do not possess the above two limitations, suffer fromthe scalability issue.
Specifically, HRM is assumed to be trained by low dimensional data. KerHRMas
the extension of HRM avoids the issue of dimensions by adopting kernel methods.
However, thetraining cost of the kernel methods can be very high if the data or pre-trained model size is large.
This work proposes a novel framework, E nvironment D iversification with multi-head N eural networkfor I nvariant L earning (EDNIL).
EDNIL is able to infer environment labels without supervisionand be jointly optimized with an invariant learning model.
The underlying neural network explicitlydiversifies the inferred environments, which is consistent with a recent study [5] revealing the benefitsof diverse environments.
Notably, the structure of the proposed neural network is similar to amulti-class classifier and can be optimized efficiently.
The advantages of EDNIL are summarized as:
• We implement this framework using various pre-trained models such as Resnet [12] and
DistilBert [24] and evaluate it with diverse data types and varied biases
The results
show that
EDNIL can constantly outperform the existing state-of-the-art models.
• The learning algorithm of EDNIL has theoretical connections to recent studies [5, 15, 17, 18]
discussing conditions of ideal environments.
• EDNIL does not have the three limitations discussed above.
The comparison between EDNIL
and other models is shown in Table
[Table]
[Equation]
✗ ✓
[Equation]
✓ ✓ ✗
[Equation]
[Equation] ✓ ✓ ✗ Ours ✓ ✓ ✓  Preliminaries and related works
The goal of EDNIL is to tackle out-of-distribution problem with invariant learning in the absenceof manual environment labels.
In Section 2.1, background knowledge about out-of-distributiongeneralization and invariant learning are introduced.
In Section 2.2, we discuss the assumptions ofvariant features and desired environments for invariant learning.
In Section 2.3, we introduce theexisting invariant learning methods that do not require human labelling.
2.1 Out-of-distribution Generalization and Invariant Learning
Following [1, 17, 18], we consider a dataset D
= { D e } e ∈ supp( E tr ) with different sources D e ={ ( x ei , y ei ) } n e i =1 collected under multiple training environments e ∈ supp( E tr ) .
Random variable E trrepresents the training environment labels.
For simplicity, X e , Y e and P e denote data, target anddistribution in environment e respectively.
With E all containing all possible indices of environments such that E all ⊃ E tr , the goal of outof-distribution generalization is to learn a predictor f ( · ) :
X → Y as Equation 1.
R e :=
E e [ l ( f ( X e ) , Y e )] measures the risk under environment e , where l ( · , · ) is the loss function.
In general,for e ∈ supp( E tr ) and e ′ ∈ supp( E all ) \ supp( E tr ) , P e ′ ( X e ′ , Y e ′ ) is different from P e ( X e , Y e ) .
[Equation]
Recently, several studies [1, 20, 22, 14, 4] attempt to tackle the generalization problem by discoveringinvariant relationships across supp( E all ) .
A commonly proposed invariance assumption is as follows: Assumption 2.1 There exists a random variable X c such that P e ( Y | X c ) = P e ′
( Y | X c ) ∀ e, e ′ ∈supp( E all ) .
X c , which is called invariant features , induces consistent conditional distributions over target Y acrossenvironments e ∈ supp( E all ) .
To learn invariant predictor Φ that purely encodes the informationof X c , IRM [1] introduces a regularization term enforcing simultaneous optimality of the classifierw ◦ Φ with E tr , where dummy classifier w = 1 . 0 is a fixed multiplier for each output dimension:(cid:88)
[Equation]
2.2 Environment Inference
Following [17, 18, 15], the existence of variant features X v is assumed as follows:
Assumption 2.2 Given X and X c satisfying Assumption 2.1, there exists a random variable X vsuch that X
= h ( X c , X v ) , where h ( · ) is a transformation function. The distribution P e ( Y | X v ) canvary across environments e ∈ supp
( E all ) .
Given a mixture of training data D , E tr is unavailable or sub-optimal for invariant learning in mostapplications [17], which limits the training of IRM. To facilitate invariant learning, our goal is to inferenvironments E learn for IRM to elicit better invariant predictor.
With Assumption 2.1 and 2.2, X cand X v are more identifiable when the dependency between X c and Y is consistent across E learn ,and that between X v and Y is diversified [17, 18, 15]. Therefore, our task is defined as Problem 2.1.
Problem 2.1 Given a dataset D
= { D e } e ∈ supp( E tr ) , the task is to generate E learn satisfying thefollowing two conditions regarding conditional entropy H and learn an invariant predictor with Equation 2 for better out-of-distribution generalization: [Equation]
2.3 Invariant learning methods with unsupervised environment inference
Here we provide a more detailed introduction of the existing unsupervised invariant learning methods:
EIIL [6], HRM [17] and KerHRM [18].
EIIL proposes formulating invariant learning as a min-maxoptimization problem.
Specifically, EIIL is composed of two objectives, EI and IL , where EI isoptimized by maximizing the training loss via labeling the data, and IL is optimized by minimizingthe training loss given the data labeled by EI .
The two-stage algorithm bypasses the difficulty ofdefining environments; however, the training result heavily relies on the initialization of the EIoptimization.
Empirically, the initialization demands an underfitted and biased ERM model, or EIILcan have a significantly weaker performance than ERM.
Another method, HRM, proposes a clustering-based method for learning plausible environments.
HRM assumes that data in each environment can be modeled by a parameterized function and thedataset is generated by the mixture of the functions.
The parameters are then estimated by employing EM algorithm.
Additionally, HRM equips a jointly learning framework which alternatively learnsinvariant predictors and improves quality of clustering results.
A known issue of HRM is anassumption that the data are represented by raw features; namely, data such as images and textrequiring non-linear neural networks to obtain representations are beyond the capability.
To extend HRM to a broader class of applications and improve the model performance, Liu et al.[18] proposes KerHRM.
The main idea is to adopt the Neural Tangent Kernel [13] method whichtransforms non-linear neural network training into a linear regression problem on the proposed Neural
Tangent Features space.
As a result, KerHRM elegantly resolves the shortcoming of HRM and isshown to be more robust and effective.
However, the kernel method also brings additional trainingcost depending on data size and pre-trained models.
In particular, for applications favoring largepre-trained models such as Resnet [12] and BERT [7], KerHRM may not be an affordable option.
Methodology
In this section, we propose a general framework to learn invariant model without manual environmentlabels.
As shown in Figure 1, our proposed method consists of two models, M EI for environmentinference and M IL for invariant learning.
Given the pooled data ( X, Y ) , M EI generates environments
E learn satisfying Condition 3 and 4, and M IL is an invariant model trained with inferred environments.
Our framework is jointly optimized with alternating updates.
The learned invariant model can provideinformation of X c to M EI , so that Condition 3 and 4 can be fulfilled simultaneously.
3.1 The graphical model and learning objectives
We propose a graphical model which is a sufficient condition (the proof is in the supplementary file
Section A) for Condition 3 and 4, and is the foundation of M EI in EDNIL. The graph is shown in
Figure 2. Based on the structure of the graphical model, we propose two objectives for learning tofind effective environments E learn . Label dependency objective: The target Y and inferred environments E learn are independent.
Environment diversification objective: H ( Y | X v ) − H ( Y | X v , E learn ) > 0 . (i.e. Condition 4)
Additionally, if jointly learning invariant representation Φ( X ) ∈ X c and environment E learn isdesired, we propose including and optimizing the third objective: Invariance preserving objective: H ( Y | Φ( X )) − H ( Y | Φ( X ) , E learn ) = 0 .
(i.e. Condition 3)
The inference model, M EI , is a deep neural network learning to optimize the objectives proposed in Section 3.1. Following the idea of HRM [17] and KerHRM [18], we assume the mapping relationbetween X and Y in environment e can be modelled by a function
f e (Ψ( X )) , where Ψ( X ) islearned representations expected to encode variant features X v .
Instead of taking f e (Ψ) as a centerof a cluster, we propose combining the environmental functions to build a multi-head network asshown in Figure 3.
Specifically, a cluster in HRM [17] or KerHRM [18] corresponds to a single-headneural network in M EI with shared parameters.
The training procedure of M EI can be divided intoinference stage and learning stage .
Inference stage The goal is to infer the environment of each training data.
Equation 5 is theproposed method, where the probability P ( e | X, Y ) is estimated via a softmax of negativel ( f e (Ψ( X ) , Y )) with a hyper-parameter τ .
The function l is expected to be the commonly used crossentropy or mean squared error loss, measuring the discrepancy between f e (Ψ( X )) and Y .
[Equation]
Learning stage The goal is to update the neural network to improve the quality of inference.
Three losses are minimized for our three objectives. The strength of each loss can be controlled byhyper-parameters β and γ .
[Equation]
Given the estimated P ( E learn | X, Y ) , L ED selects the most probable environment and its corresponding network for optimization:
[Equation]
Although max e ∈ supp( E learn ) P ( e | X, Y ) selects only one environment for the minimization, thesoftmax simultaneously propagates gradient to maximize l ( f e ′ (Ψ( X ) , Y ))[-for e 0 
 = e i-] .
As a result, the networklearns to predict Y by relying on both Ψ( X ) and E learn , consistent with the environment diversification objective that H ( Y | X v ) − H ( Y | X v , E learn ) > 0 .
In practice, we utilize rescaling weightw i inversely proportional to the size of e i .
The importance of smaller environments will be thusenhanced within summation.
The second term, L LD , constraints the mutual information I ( Y ; E learn ) measuring the dependencybetween Y and E learn .
With satisfying the label dependency objective, L LD also prevents a naivesolution that environments are purely split by labels without considering X v , which is undesirable forinvariant learning.
To minimize I ( Y ; E learn ) , it can be verified that it is equivalent to minimizing Equation 8 given P ( Y ) is known.
[Equation]
The last term, L IP , promotes E learn to invariance preserving objective as M IL jointly learns invariantfeatures to some extent.
To prevent existing invariant relationship from being diversified by L ED , itrestricts the variance of expected loss from invariant model Φ between environments.
[Equation]
[Equation]
In addition, before minimizing L EI , we pre-train our Ψ and one arbitrary f e with ERM.
In general,it empirically facilitates better feature extraction.
Unlike [6] taking ERM as reference model heavilyrelying on spurious features, EDNIL performs consistently under various choices of ERM.
Inparticular, it is not restricted to under-fitted ERM training.
We verify the argument in Section 4.
3.3 The invariant learning model
Invariant Learning identifies invariance across given environments.
For our learning model, M IL , IRM [1] is selected as our base algorithm.
We assign environment label e ∈ supp( E learn ) with largest P ( e | x i , y i ) to each data ( x i , y i ) for training.
However, it is inevitable that there exist some noises inautomatically inferred environments, especially in the beginning of joint optimization.
To reduce theimpact of immature environments on invariant learning, we calculate the confidence score c e for eachenvironment e ∈ supp( E learn ) , i.e. P ( e | X e , Y e ) .
Our training objective is modified to minimize theweighted average of the environment risks:
[Equation] [Equation]
Experiments
We empirically validate the proposed methods on four biased datasets, Adult-Confounded, CMNIST,
Waterbirds and SNLI.
The generation of spurious correlations mainly follow the protocols proposed by[6, 18, 23, 8].
In Section 4.1, Adult-Confounded and CMNIST are tested with multilayer perceptron(MLP).
In Section 4.2, two more complex datasets, Waterbirds and SNLI, are taken for evaluating theintegration of representation learning.
Modern deep models with pretrained weights will serve as theencoders of variant and invariant features.
For hyper-parameter tuning, we split 10% of training data as in-distribution validation set. In eachdataset, several testing subsets with different distributions are listed to evaluate the robustness ofeach method, and we mainly take worst-case performance for assessment.
† will be marked thedifference is statistically significant value -best result.
• Empirical Risk Minimization (ERM): ERM chooses model parameters that minimize theaverage loss on the training set.• Invariant Risk Minimization (IRM [1]): Given environments, IRM seeks the invariancerelationship with Equation 2.
As indicated in [5, 18, 8], larger discrepancy of spurious correlations between environments benefits invariant learning more. Therefore, unlike [6, 17, 18]who take sub-optimal E tr for IRM, we re-label E oracle on given biased training set, which ex-
tremely diversifies the correlations between variant features and target, to provide upper-boundperformance with IRM. • Environment Inference for Invariant Learning (EIIL [6]): EIIL partitions environments bymaximizing Equation 2 with ERM as reference model. Subsequently, IRM is taken to learninvariant predictor with inferred environments. • Kernelized Heterogeneous Risk Minimization (KerHRM [18]): KerHRM jointly exploresheterogeneity by clustering algorithm and learns invariant parameters in kernel space.
4.1 Simple datasets with MLP
This section includes two simple datasets, Adult-Confounded and CMNIST, where spurious correlations are synthetically produced with the pre-defined strength.
For all competitors, MLP is takenas the base model and full-batch training is implemented.
Since KerHRM performs unstable overrandom seeds, we first average the results after 10 runs as its first score, and select top-5 among themas the second one, which will be marked with a star ( ∗ ) in each table.
4.1.1 Discussion on Adult-Confounded dataset
We take UCI Adult dataset 1 to predict binarized income levels (above or below $50,000 per year)based on demographic information.
Following [6], individuals are re-sampled according to sensitivefeatures race and sex to simulate spurious correlations.
Specifically, with binarized race (Black/NonBlack) and sex (Female/Male), four possible subgroups are constructed: Non-black Males (SG1), Non-black Females (SG2), Black Males (SG3), and Black Females (SG4).
Keeping original train/testsplit and subgroup size from UCI Adult, we sample data for the given label distributions in eachsensitive subgroup as shown in Table 2.
In this task, MLP with one hidden layer of 96 neurons isconsidered.
For IRM, four environments comprise E oracle , where the correlations between variantfeatures ( race, sex ) and target Y are distributed without overlapping.
More details are provided inthe supplementary file Section B.
Results The results are shown in Table 3.
First of all, notsurprisingly we found ERM adopts poorly to distribution shift.
Among all invariant learning methods without prior environment labels, EDNIL can perfectly identify variant features andgenerate diversified environments.
Therefore, it achieves themost invariant performance over different testing distributions.
EIIL can improve consistency to some degree, but not as strongas EDNIL.
A possible explanation is that some predictable invariant features risk being diversified across environments andbeing discarded by invariant learning since empirically trainedreference model is not guaranteed to be purely variant [6].
For KerHRM, it performs inconsistently across random seeds,which can be reflected on large standard deviation.
In somecases, the performance hardly improves over iterations, as alsoobserved in Liu et al.
Ablation Study for M EI
We first claim the importance of constraining label dependency byremoving the coefficient β of L LD .
As discussed in Section 3, the resulting environments are purelysplit by labels, which leads to inferior performance for invariant learning as shown in Table 3.
Figure4 shows the effectiveness of joint optimization.
The regularization L IP promotes environmentinference, so that the worst-case performance improves and remains stable over iterations.
If thecoefficient γ is turned off, feedback generated by M IL will be ignored and the effect of invariant
[Table]
Three subsets are defined in Table[-2.  = 
 and-]
[Table]
learning risks being undesirable.
In Table 3, the degradation occurs on all subsets.
It indicatesthat some invariant features are eliminated since Condition 3 cannot be realized in the splitting ofenvironments [Table]
4.1.2 Discussion on CMNIST dataset
We report our evaluation on a noisy digit recognition dataset, CMNIST.
Following [1], we first assign
Y = 0 to those whose digits are larger than 5 and Y = 1 to the others.
Next, we apply label noise byrandomly flipping Y with probability 0.2.
Finally, the digits are colored with color label C , which isgenerated by randomly flipping Y with probability e .
For training, e is set to 0.15, which is equivalentto merging two training environments from [1].
For testing, three situations are considered when e is0.1, 0.5 or 0.9 respectively.
Note that when e = 0 .
1 , the spurious correlation is much aligned withtraining set.
On the other hand, e = 0 .
9 defines the most challenging scenario since the correlationshifts most dramatically from training.
[Table]
[Table]
For all competitors except KerHRM, weselect MLP with two hidden layers of390 neurons, and consider the wholedataset (50,000 samples) for training.
Due to massive computing resources required by KerHRM, we follow the settings recommended by [18].
Specifically, we randomly select 5,000 samplesand train MLP with one layer of 1024neurons.
To construct most diversified
E oracle for IRM, we pack all exampleswith C =
Y into one environments, and
C [-
-]̸
= Y into the other.
Results The results are shown in Table 4.
With large spurious correlation at train time, ERMobtains high accuracy as the correlation remains aligned; however, its generalization to other testingdistribution is limited.
Among all invariant learning methods without manual environment labels,
EDNIL gets closest to IRM with E oracle , achieving consistent and robust performance in this dataset.
Number of Environments As shown in Figure 5, EDNIL is not sensitive to predefined numberof environments.
Specifically, when the environment number is larger than the oracle (i.e. 2), someenvironment classifiers become redundant.
Each of them owns a high constant loss, taking up a fixedand ignorable space in softmax function.
The comparison between E oracle and E learn with size of5 is visualized in Figure 6.
Additionally, training EDNIL with more environments is much moreefficient than clustering-based methods, such as the one proposed in HRM and KerHRM.
4.2 Complex datasets with pre-trained deep learning models
This section extends MLP to deep learning models with pre-trained weights on more complex data.
With mini-batch fine-tuning, we consider all competitors except KerHRM due to its limitation todeep representation learning [18].
In Section 4.2.1, image dataset, Waterbirds [23], with controlledspurious correlations is selected for evaluating the generalization on more high-dimensional images.
In Section 4.2.2, a real-world NLP dataset, SNLI [3], is considered.
The biases in SNLI are naturallyderived from the procedure of data collection, and we define biased subsets for evaluation following
Dranker et al.
4.2.1 Discussion on Waterbirds dataset
In Waterbirds [23], each bird photograph, from CUB dataset [25], is combined with one backgroundimage, from Places dataset [26].
Both birds and backgrounds are either from land or water, and ourtarget is to predict the species of birds.
At training time, landbirds and waterbirds frequently presentin land and water backgrounds respectively.
Therefore, empirically trained models are prone to learncontext features, and fails to generalize as background varies [2, 6, 9, 16, 23].
To split validation set so that its distribution is i.i.d.
to training set, we merge original trainingand validation data and split 10% for hyper-parameter tuning.
For testing, we observe all fourcombinations of birds and backgrounds in the original testing set.
Among them, the minor subgroup(waterbirds on land) contributes the most challenging scenario.
In this task, Resnet-34 [12] ischosen with mini-batch fine-tuning.
Given two binary labels ( target, background ) , we distributetarget = background and target ̸ = background into two different environments and applybalanced class weights for the oracle settings of IRM.
Results The results are shown in Table 5.
As observed in [6, 23], ERM suffers in the hardest scenario(i.e.
waterbirds on land).
EIIL also performs poorly in this scenario.
With more sophisticated learningframework, EDNIL narrows the gaps between subgroups and uplifts the worst-case performance.
The results show that EDNIL is more resistant to distribution shifting.
Choice of Initialization Both EIIL and EDNIL take ERM as initialization.
As mentioned in Section1, heavy dependency on initialization is risky when testing distribution is unavailable.
Therefore, wetake ERM with different training steps for EIIL and EDNIL to verify the stability.
The results areshown in Figure 7.
As suggested by [6], EIIL works only with underfitted reference model (trainingsteps fewer than 400).
If the reference model is well-fitted, the performance of EIIL will greatlydecline since ERM learns beyond variant features.
One can be misled into undesirable choice for
EIIL when seeking hyper-parameters with in-distribution validation set.
For instance, the validationscore of EIIL with 500-step reference model (96.9%) is much higher than that with 100-step (94.9%),which is not consistent with their performances on testing.
In comparison, EDNIL remains consistentacross different pre-trained steps, which accentuates our strength of effortless initialization .
4.2.2 Discussion on SNLI dataset
The target of SNLI [3] is to predict the relation between two given sentences, premise and hypothesis.
Recent studies [11, 19, 21] reveal hypothesis bias in SNLI, which is characterized by patterns inhypothesis highly correlated with a specific label.
One can achieve low empirical risk without
The subgroup (Water, Land) contributes the worst-case performance.
[Table]
considering premises during prediction.
However, as the bias no longer holds, the performancedegradation occurs [10, 19].
We sample 50,000 examples and select two major classes, entailment and contradiction , for ourexperiment.
Following [8], we define three subsets, unbiased, bias aligned and bias misaligned, bytraining a biased model with hypothesis as its only input.
Specification of the subsets is as follows:
• Unbiased: Examples whose predictions from biased model are ambiguous.
Aligned:
Examples that biased model can predict correctly with high confidence.
Misaligned: Examples that biased model can predict incorrectly with high confidence.
The proportions of the three subsets are 22%, 67% and 11% respectively.
Due to the minority ofmisaligned bias subset, it is more likely to be ignored and thus defines the worst-case performance.
[Table]
For all methods, DistilBERT [24] is taken as thepre-trained model for further mini-batch finetuning.
For E oracle , we assign the aligned biassubset to the first environment, and the misaligned bias subset to the second.
In order tomake bias prevalence equal in the two environments, samples in the unbiased subset are scattered according to their sizes.
Results Table 6 shows our results with thethree subsets.
As reported by [8], ERM receivesremarkable score on the bias aligned subset, butit fails on the bias misaligned case.
Althoughboth EIIL and EDNIL degrade on the unbiased and bias aligned subsets, only EDNIL improvessignificantly on the bias misaligned subset.
Even though the definition of biases are in high level, EDNIL is still capable of encoding and diversifying possible variant features.
Limitation
The learning algorithm of EDNIL is based on the graphical model plotted in Figure 2.
As data arenot necessarily generated by the presumed process, there can exist biases that cannot be captured bythe proposed neural network.
In the paper, we provide empirical studies of model effectiveness ondiverse datasets, while we are aware that a stronger guarantee of performance is required.
Conclusion
This work proposes EDNIL for training models invariant to distributional shifts.
To infer environmentswithout supervision, we propose a multi-head neural network structure with a two-stage learningprocedure to identify and diversify plausible environments.
With joint optimization, the trainedmodels are shown to be more robust than existing solutions on data having distinct characteristicsand different strengths of biases.
We attribute the effectiveness to the underlying learning objectivesconsistent with recent studies of ideal environments.
Additionally, we note that the classifier-likestructure makes EDNIL easy to combine with off-the-shelf pre-trained models and trained efficiently.