Provable FDR Control for Deep Feature Selection: Deep MLPs and Beyond

AI summary

Overview

Research area: High-dimensional statistics and machine learning theory, at the intersection of false discovery rate (FDR) control and explainable AI (XAI). The paper sits in the stat.ML category (arXiv:2512.04696v2).

Technical level: Advanced. The paper is a theoretical statistics contribution built on asymptotic normality, orthogonal invariance, and random-matrix-style arguments, supported by deep learning experiments.

Scope in one sentence: The paper develops a feature-selection procedure for deep neural networks whose first layer is fully connected and proves that it asymptotically controls the false discovery rate, while allowing arbitrary width, depth, and many modern architectural components from the second layer onward.

What This Paper Is About

Feature selection—identifying which input features genuinely matter for a response—has traditionally been split between high-dimensional statistics (which gives rigorous FDR guarantees under strong modeling assumptions like linear or generalized linear models) and explainable AI (which offers flexible attribution scores such as LIME, SHAP, and saliency maps but no Type-I error guarantees). The author proposes a procedure that uses the gradient of a trained network's output with respect to its inputs ("input sensitivity"), combined with data splitting and a symmetry-based thresholding rule, and proves that this procedure approximately controls the FDR under a flexible multi-index data-generating model, regardless of the network's depth, width, or the number of training iterations.

Key Contributions

1. Flexible scope across architectures. The method applies to multilayer perceptrons of arbitrary width and depth, convolutional and recurrent networks, attention mechanisms, residual connections, and dropout, provided the first layer is dense and fully connected.

2. Agnostic to the training protocol. The guarantees accommodate stochastic gradient descent with data-independent initializations and learning rates, and hold at each training iteration t, so early stopping is permitted. The theory does not require existence or uniqueness of a loss minimizer, convexity of the loss, or convergence of the optimization path.

3. Normal approximation of input sensitivity. The paper shows that the gradient-based feature-importance vector is asymptotically normal in the regime where the feature dimension n diverges with q* = o(n) (latent dimension grows slower than feature dimension), for arbitrary sample size m and arbitrary network parameters. The proof technique relies on recursive inheritance of orthogonal invariance of the input sensitivity.

4. Asymptotic FDR control via a simple pipeline. Aggregating input sensitivity across data splits yields a scale-free importance statistic and a thresholding rule, and the paper proves FDR control at any nominal level under stated conditions. The author notes this is the first work to provide a theoretical FDR guarantee for feature selection in such a general deep learning setting.

Main Findings

• Spherical uniformity under the null (Proposition 1). Under Assumptions 1–4, the normalized projection P_B^⊥ ξ^(t) / ||P_B^⊥ ξ^(t)|| is uniformly distributed on the unit sphere inside the orthogonal complement of the column space of B, for each m, n, q* and each SGD iteration t, conditioning on the learning-rate and mini-batch schedule.

• Marginal asymptotic normality (Theorem 1). For any null feature index j ∈ S^c, √n ξ_j^(t) / ||P_B^⊥ ξ^(t)|| converges in distribution to N(0,1) as n → ∞ with q* = o(n), and the convergence is uniform over j ∈ S^c in the Kolmogorov sense.

• Asymptotic FDR control (Theorem 2). Under Assumptions 1–7 and n_0/n → π_0 ∈ (0,1], Algorithm 1 satisfies FDP ≤ α + o_p(1) and limsup FDR ≤ α for any nominal level α ∈ (0,1).

• Scale-free property of the statistics. Because only symmetry of the null distribution is needed, the procedure does not require estimating the asymptotic variance or convergence rate—information about B is not needed.

• Numerical confirmation of normality. Histograms of the latter half of components of √n ξ^(10)/||P_B^⊥ ξ^(10)|| were compared against the standard normal density and a fitted normal density across a 5-layer MLP, a 1D-CNN, an LSTM, and a Transformer, in regimes n ≪ m with (m,n) = (100000, 1000), n ≍ m with (2000, 1000), and n ≫ m with (10, 1000).

• FDR and power behave as predicted. In experiments with (m,n) = (1600, 400) and α = 0.1 using ψ(u,v) = min(u,v), the FDR was brought under control once training progressed and power reached a reasonable level, consistent with the requirement in Assumption 7. The curves in Figure 3 are averages over 20 independent runs with shaded areas indicating one standard deviation.

• A stated limitation. The theory assumes B-right orthogonal invariance (B-ROI) of the design matrix. This accommodates row-wise dependence, heavy-tailed marginals, and low-rank structures, but rules out certain forms of column dependence, discrete-valued entries, and multi-modal distributions. Broader correlation structures are discussed in Appendix C.

• A second stated limitation. The analysis provides no theoretical guarantee on Type-II error (power); only numerical evaluations are reported. The author notes a procedure that selects nothing attains FDR zero but Type-II error of one.

• A power caveat (Remark 5.1). In a single-index model, if the factor E[g'(Z)] is close to zero, the signal direction is cancelled and power of Algorithm 1 is reduced, though FDR validity is unaffected.

Methodology in Plain English

The procedure has two stages.

Stage I — build an importance statistic. After each SGD update, the method computes an "input sensitivity" vector: for each feature j, it sums the partial derivatives of the trained network's output with respect to that feature across all samples. This measures how much the network's prediction responds to each input. The data are randomly split into two equal halves, the sensitivity is computed separately on each half, and the two are combined into a single score via M_j = sign(ξ_{j1}^(t) ξ_{j2}^(t)) · ψ(|ξ_{j1}^(t)|, |ξ_{j2}^(t)|), where ψ is non-negative, symmetric, positive homogeneous, and monotone in each input (examples given: uv, min(u,v), u+v).

Stage II — threshold. The cutoff is chosen as τ_α = min{u > 0 : FDP^(u) ≤ α}, where FDP^(u) is the count of statistics below -u divided by the count above +u (or 1). This works because under the null the statistic is symmetric around zero, so the negative tail acts as a proxy for the unknown number of false discoveries in the positive tail. No p-values, no null distribution estimation, and no knowledge of B are needed.

Why the theory holds. The design matrix is assumed B-right orthogonally invariant, meaning rotating the feature space in any way that leaves B unchanged leaves the distribution of X unchanged. This invariance is inherited by the first-layer weights, then by the sensitivity vector through the chain rule, and is preserved recursively across SGD updates. Under this invariance the noise part of the sensitivity vector points in a uniformly random direction, which yields a Gaussian limit for each coordinate. Because the asymptotic regime keeps the sample size m, the iteration count t, the depth, and the widths unrestricted, the guarantee does not depend on the network converging.

Experiments. Synthetic data come from a multi-index model with q* = 8 latent directions, a quadratic nonlinearity g(x) = (x-2)^2, and a ReLU-type interaction h(x) = max(x,0). Four architectures were trained to test whether the theory survives real deep architectures, and a comparison was run against Neural Gaussian Mirror (NGM) and DeepLINK at nominal level α = 0.1, with n = 500 and m ∈ {2000, 1000, 500}, under Gaussian, heavy-tailed t(3), and spiked design matrices.

Why This Matters

Impact on research. The paper connects two research threads that have developed largely in isolation: FDR-controlling methodology from high-dimensional statistics and attribution-based feature selection from explainable AI. It extends the orthogonal-invariance theory of marginal asymptotic normality—previously developed for loss minimizers (M-estimators) by Zhao et al. and Sawaya et al.—to individual iterates of optimization algorithms, which matters because deep networks almost never reach a global minimum. The recursive orthogonal-invariance proof technique is presented as potentially of independent interest.

Real-world applications (as framed by the paper):

• Identifying genes associated with Alzheimer's disease—the one concrete example the paper names, under the general heading of scientific discovery.
• Other application areas are not enumerated in the paper; the text frames the general setting as scientific discovery plus improving the interpretability of predictive models.

Industry relevance. The paper itself does not discuss industry use, but the setting it targets—deciding which input variables to report as important from a trained deep model—is common wherever feature selection feeds a downstream decision or a report. The guarantees apply without assuming a linear or generalized linear truth, and the assumptions permit heavy-tailed, dependent, and low-rank data, which is a broader modeling scope than classical FDR methods that assume a known feature distribution or a known response model.

Practical reproducibility. All code and scripts are available at https://github.com/sawaya-ka/deep-feature-selection.

Future Directions

1. Relaxing B-ROI. The paper identifies the design-matrix assumption as the main limitation and discusses extensions to general correlation structures in Appendix C, including robustness to elliptical designs. Achieving FDR control under arbitrary feature correlations—without giving up the flexibility of the multi-index model—remains open.

2. Theoretical power guarantees. The paper gives no theoretical control over Type-II error. Since its own Remark 5.1 shows the signal can be cancelled when E[g'(Z)] is near zero, characterizing when power is achievable, and how it scales with m, n, q*, and t, is a natural next step.

3. Better use of data. Data splitting halves the effective sample size and may reduce power. The paper points to aggregating results across multiple random splits (as in stability selection) and cites Du et al., Ren and Barber, and Takahashi, the last showing theoretically that such stabilization increases power in proportional asymptotics. Designing and analyzing optimal aggregation schemes for this pipeline is an open direction.

4. Generalizing the first layer. The architecture assumption requires the first layer to be dense and fully connected and all input dependence to flow through W_1^T x. Extending the guarantee to architectures with more complex input pathways (for example, direct input-to-deeper-layer connections) is not addressed.

Target Audience

This paper is for researchers and graduate students in high-dimensional statistics, statistical machine learning, and the theory side of explainable AI who care about error-rate guarantees for feature selection. It will also interest methodologists who work with knockoff filters, Gaussian mirror, or data-splitting approaches and want to know how far those ideas extend into deep architectures. Practitioners who need a feature-selection routine with a defensible error guarantee on deep models will find the algorithm simple to implement, but the paper assumes familiarity with FDR, orthogonal invariance, and asymptotic normality arguments, so an advanced statistical background is helpful.

Authors’ abstract

We develop a flexible feature selection framework based on deep neural networks that approximately controls the false discovery rate (FDR), a measure of Type-I error. The method applies to architectures whose first layer is fully connected. From the second layer onward, it accommodates multilayer perceptrons (MLPs) of arbitrary width and depth, convolutional and recurrent networks, attention mechanisms, residual connections, and dropout. The procedure also accommodates stochastic gradient descent with data-independent initializations and learning rates. To the best of our knowledge, this is the first work to provide a theoretical guarantee of FDR control for feature selection within such a general deep learning setting. Our analysis is built upon a multi-index data-generating model and an asymptotic regime in which the feature dimension $n$ diverges faster than the latent dimension $q^{*}$, while the sample size, the number of training iterations, the network depth, and hidden layer widths are left unrestricted. Under this setting, we show that each coordinate of the gradient-based feature-importance vector admits a marginal normal approximation, thereby supporting the validity of asymptotic FDR control. As a theoretical limitation, we assume $\mathbf{B}$-right orthogonal invariance of the design matrix, and we discuss broader generalizations. We also present numerical experiments that underscore the theoretical findings.