Normalizing Flows for Knockoff-free Controlled Feature Selection
Abstract
Controlled feature selection aims to discover the features a response depends on while limiting the false discovery rate (FDR) to a predeﬁned level.
Recently, multiple deep-learning-based methods have been proposed to perform controlled feature selection through the Model-X knockoff framework.
We demonstrate, however, that these methods often fail to control the FDR for two reasons.
First, these methods often learn inaccurate models of features.
Second, the “swap” property, which is required for knockoffs to be valid, is often not well enforced.
We propose a new procedure called F LOW S ELECT to perform controlled feature selection that does not suffer from either of these two problems.
To more accurately model the features, F LOW S ELECT uses normalizing ﬂows, the state-of-the-art method for density estimation.
Instead of enforcing the “swap” property, F LOW S ELECT uses a novel
MCMC-based procedure to calculate p-values for each feature directly.
Asymptotically, F LOW S ELECT computes valid p-values.
Empirically, F LOW S ELECT consistently controls the FDR on both synthetic and semi-synthetic benchmarks, whereas competing knockoff-based approaches do not.
F LOW S ELECT also demonstrates greater power on these benchmarks.
Additionally, F LOW S ELECT correctly infers the genetic variants associated with speciﬁc soybean traits from GWAS data.
Introduction
Researchers in machine learning have made much progress in developing regression and classiﬁcation models that can predict a response based on features.
In many application areas, however, practitioners need to know which features drive variation in the response, and they need to do so in a way that limits the number of false discoveries.
For example, in genome-wide association studies (GWAS), scientists must consider hundreds of thousands of genetic markers to identify variants associated with a particular trait or disease.
The cost of false discoveries (i.e., selecting variants that are not associated with the disease) is high, as a costly follow-up experiment is often conducted for each selected variant.
Another example where controlled feature selection matters is analyzing observational data about the effectiveness of educational interventions.
In this case, researchers may want to select certain educational programs to implement on a larger scale and require conﬁdence that their selection does not include unacceptably many ineffective programs.
As a result, researchers are interested in methods that model the dependence structure of the data while providing an upper bound on the false discovery rate (FDR).
Model-X knockoffs (Candès et al., 2018) is a popular method for controlled variable selection, offering theoretical guarantees of FDR control and the ﬂexibility to use arbitrary predictive models.
However, even with knowledge of the underlying feature distribution, the Model-X knockoffs method is not feasible unless the feature distribution is either a ﬁnite mixture of Gaussians (Gimenez et al., 2019) or has a known Markov structure (Bates et al., 2020).
Hence, a body of research explores the use of empirical approaches that use deep generative models to estimate the distribution of X and sample knockoff features (Jordon et al., 2019; Liu & Zheng, 2018;
Romano et al., 2020; Sudarshan et al., 2020).
The ability of these methods to control the FDR is contingent on their ability to correctly model the distribution of the features.
By itself, learning a sufﬁciently expressive feature model can be challenging.
However, the knockoff procedure requires learning a knockoff distribution that satisﬁes the swap property , which is a much stronger requirement.
Formally, let X ∈ R D be a sample from the feature distribution and ˜ X ∈ R D be a sample from the knockoff distribution conditioned on X .
The swap property stipulates that the joint distribution ( X, ˜ X ) ∈ R 2 D must be invariant to swapping the positions of any subset of features S ∈ { 1 , . . .
, D } :
[Equation]
Here, swap ( S ) means exchanging the positions of X j and ˜
X
j for all j ∈ S .
For example, in the case D = 3 and S = { 1 , 3 } , the joint distribution is ( X, ˜ X ) =
( X 1 , X 2 , X 3 , ˜ X 1 , ˜ X 2 , ˜ X 3 ) , and the swapped joint distribution is ( X, ˜ X ) swap ( S )
=
( ˜ X 1 , X 2 , ˜ X 3 , X 1 , ˜ X 2 , X 3 ) .
Note that, for
S = { 1 , . . .
, D } , the swap property implies that ˜ X D
= X .
See Candès et al.
for a more detailed description of the swap property.
Even if a distribution were found satisfying the swap property, it may not provide enough power to make discoveries.
For example, both properties are trivially satisﬁed by constructing exact copies of the features as knockoffs, but the resulting procedure has no power.
In situations where a valid knockoff distribution is available to sample from, knockoffs are computationally appealing because they require only one sample from a knockoff distribution to assess the relevance of all p features.
However, in situations where the joint density of the features is unknown, we show that empirical approaches to knockoff generation (Jordon et al., 2019; Liu & Zheng, 2018;
Romano et al., 2020; Sudarshan et al., 2020) fail to characterize a valid knockoff distribution and therefore do not control the FDR.
We further show that even with a known covariate model, it is not straightforward to construct a valid knockoff distribution unless a speciﬁc model structure is known.
We propose a new feature selection method called F LOW S ELECT (Section 3), which does not suffer from these problems.
F LOW S ELECT uses normalizing ﬂows to learn the joint density of the covariates.
Normalizing ﬂows is a state-of-the-art method for density estimation; asymptotically, it can approximate any distribution arbitrarily well (Papamakarios et al., 2021;
Kobyzev et al., 2020; Huang et al., 2018).
Additionally, F LOW S ELECT circumvents the need to sample a knockoff distribution by instead applying a fast variant of the conditional randomization test (CRT) introduced in Candès et al.
Samples from the complete conditionals are drawn using MCMC, ensuring they are unbiased with respect to the learned data distribution.
Asymptotically, F LOW S ELECT computes correct p-values to use for feature selection (Section 4).
Our proof assumes the universal approximation property of normalizing ﬂows and the convergence of MCMC samples to the Markov chain’s stationary distribution.
Under the same assumptions as the CRT, which includes a multiple-testing correction as in Benjamini & Hochberg (1995), a selection threshold can be picked which controls the FDR at a pre-deﬁned level.
Empirically, on both synthetic (Gaussian) data and semi-synthetic data (real predictors and a synthetic response), F LOW S ELECT controls the FDR where other deep-learning-based knockoff methods do not.
In cases in which competing methods do control the FDR, F LOW S ELECT shows higher power (Section 5).
Finally, in a challenging real-world problem with soybean genome-wide association study (GWAS) data, F LOW S ELECT successfully harnesses normalizing ﬂows for modeling discrete and sequential GWAS data, and for selecting genetic variants the traits depend on (Section 5.4).
Background
F LOW S ELECT brings together four existing lines of research, which we brieﬂy introduce below.
Normalizing ﬂows Normalizing ﬂows is a general framework for density estimation of a multidimensional distribution with arbitrary dependencies (Papamakarios et al., 2021).
A normalizing ﬂow starts with a simple probability distribution (e.g., Gaussian or uniform), which is called the base distribution and denoted Z , and transforms samples from this base distribution through a series of invertible and differentiable transformations, denoted G , to deﬁne the joint distribution of
X ∈ R D ∼ P X .
A normalizing ﬂow with enough transformations can approximate any multivariate density, subject to regularity conditions detailed by Kobyzev et al.
Compared to other densityestimation methods, normalizing ﬂows are computationally efﬁcient.
Details about the speciﬁc normalizing ﬂow architecture used in F LOW S ELECT are provided in Appendix A.
Controlled feature selection Consider a response Y which depends on a vector of features X ∈ R D .
Depending on how the features are chosen, it is plausible that only a subset of the features contains all relevant information about Y .
Speciﬁcally, conditioned on the relevant features in X ,
Y is independent of the remaining features in X (i.e. the null features).
The goal of the controlled feature selection procedure is to maximize the number of relevant features selected while limiting the number of null features selected to a predeﬁned level.
If we denote the total number of selected features R , then we can decompose R into V , the number of relevant features selected, and S , the number of null features selected.
Conditional randomization test Controlled feature selection can be seen as a multiple hypothesis testing problem where there are p null hypotheses, each of which says that feature X j is conditionally independent of the response Y given all the other features X − j .
Explicitly, the test of the following hypothesis is conducted for each feature
j = { 1 , . . .
, D } :
[Equation]
To test these hypotheses, one can use a conditional randomization test (CRT) (Candès et al., 2018).
For each feature tested in a conditional randomization test, a test statistic T j (e.g., the LASSO coefﬁcient or another measure of feature importance) is ﬁrst computed on the data.
Then, the null distribution of T j is estimated by computing its value ˜ T j based on samples ˜
X j drawn from the conditional distribution of X j given X − j .
Finally, the p-value is calculated based on the empirical
CDF of the null test statistics, and features whose p-values fall below the threshold set by the
Benjamini-Hochberg procedure (Benjamini & Hochberg, 1995) are selected.
Though the CRT is introduced as a computationally inefﬁcient alternative to knockoffs, the CRT nonetheless has appeal because it requires only knowledge of the feature distribution, which can be learned empirically by maximum likelihood.
Holdout randomization test The holdout randomization test (HRT)
(Tansey et al., 2021) is a fast variant of the CRT; it uses a test statistic that requires ﬁtting the model only once.
Let θ represent the parameters of the chosen model, and let T ( X, Y, θ ) be an importance statistic calculated from the model with input data.
For example, T , could be the predictive likelihood P θ ( Y test | X test ) or the predictive score R 2 .
To use the HRT, ﬁrst ﬁt model parameters ˆ θ based on the training data.
Next, for each covariate j , calculate the test statistic T ∗ j ← T ( X test , Y test , ˆ θ ) .
Then, generate k null samples and compute T j,k ← T ( X test ( j ← j k ) , Y test , ˆ θ ) , where X test ( j ← j k ) replaces the j -th covariate with the k -th generated null sample.
Finally, calculate the p-value as in the CRT, based on the empirical CDF of the null test statistics.
Methodology
F LOW S ELECT implements the CRT for arbitrary feature distributions by using a normalizing ﬂow to ﬁt the feature distribution and Markov chain Monte Carlo (MCMC) to sample from each complete conditional distribution.
Performing controlled feature selection with F LOW S ELECT consists of the three steps below.
Step 1: Model the predictors with a normalizing ﬂow
Starting with the observed samples of the features X 1 , . .
, X N ∼ P X , we ﬁt the parameters of a normalizing ﬂow G θ to maximize the log likelihood of the data with respect to a base distribution p Z :
[Equation]
N
[Equation]
where p θ ( X i ) = p
[Equation]
G
[Equation]
X [Equation] [Equation]
det
[Equation] [Equation] The resulting density p ˆ θ is a ﬁtted approximation to the true density P X .
The speciﬁc normalizing ﬂow architecture we use in our ﬁrst two experiments consists of a single Gaussianization layer (Meng et al., 2020) followed by a masked autoregressive ﬂow (MAF) (Papamakarios et al., 2017).
The ﬁrst layer can learn complex marginal distributions for each covariate, while the MAF learns the dependencies between them.
More detail on normalizing ﬂows and on this particular architecture can be found in Appendix A.
Step 2: Sample from the complete conditionals with MCMC
For each feature j , we aim to sample corresponding null features
˜ X
i,j,k for all k ∈ { 1 , . . .
, K } that are equal in distribution to p ˆ θ ( X i,j | X i, − j ) , but independent of Y i .
However, directly sampling from this conditional distribution is intractable.
Instead, we implement an MCMC algorithm that admits it as a stationary distribution.
The samples drawn from MCMC are autocorrelated, but any statistic calculated over these samples will converge almost surely to the correct value.
The choice of the
MCMC proposal distribution q j is ﬂexible.
Because each Markov chain is only one-dimensional, a
Metropolis-Hastings Gaussian random walk with the standard deviation set based on the covariance can be expected to mix rapidly.
Alternatively, information from p ˆ θ , such as higher-order derivatives, could be used to construct a more efﬁcient proposal.
Algorithm 1 details how to implement step 2.
Step 3: Test for signiﬁcance with the HRT
As in the CRT, feature j has high evidence of being signiﬁcant if, under the assumption that j is a null feature, the probability of realizing a test statistic greater than the observed T j ( X ) is low.
Formally, letting [ ˜ X j , X − j ] be the observed feature matrix with the observed feature X j swapped out with the null feature ˜ X j , we can write this as a p-value α j :
[Equation]
However, the above p-value α j is not tractable.
For each sample ˜ X · ,j,k drawn using MCMC, we calculate the corresponding feature statistic and compare it to the real feature statistic, leading to an Algorithm 1 Step 2 of the F LOW S ELECT procedure for drawing K null features ˜ X
i,j | X i, − j for feature j at observation i .
Input: Feature matrix X ∈ R N × D , observation index i , feature index j , number of samples K , ﬁtted normalizing ﬂow p ˆ θ , MCMC proposal q j
Output: Null features ˜
X i,j,k for k = 1 , . . . , K for k = 1 , . . . K do
Propose: X (cid:63)i,j,k ∼ q j ( ·| ˜ X
i,j,k − 1 , X i, − j ) r i,j,k ← p ˆ θ ( X (cid:63)i,j,k ,X i, − j ) q j ( ˜ X i,j,k − 1 | X (cid:63)i,j,k ,X i, − j ) p ˆ θ ( ˜ X i,j,k − 1 ,X i, − j ) q j ( X (cid:63)i,j,k | ˜
X i,j,k − 1 ,X i, − j )
Sample: U i,j,k ∼ Bernoulli ( r i,j,k ∧ 1) if U i,j,k = 1 then ˜
X i,j,k ← X (cid:63) i,j,k else ˜
X
i,j,k ← ˜ X i,j,k −end if end for approximated p-value ˆ α j : [Equation] To control the FDR, we use the Benjamini-Hochberg procedure to establish a threshold for the observed p-values.
Speciﬁcally, we set the threshold to s ( γ ) (cid:44) max j { ˆ α j : ˆ α j ≤ jD γ } , and select all features j such that α j ≤ s ( γ ) .
The Benjamini-Hochberg correction only guarantees FDR control provided that the p-values have either positive or zero correlation.
Thus, the FDR control of F LOW S ELECT depends on these assumptions being met.
A more conservative correction from Benjamini & Yekutieli (2001) allows for arbitrary dependencies in p-values, but it suffers from low power.
The Benjamini-Hochberg correction is widely used and empirically robust (Tansey et al., 2021), so we report results using it.
Across our synthetic and semi-synthetic benchmarks in Section 5, we also ﬁnd that F LOW S ELECT maintains empirical FDR control.
Provided that the Benjamini-Hochberg assumptions are met, the FDR will be controlled, but the power of the test depends on T j being higher when j is a signiﬁcant feature.
For example, if Y is expected to vary approximately linearly with respect to X , T j ( X ) could be the absolute estimated regression coefﬁcient | ˆ β j | for the linear model Y =
Xβ + (cid:15) .
Another choice is the HRT feature statistic described earlier.
Asymptotic results
The ability of F LOW S ELECT to control the FDR relies on its ability to produce estimated p-values that converge to the correct p-values for the hypothesis test in Equation (2).
Theorem 1.
Let X ∈ R N
×
D be a random feature matrix, where each row X i, · is independent and identically distributed; x ∈ R N × D be the observed feature matrix; and α j be the p-value as deﬁned in Equation (4) with test statistic T j ( X ) .
Suppose there exists a sequence of functions ( G n ) ∞ n =1 and a base
random variable
Z satisfying the following conditions:
Each G n is continuously differentiable and invertible.
G n → G pointwise for some map G that is triangular, increasing, continuously differentiable, and satisﬁes G
( X i, · ) D =
Z .
For n = 1 , 2 , . . .
, let X n be the random feature matrix where each row i is independent and has distribution X ni, · = ( G n ) − 1 ( Z ) .
Then, the p-value in Equation (5) calculated using K MCMC samples targeting X n · ,j | X n · , − j = x · , − j converges to the correct p-value α j with probability 1 .
Here we sketch the proof.
A full proof can be found in Appendix B. First, by construction each G n deﬁnes a distribution X ni, · D = ( G n ) − 1 ( Z ) that in turn implies a conditional distribution X n · ,j | X n · , − j = x · , − j .
We show these conditional distributions converge to the true conditional distribution of X · ,j given X · , − j = x · , − j .
Consequently, the probability of observing a higher test statistic under the approximated null distribution ˜ X n · ,j D = X n · ,j , written α nj , will converge to the probability under the true null distribution ˜ X · ,j | X · , − j = x · , − j , i.e. α j .
Next, the Cesaro average of K samples from an
MCMC algorithm targeting ˜ X n · ,j | X · , − j = x · , − j , written ˆ α j,K,n will converge to α nj with probability 1 as K → ∞ .
Combining these two convergences leads to the stated result.
Assuming the limiting p-values { α j } satisfy the chosen multiple-hypothesis-testing assumptions, Theorem 1 speciﬁes additional conditions that are sufﬁcient for FDR control.
These conditions are not strictly fewer than those required for empirical model-X knockoff-based methods to control FDR, but they may be easier to satisfy adequately in practice. For example, the condition that there exists a sequence ( G n ) ∞ n =1 converging to the true mapping G is satisﬁed asymptotically by many ﬂow architectures that are universal distribution approximators, including the Gaussianization Flows and Masked Autoregressive Flows used in our experiments (Huang et al., 2018;
Meng et al., 2020; Kobyzev et al., 2020).
In practice, it is unlikely that an exact mapping G will be learned, as doing so
Nonetheless, normalizing ﬂows work extremely well in practice; Theorem 1 gives intuition for the good performance of F LOW S ELECT that we observe empirically.
Experiments 5.1 Synthetic experiment with a mixture of highly correlated Gaussians
We compare F LOW S ELECT to the aforementioned knockoff methods with synthetic data drawn from a mixture of three highly correlated Gaussian distributions with dimension D = 100 .
For each knockoff method, we use the exact implementation described in their respective papers, and we utilize the code made publicly available by the authors (c.f. Appendix D.3 for further details).
For further comparison, we also implement the MASS knockoff procedure from Gimenez et al.
(2019) and the
RANK knockoff procedure from Fan et al.
These methods estimate the unknown feature distribution using either a mixture of Gaussians (MASS) or a sparse precision matrix (RANK), and then sample the knockoffs directly as in Candès et al.
To generate the data, we draw N = 100 , 000 highly correlated samples.
For i = 1 , . . .
, N , we sample
[Equation]
i i
[Equation]
i
[Equation]
(6) with mixing weights π = (0 . 371 , 0 .
258 , 0 . 371) , mean vector µ = (0 , 20 , 40) , and covariance matrices Σ m .
Each covariance Σ m follows an AR(1) pattern such that (Σ m ) i,j = ρ | i − j | m where ρ = (0 . 982 , 0 . 976 , 0 . 970) .
The response Y i is linear in f i ( X i ) for some function f
i and coefﬁcient vector β i.e., Y i =
f i ( X i ) β + (cid:15) i .
Each coefﬁcient β
j equals 100 √ N B j , where B j = 0 with probability 0 .
8 , B j = 1 with probability 0.1, and B j = − 1 with probability 0.1.
We consider two different schemes for the f i that connect the features to the response.
In our linear setting, f i is equal to the identity function.
In our nonlinear setting, f i ( x ) is set equal to sin(5 x ) for odd i and f i ( x ) = cos(5 x ) for even i .
The experimental setting we have described so far is adapted from Sudarshan et al.
However, we found that the N = 2000 they used was too few observations for any of the methods to do well in a general non-linear setting.
Moreover, in many situations where controlled feature selection is deployed, neighboring features will be highly correlated.
To reﬂect this, we also increased the base correlation between features within each mixture to create a more challenging example.
We show results under the original settings of Sudarshan et al.
(2020) in Appendix K.
For each model, we use 90% of the data for training to generate null features and the remaining 10% for calculating the feature statistics.
To deﬁne the feature statistics, we use the holdout randomization test (HRT) described at the end of Section 2.
For the HRT, we employ different predictive models for each response type (“linear” and “nonlinear”).
Speciﬁcally, for the linear response, we use the predictive log-likelihood from the LASSO (Tibshirani, 1996), and for the nonlinear response, we use the predictive negative mean-squared error from a random forest regressor (Breiman, 2001). scRNA-seq
First, we look at how each procedure models the covariate distribution in Figure 1.
In order to be valid knockoffs, the distribution of two knockoff features needs to be equal to that of the covariates.
In this challenging example, each of the empirical knockoff methods fails to match the ground truth.
In particular, DDLK and DeepKnockoffs are over-dispersed, while KnockoffGAN suffers from mode collapse.
These ﬁndings for DeepKnockoffs and KnockoffGAN are similar to those reported by
Sudarshan et al.
Other than MASS, which directly ﬁts a mixture of Gaussians, F LOW S ELECT is the only method that matches the basic structure of the ground truth.
Figure 2 shows that the empirical knockoff procedures fail to control the FDR for both linear and nonlinear responses.
One explanation for this lack of FDR control is the inability of the deeplearning-based methods to accurately model a knockoff distribution (c.f., Figure 1).
As a result, the assumptions for the knockoff procedure will not hold, and FDR control is not guaranteed.
The effects of misspeciﬁcation are clearly visible in the case of RANK, which approximates the mixture-of-Gaussians data with a multivariate Gaussian.
However, even MASS, when given access to the correct data distribution, does not achieve across-the-board FDR control.
This highlights the potential sensitivity of knockoffs to parameter misﬁt even when the underlying distributional family of the features is known.
This is conﬁrmed by the fact that, when provided with the true parameters, the oracle Model-X maintains FDR control, though with signiﬁcantly less power than F LOW S ELECT .
(c.f. Appendix H).
5.2 Semi-synthetic experiment with scRNA-seq data
In this experiment, we use single-cell RNA sequencing (scRNA-seq) data from 10x Genomics (10x
Genomics, 2017).
Each variable X n,g is the observed gene expression of gene g in cell n .
These data provide an experimental setting that is both realistic and, because gene expressions are often highly correlated, challenging.
More background information about scRNA-seq data can be found in
Agarwal et al.
We normalize the gene expression measurements to lie in [0 , 1] , and we add a small amount of
Gaussian noise so that the data is not zero-inﬂated.
As in the semi-synthetic experiment from
Sudarshan et al.
(2020), we pick the 100 most correlated genes to provide a challenging, yet realistic example.
We simulate responses that are both linear and nonlinear in the features.
Figure 2 shows that F LOW S ELECT maintains FDR control across multiple FDR target levels, feature statistics, and generated responses.
In cases in which the knockoff methods control FDR successfully, F LOW S ELECT has higher power in discovering the features the response depends on.
An advantage of knockoffs over CRT-based methods like F LOW S ELECT is that the predictive model only needs to be evaluated once.
Hence, while F LOW S ELECT has a faster runtime than DDLK for this experiment, it is slower than DeepKnockoff and KnockoffGAN.
However, Figure 2 shows that these two models fail to reliably control FDR and have much less power than F LOW S ELECT ; it is not clear how additional computational resources could be leveraged to improve the performance of these competing methods.
A full table of runtimes on the scRNA-seq dataset can be found in Appendix F.
The need to compute a different predictive model for each feature within the CRT is mitigated by using efﬁcient feature statistics such as the HRT (Tansey et al., 2021) and the distilled CRT (Liu et al., 2020).
These methods ﬁt a larger predictive model once, then evaluate either the residuals or test mean-squared-error for each feature individually.
Moreover, the ability to scale to large feature dimensions D is more limited by ﬁtting the feature distribution than computational burden, a trait shared by both knockoff- and CRT-based methods.
F LOW S ELECT provides asymptotic guarantees of FDR control assuming sufﬁcient MCMC samples have been drawn for the p-values to converge.
In this experiment, the consequence of terminating
MCMC sampling before convergence is low power, rather than loss of FDR control (see Figure 7 in
Appendix J).
Even for small numbers of MCMC samples, the FDR stabilizes below the target rate, while the power steadily increases with the number of samples.
Because the MCMC run is initialized at the true features, we speculate that the sampled features will be highly correlated with the true features in the beginning of the run, making it harder to reject the null hypothesis that a feature is unimportant.
5.3 Ablation Study
F LOW S ELECT differs from the competing knockoff-based approaches in two ways: using normalizing ﬂows with MCMC to model the feature distribution for sampling null features and using the CRT for feature selection.
To illustrate the impact of each of these components separately, we compare to the procedure used in Tansey et al.
(2021), which uses mixture density networks (MDNs) to model the complete conditional distribution of each feature P ( X j |
X − j ) separately.
They then sample null features from these learned distributions directly and use the HRT for feature selection.
Since both
F LOW S ELECT and this procedure utilize the HRT, this allows us to evaluate whether the performance improvement of F LOW S ELECT over empirical knockoffs is solely due to use of the HRT.
We compare the MDN-based approach to F LOW S ELECT on the mixture-of-Gaussians (Section 5.1) and scRNA-seq (Section 5.2) datasets.
A plot of this comparison can be found in Appendix G.
While the MDN-based approach was able to match the performance of F LOW S ELECT on the scRNAseq dataset, it failed to control FDR at any level on the Mixture-of-Gaussians dataset, indicating that MDNs are less ﬂexible than normalizing ﬂows.
In aggregate, these results show that both the normalizing ﬂows paired with MCMC and the use of the HRT for signiﬁcance testing are key to the performance of F LOW S ELECT .
5.4 Real data experiment: soybean GWAS
Genome-wide association studies are a way for scientists to identify genetic variants (single-nucleotide polymorphisms, or SNPs) that are associated with a particular trait (phenotype).
We tested F LOW S ELECT on a dataset from the SoyNAM project (Song et al., 2017), which is used to conduct GWAS for soybeans.
Each feature X j takes on one of four discrete values, indicating whether a particular
SNP is homozygous in the non-reference allele, heterozygous, homozygous in the reference allele, or missing.
A number of traits are included in the SoyNAM data; we considered oil content (percentage in the seed) as the phenotype of interest in our analysis.
There are 5,128 samples and 4,236 SNPs in total.
To estimate the joint density of the genotypes, we used a discrete ﬂow (Tran et al., 2019).
Modeling of genomic data is typically done with a hidden Markov model (Xavier et al., 2016); however, such a model may fail to account for long range dependence between SNPs, which a normalizing ﬂow is better suited to handle.
Having a more ﬂexible model of the genome enables F LOW S ELECT to provide better FDR control for assessing genotype/phenotype relationships.
For the predictive model, we used a feed-forward neural network with three hidden layers.
Additional details of training and architecture are presented in Appendix E.
{+Chromosome - l og 10 ( p )  3    911 12 13 14 15 16 
+} A graphical representation of our results is shown as a Manhattan plot in Figure 3, which plots the negative logarithm of the estimated p-values for each SNP.
At a nominal FDR of 20%, we identiﬁed[-Chromosome - l og 10 ( p ) 
-] seven SNPs that are associated with oil content in soybeans.
We cross-referenced our discoveries with other publications to identify SNPs that have been previously shown to be associated with oil content in soybeans.
For example, F LOW S ELECT identiﬁes one SNP on the 18th chromosome,
Gm18_1685024 , which is also selected in Liu et al.
F LOW S ELECT also selects a SNP on the 5th chromosome, Gm05_37467797 , which is near two SNPs ( Gm05_38473956 and Gm05_38506373 ) identiﬁed in Cao et al.
but which are not in the SoyNAM dataset.
Sonah et al.
(2014) identiﬁes eight SNPs near the start of the 14th chromosome, and we select multiple SNPs in a nearby region on the 14th chromosome (seen in the peak of dots on chromosome 14 in Figure 3).
However, the dataset in Sonah et al.
(2014) is much larger ( ≈ 47 , 000 SNPs), which prevents an exact comparison.
A list of all SNPs selected by our method is provided in Appendix E. For this experiment, F LOW S ELECT tests over 4000 features in 10 hours using a single GPU.
None of the empirical knockoff procedures (Sudarshan et al., 2020; Jordon et al., 2019; Romano et al., 2020) tested more than 387 features.
This shows the potential for F LOW S ELECT for high-dimensional feature selection with FDR control in a reasonable amount of time.
Additional details about this experiment are available in Appendix E.
Discussion
F LOW S ELECT enables scientists and other practitioners to discover features a response depends on while controlling false discovery rate, using an arbitrary predictive model; even large-scale nonlinear machine learning models can be utilized.
By making fewer false discoveries for a ﬁxed sensitivity level, F LOW S ELECT can reduce the cost of follow-up experiments by limiting the number of irrelevant features considered.
In contrast to the original model-X knockoffs method, F LOW S ELECT does not require the feature distribution to be known a priori, nor does it require the feature distribution to have a particular form (e.g., Gaussian).
Neither of these conditions are often satisﬁed in practice.
One limitation shared by both the conditional randomization test (CRT) and knockoffs is low power in cases in which important features are highly correlated with other important features.
To mitigate this limitation, the CRT can be applied to test the signiﬁcance of groups of correlated features rather than individual features.
Within the F LOW S ELECT framework, this entails modifying the MCMC step to draw null samples of groups of features conditioned on the others. The group’s p-value can then be calculated with the same holdout randomization test (HRT) statistic used for testing individual features. Group feature selection has also been explored for knockoffs (Dai & Barber, 2016; Liu et al., 2020). Another limitation of F LOW S ELECT stems from its reliance on normalizing ﬂows.
The ﬂexibility of normalizing ﬂows, though often beneﬁcial, comes at a cost: sufﬁcient training examples are needed to learn the feature distribution, limiting applicability in data-starved regimes.
Fortunately, as we show in Appendix K, F LOW S ELECT fares no worse than competing methods in low-data settings.
In these regimes, F LOW S ELECT could also use other density estimation techniques such as autoregressive models.
Furthermore, learning the feature distribution (potentially from limited data) is not the sole difﬁculty that the deep-learning-based knockoff methods face.
To demonstrate that there are additional sources of difﬁcult for knockoff-based methods, we gave DDLK, which typically ﬁts the data distribution as part of its training procedure, access to the the exact joint density ; neither the empirical FDR nor the power improved signiﬁcantly (c.f. Appendix I).
This result points to a failure of DDLK to enforce the swap property, which is a challenging task as the number of swaps grows exponentially with the number of features.
F LOW S ELECT , on the other hand, achieves FDR control under a different set of conditions that often are simpler to satisfy adequately in practice. Acknowledgments and Disclosure of Funding Derek Hansen acknowledges support from the National Science Foundation Graduate Research
Fellowship Program under grant no. 1256260. Any opinions, ﬁndings, and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reﬂect the views of the National Science Foundation.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] Runtime discussion in section 5. and limitations of normalizing ﬂows in section(c)
Did you discuss any potential negative societal impacts of your work?
[N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] Theoremstatement lists all assumptions.
(b) Did you include complete proofs of all theoretical results?
[Yes] Yes; Appendix B contains the proof to Theorem 1. 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] Code included as supplement, and speciﬁc details were described in the Appendix.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] In Appendices D and E.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] All ﬁgures comparing FDR and Power results (Figure 2 in main text, Figures 4-6, 8 in the Appendix) contain error bars indicating results across runs.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] We specify hardware used in Appendix F. 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] See Appendix C (b) Did you mention the license of the assets?
[Yes] See Appendix C (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A] Data are publicly available (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
Normalizing ﬂows (Papamakarios et al., 2021) represent a general framework for density estimation of a multi-dimensional distribution with arbitrary dependencies. Brieﬂy, suppose X ∼ P X is a random variable in R d . Now, let Z ∼ N (0 , I d ) be a multivariate standard normal distribution.
We assume there exists a mapping G that is triangular, increasing, and differentiable such that G ( X ) = Z. A formal treatment of when such a G exists can be found in Bogachev et al.
However, a sufﬁcient condition is that the density of X is greater than 0 on R d and the cumulative density function of X j , conditional on the previous components X ≤ j , is differentiable with respect to
X j , X ≤ j (Papamakarios et al., 2021):
[Equation]
From this construction, each U i is independent of all previous U i and has distribution Unif[0 , 1] .
From there, we simply set Z i =
Φ − 1 ( U i ) , where Φ is the CDF of the standard normal.
Since G i ( X ) depends only on the elements in X up to i , it is triangular.
Because p X > 0 , the conditional cdfs are strictly increasing, so G is an increasing map.
Finally, since each cdf is differentiable, the entire map G is differentiable, and its Jacobian is non-zero.
Because of the inverse mapping theorem, G is invertible and we can write
X =
G ( Z ) .
Normalizing ﬂows are a collection of distributions that parameterize a family of invertible, differentiable transformations G θ from a ﬁxed base distribution Z to an unknown distribution X .
Using the change-of-variables theorem, we can express the distribution of X in terms of the base distribution density p Z and the transformation G θ :
[Equation]
where ∂G θ ( X ) ∂X is the Jacobian of G .
The goal is to ﬁnd a parameter value ˆ θ that maximizes the likelihood of the observed X :
[Equation]
A key feature of normalizing ﬂows is that they are composable.
A.1 Flow Architecture
In experiments, the ﬁrst layer G is a Gaussianization ﬂow (Meng et al., 2020) applied elementwise:
[Equation]
where Φ − 1 is the standard normal inverse CDF.
With sufﬁciently large M , this Gaussianization layer can approximate any univariate distribution.
This is composed with a Masked Autoregressive
Flow (MAF) F (Papamakarios et al., 2017), which consists of MADE layers interspersed with batch normalization and reverse permutation layers:
MADE j,k =
( X j − µ j,k ) exp( − α j,k ) where µ j = f µ j,k
( X <j ) α j = f α j,k ( X <j )
F = MADE j,K ◦ BatchNorm ◦ Reverse ◦ MADE
j,K − 1 ◦ · · · BatchNorm ◦ Reverse ◦ MADE j,
Here, f µ j and f α j are fully connected neural networks.
B Proof of convergence
Theorem 1.
Let X ∈ R N
×
D be a random feature matrix, where each row X i, · is independent and identically distributed; x ∈ R N × D be the observed feature matrix; and α j be the p-value as deﬁned in Equation (4) with test statistic T j ( X ) .
Suppose there exists a sequence of functions ( G n ) ∞ n =1 and a base
random variable
Z satisfying the following conditions:
Each G n is continuously differentiable and invertible.
G n → G pointwise for some map G that is triangular, increasing, continuously differentiable, and satisﬁes G
( X i, · ) D =
Z .
For n = 1 , 2 , . . .
, let X n be the random feature matrix where each row i is independent and has distribution X ni, · = ( G n ) − 1 ( Z ) .
Then, the p-value in Equation (5) calculated using K MCMC samples targeting X n · ,j | X n · , − j = x · , − j converges to the correct p-value α j with probability 1 .
Proof of Theorem 1.
Without loss of generality, we consider the ﬁrst feature, which is indexed by j = 1 .
Let p
X
be the density of each row of the matrix X i, · and p Z the density of the base variable
Z .
For each i.i.d observation at i = 1 , . . .
{+, N , we deﬁne F to be the cumulative distribution function of X i, 1 conditional on the other features X i, − 
+}
= x i, − 1 :
[Equation]
For a particular mapping G n , we deﬁne F n analogously:
[Equation]
Since G n and G are continuously differentiable,
[Equation]
[Equation]
Then, by the dominated convergence theorem, F n → F pointwise.
Let X ni, 1 ∼ F n .
Since F n → F pointwise, and F is a distribution function, X ni, 1 converges in distribution to X i, 1 |
X i, − 1 = x
i, − 1 .
Likewise, the joint distribution across all independent observations, written X n · , 1 , converges in distribution to X · , 1 | X · , − 1 = x · , − 1 .
Now, let ˜ X n · , 1 be equal in distribution to X n · , 1 , but sampled such that it is independent of the outcome
Y .
It follows from the reasoning above that ˜ X n · , 1 converges to the desired null distribution ˜ X · , 1 | X · , −as n → ∞ .
{+Deﬁne g 1 (˜ x · , 1 ) (cid:44) 
+}
[ T 1 ( X ) < T 1 ([˜ x · , 1 , X · , − 1 ])] .
With the regularity condition that T 1 is discontinuous on a set of measure zero, the expectation converges:
[Equation]
The Cesaro average of g calculated over MCMC samples that target the distribution of ˜ X n · , 1 under the probability law of G n converges almost surely to E ˜ X n · , 1 ( g 1 ) (Smith & Roberts, 1993).
That is,
[Equation]
Combining Equation (10) and Equation (11) gives the desired result.
C Feature datasets
[Table]
Licensing All of the data used is available for personal use.
Terms for the scRNA-seq data can be found here: https://www.10xgenomics.com/terms-of-use .
The scRNA-seq data was accessed using scvi-tools (Gayoso et al., 2021), distributed under the BSD 3-Clause license.
The soybean data is part of the SoyNAM R package (Xavier et al., 2019), distributed under the GPL-3 license.
D Architecture and training details for synthetic experiments D.1 FlowSelect
For F LOW S ELECT , the joint distribution was ﬁtted with a GaussMAF normalizing ﬂow as described in Appendix A.
The ﬁrst Gaussianization layer consisted of M = 6 clusters, followed by 5 layers of
MAF.
Within each MAF layer, the neural network consisted of three masked fully connected residual layers with 100 hidden units, followed by a BatchNorm layer.
We trained the Gaussianization layer ﬁrst with 100 epochs and learning rate 1 × 10 − 3 within the
ADAM optimizer.
This allowed the Gaussianization layer to learn the marginal distribution of each feature.
Then, we jointly trained the whole architecture with 100 epochs and learning rate 1 × 10 −using ADAM.
MCMC We draw 1000 samples using a Metropolis-Hastings procedure.
The proposal distribution is a random walk:
[Equation]
where ˆ σ 2 j is the sample conditional variance:
[Equation]
where ˆΣ j = (cid:99)
Var ( X )
D.2 Variable selection methods
Linear For the linear response, we estimate a linear model with an L1 penalty (aka the LASSO) on training data:
[Equation]
The penalization term λ is selected via 5-fold cross-validation.
Nonlinear For the nonlinear response, we ﬁt a random forest on the training data.
The hyperparameters are the defaults in the scikit-learn implementation.
Feature statistic If ˆ f ( X ) is the ﬁtted regression function, then the feature statistic is the negative mean-squared error:
[Equation]
D.3 Competing methods
For DDLK (Sudarshan et al., 2020), KnockoffGAN (Jordon et al., 2019), and DeepKnockoffs, (Romano et al., 2020), we used the exact architecture and hyperparameter settings from their respective papers.
For the ablation study in Section 5.3, we use the exact implementation in Tansey et al.
For these methods, we used the code that the researchers graciously made publicly available:
Method
Link
DDLK https://github.com/rajesh-lab/ddlk/
DeepKnockoffs https://github.com/msesia/deepknockoffs/
HRT (MDN) https://github.com/tansey/hrt/
KnockoffGAN https://github.com/firmai/tsgan/tree/master/alg/knockoffgan
For MASS (Gimenez et al., 2019), we followed their described procedure and ﬁt a mixture of
Gaussians to the feature distribution using scikit-learn, selecting the number of components via the
Akiake Information Criterion (AIC).
We then used the knockoffs R package, available on CRAN, to sample knockoffs using the estimated parameters for each component.
For RANK (Fan et al., 2020), we estimate the sparse precision matrix using the Graphical LASSO (Friedman et al., 2008) implemented in sci-kit learn, using cross-validation to tune the regularization parameter.
We then use the knockoffs R package to sample the knockoffs with this covariance.
E Architecture and training details for soybean GWAS
Discrete ﬂows For the discrete ﬂows in the soybean example, we use a single layer of MADE which outputs a dimension of size 4 .
µ is then set equal to the argmax of this output.
For training the ﬂows, we use a relaxation of argmax with temperature equal to 0 .
Discrete MCMC Each feature has K = 4 values, so we can enumerate all four possible states for each proposal and sample in proportional to these probabilities via a Gibbs Sampling procedure.
Setting the probabilities leads to an acceptance rate of 1 , and the samples are uncorrelated since the previous sample doesn’t enter into the proposal distribution
Predictive model For the predictive model of each trait conditional on the SNPs, we use a fully connected neural network.
This network has three hidden layers of size 128, 256, and 128.
ReLU activations are used between each fully connected layer.
Dropout is used on both the input layer and after each hidden layer with p = 0 .
The learning rate in ADAM was set to 1 × 10 − 5 , with early stopping implemented using a held-out validation set.
The feature statistic for each sample is the negative mean-squared error (MSE) for each observation.
Runtime To obtain sufﬁcient resolution on roughly 4200 simultaneous tests, we drew 100, samples from our model.
The runtime was 10 hours using a single NVIDIA 2080 Ti.
Selected SNPs Table 1 shows the SNPs selected by F LOW S ELECT that are associated with oil content in soybeans.
[Table]
F
each controlled feature selection method
[Table]
G Comparison to Holdout Randomization Test of Tansey et al. (2021)
[Figure]
I DDLK with true joint distribution
K Mixture-of-Gaussians results for FDR and Power under Sudarshan et al. (
[Figure]
settings
L Learned normalizing ﬂow mapping on mixture-of-Gaussians and scRNA-seq datasets