MAFS: Multi-head Attention Feature Selection for High-Dimensional Data via Deep Fusion of Filter Methods

TL;DR

MAFS addresses the instability and limited interpretability of existing high-dimensional feature selection methods by integrating multiple filter priors with a deep multi-head attention framework. It uses a filter-weight initialized external attention to form diverse, stable feature representations, followed by a reordering module that consolidates head outputs into a robust ranking. Across simulations with varying relationships and dimensionality and real datasets from cancer gene expression and ADNI, MAFS achieves superior coverage, stability, and often higher predictive correlations with fewer selected features, while remaining more scalable than graph-based baselines. This hybrid approach provides interpretable feature importance scores and practical utility for precision medicine applications where data are ultra-high-dimensional and heterogeneous.

Abstract

Feature selection is essential for high-dimensional biomedical data, enabling stronger predictive performance, reduced computational cost, and improved interpretability in precision medicine applications. Existing approaches face notable challenges. Filter methods are highly scalable but cannot capture complex relationships or eliminate redundancy. Deep learning-based approaches can model nonlinear patterns but often lack stability, interpretability, and efficiency at scale. Single-head attention improves interpretability but is limited in capturing multi-level dependencies and remains sensitive to initialization, reducing reproducibility. Most existing methods rarely combine statistical interpretability with the representational power of deep learning, particularly in ultra-high-dimensional settings. Here, we introduce MAFS (Multi-head Attention-based Feature Selection), a hybrid framework that integrates statistical priors with deep learning capabilities. MAFS begins with filter-based priors for stable initialization and guide learning. It then uses multi-head attention to examine features from multiple perspectives in parallel, capturing complex nonlinear relationships and interactions. Finally, a reordering module consolidates outputs across attention heads, resolving conflicts and minimizing information loss to generate robust and consistent feature rankings. This design combines statistical guidance with deep modeling capacity, yielding interpretable importance scores while maximizing retention of informative signals. Across simulated and real-world datasets, including cancer gene expression and Alzheimer's disease data, MAFS consistently achieves superior coverage and stability compared with existing filter-based and deep learning-based alternatives, offering a scalable, interpretable, and robust solution for feature selection in high-dimensional biomedical data.

Figures (7)

• Figure 1: Overview of MAFS. (a) Overall architecture of the proposed Multi-head Attention-based Feature Selection (MAFS) framework, showing three main modules: Filter Module, Multi-head Attention Module, and Reordering Module. (b) Multi-head attention module where multiple attention heads process the input in parallel. Each head has its own set of learnable parameters to capture diverse feature patterns and interactions. (c) Reordering module that consolidates outputs from individual attention heads to identify the most informative features.
• Figure 2: Results of relationship-specific feature selection for continuous outcomes. Coverage rates under seven different feature–response relationships at selection ratios 0.5%, 1%, 1.5% and 2% in the high-dimensional setting ($n=2{,}000$, $p=100{,}000$) for continuous outcomes.
• Figure 3: Results of relationship-specific feature selection for binary outcomes. Coverage rates under seven different feature–response relationships at selection ratios 0.5%, 1%, 1.5% and 2% in the high-dimensional setting ($n=2{,}000$, $p=100{,}000$) for binary outcomes.
• Figure 4: Feature selection performance across different dimensionalities. Coverage rates under different input feature dimensions with moderate sample size ($n=2{,}000$). Each subfigure compares five feature selection methods across three feature dimensionalities (25K, 50K, and 100K features). Rows correspond to six data type combinations defined by feature types (continuous, categorical, and combined) and response types (continuous and binary). Columns correspond to three selection criteria: top 2%, top 100, and top 500 features.
• Figure 5: Evaluation on Cancer Gene Expression Datasets. Comparison of AUROC performance for feature selection methods across six cancer datasets. Each row corresponds to one dataset, and each column represents a different classifier: Support Vector Machine (left), K-Nearest Neighbors (middle), and Multi-Layer Perceptron (right).