Domain-Aware Feature Selection

Updated 10 February 2026

Domain-aware feature selection is a family of methods that integrate explicit domain information to choose features robust to shifts and heterogeneous data distributions.
These approaches employ techniques such as optimal transport, selective inference, and memory-augmented architectures to align feature distributions and validate selection.
Applications span transfer learning, multi-task recommendation, and biomarker discovery, enhancing interpretability and performance in high-dimensional settings.

Domain-aware feature selection refers to the family of methods that incorporate explicit information about the distributional, structural, or semantic properties of domains (or tasks) into the selection of features for statistical or machine learning models. These algorithms are developed to address the limitations of standard feature selection, which often operate in a domain-agnostic manner and may fail to identify features that are either robust to domain shifts or specifically informative for distinguishing across (or within) heterogeneous domains. Research in this area covers supervised and unsupervised domain adaptation, multi-domain representation, recommendation, multi-task learning, and statistical inference for transfer settings, integrating techniques such as optimal transport, selective inference, memory-augmented neural architectures, and retrieval-augmented LLM guidance. Domain-aware feature selection is foundational for settings with domain shift, multi-source data, high-dimensional predictors, and stringent requirements for interpretability or reliability.

1. Conceptual Foundations and Problem Formulations

Domain-aware feature selection generalizes classical feature selection by leveraging information beyond the target data distribution, such as knowledge from related domains, labeled or unlabeled. This is essential in situations where there are multiple domains (e.g., datasets collected under different conditions, populations, or modalities), or where the statistical properties of features vary across data sources.

Formally, in a canonical regression/adaptation setting, two data blocks are observed:

Source domain: $\{(X_i^s, Y_i^s)\}_{i=1}^{n_s}$ with $X^s \in \mathbb{R}^{n_s \times p}$ , $Y^s \sim N(\mu^s, \Sigma^s)$
Target domain: $\{(X_j^t, Y_j^t)\}_{j=1}^{n_t}$ with $X^t \in \mathbb{R}^{n_t \times p}$ , $Y^t \sim N(\mu^t, \Sigma^t)$

Typically, $n_s \gg n_t$ and $p \gg n_t$ , encapsulating the high-dimensional, small-sample regime and the need for transfer-learning (Loc et al., 17 Jan 2025, Loi et al., 2024).

In the multi-domain multi-task context, the objective is to extract sparse masks or weighting vectors $W \in \mathbb{R}^d$ such that the same (or overlapping) feature subset enables discrimination across domains (e.g., tissues, experimental conditions), while possibly preserving domain-unique structure through auxiliary objectives (Salta et al., 2024).

2. Methodologies for Domain-Aware Feature Selection

A diverse set of methodology classes has emerged, each integrating domain information through different mechanisms:

2.1 Optimal Transport–Based Selection

Optimal transport (OT) is leveraged to realign or compare feature marginal distributions between domains. OT-based domain-aware selection ranks features by their distributional similarity across domains:

Diagonal entries of the coupling matrix $\gamma^{*F}$ between feature marginals indicate features where the source and target distributions are similar; large $\gamma^{*F}_{ii}$ implies robustness to domain shift, while small values highlight domain-variant features (Gautheron et al., 2018).
For adaptation, transported source samples are merged with target data, and downstream selection (e.g., via Lasso or sequential FS) is conducted on this aligned space (Loc et al., 17 Jan 2025, Loi et al., 2024).

2.2 Selective Inference and Reliability Control

Selective inference (SI) frameworks extend classical FS to provide valid, conditional statistical guarantees post-selection, accounting for both the feature selection process and any prior domain adaptation:

SI methods condition on the data-dependent selection event (e.g., the combined effect of OT alignment and selection steps), and compute valid $p$ -values by characterizing solutions along one-dimensional projections subject to unions of quadratic and linear constraints.
Full power is obtained by characterizing the entire feasible region in projected outcome space ("divide and conquer"), in contrast to over-conditioning on a single outcome, which reduces statistical power (Loc et al., 17 Jan 2025, Loi et al., 2024).

2.3 Multi-Domain Representations and Task-Specific Masking

For few-shot and heterogeneous data, feature banks collecting per-domain representations enable per-task reweighting or masking:

A multi-domain feature bank $\mathcal{F} = \{ f_1, \dots, f_K \}$ is constructed by training separate or modular extractors on different domains. At test time, a mask $\lambda \in [0,1]^K$ is optimized against a small support set to select features most relevant for the current task (Dvornik et al., 2020).
Continuous masks or scoring functions induce sparsity and adaptation, often optimized via gradient-based procedure on a loss incorporating nearest-centroid classifiers.

2.4 Domain-Sensitive Attribution and Memory-Augmented Architectures

Features whose empirical distribution and predictive impact differ sharply across domains—termed "domain-sensitive features"—are identified through joint statistical and attribution analysis:

Integrated gradients (IG) or related attribution techniques quantify per-example and per-feature influence on prediction outcomes across domains.
Jensen–Shannon divergence (for categorical) or Wasserstein distance (for numerical features) of effect-weighted empirical distributions yield domain-sensitivity scores (Zhao et al., 2024).
Selected features are stored in memory modules, which, via interposed linear attention-based retrievals, re-inject domain-discriminative information at multiple network levels. This enhances both performance and domain distinction awareness while maintaining online efficiency.

2.5 Retrieval-Augmented LLM Guidance and Weighted Regularization

Recent approaches supplement data-driven FS with domain knowledge mined via LLMs:

LLMs, supplied with feature lists, prompts, and optional retrieved scientific context, emit per-feature relevance scores.
Scores are converted into penalty weights $w_j$ for Lasso regularization via a tunable mapping (e.g., $w_j = (s_j)^{-η}$ ), interpolating between pure data-driven and pure domain-driven selection. Internal cross-validation tunes trust in the domain signal (Zhang et al., 15 Feb 2025).
Retrieval-augmented generation (RAG) architectures further ground LLM assessments in structured or unstructured knowledge bases, e.g., OMIM, for improved interpretability and robustness.

2.6 Multi-Domain Multi-Task Sparsification

Deep learning approaches use joint sparsification layers, shared classifiers, and domain-specific encoders/decoders:

A shared mask $W$ is applied to all domains, with domain-unique VAEs reconstructing residual patterns and a common task classifier ensuring cross-domain predictive alignment.
The objective is a weighted combination of reconstruction error, KL-divergence, classification loss, and $\ell_1$ penalty on $W$ , balancing domain consensus and uniqueness (Salta et al., 2024).

3. Statistical Properties and Theoretical Guarantees

Contemporary domain-aware feature selection methods emphasize statistical validity, computational tractability, and interpretability:

Selective inference frameworks for post-domain adaptation or after optimal transport rigorously control false positive rates (FPR) at a user-specified level $\alpha$ (e.g., $0.05$), even under complex selection and adaptation processes (Loc et al., 17 Jan 2025, Loi et al., 2024).
Power enhancement via full-region conditioning restores maximal conditional sample spaces, maximizing the discriminative power of the selection test.
OT-based methods provide an explicit link to generalization bounds in domain adaptation; minimizing Wasserstein domain discrepancy reliably reduces joint risk upper bounds (Gautheron et al., 2018).
Robustness is achieved in LLM-guided approaches by calibrating the influence of the LLM via cross-validation, so that model performance degrades gracefully if the domain-knowledge signal is weak, adversarial, or hallucinated (Zhang et al., 15 Feb 2025).

4. Empirical Studies and Domain-Specific Insights

Synthesized results across domain-aware FS methodologies reveal consistent empirical advantages:

OT-based FS, followed by standard adaptation/classification pipelines, achieves equivalent or increased accuracy compared to full-dimensional approaches, with dramatically reduced computational cost. For example, using OT-selected features in DA tasks led to a $270\times$ speedup and increased accuracy in visual object recognition benchmarks (Gautheron et al., 2018).
SI-SeqFS-DA and SFS-DA methods maintain FPR $\leq \alpha$ across high-dimensional and real-world datasets, outperforming both naive and Bonferroni-corrected approaches, and yielding highest true positive rates (Loc et al., 17 Jan 2025, Loi et al., 2024).
Multi-domain memory-augmented recommendation achieves higher AUC than prior hard- and soft-sharing architectures, with tailored ablations underscoring the necessity of domain-sensitive attribution and cross-attention memory retrieval (Zhao et al., 2024).
Multi-task, multi-domain genomics models reliably uncover cross-tissue biomarker sets that are stable, interpretable, and outperform single-domain approaches. Notably, 20% of features selected in the across-domain setting would not have been identified in any single-tissue analysis (Salta et al., 2024).
LLM-Lasso, integrating domain reasoning, consistently outperforms pure-filter, unweighted Lasso, and wrapper-based FS, with RAG mode further enhancing prioritization of known drivers in biomedical genomes (Zhang et al., 15 Feb 2025).
GRU/LSTM-based continuous embedding models for feature selection generalize across domains by mapping discrete masks to low-dimensional, smooth spaces and perform efficient refinement search via gradient ascent in embedding space, delivering state-of-the-art compactness and accuracy (Xiao et al., 2023).

5. Implementation Guidelines and Practical Considerations

The design and deployment of domain-aware feature selection systems must attend to several crucial factors:

For OT-based strategies, standardize data, select feasible sample/feature subset sizes, and utilize entropic OT (e.g., Sinkhorn) solvers for efficiency. Diagonal coupling entries provide the ranking scores for selection (Gautheron et al., 2018).
Implement selective inference via divide-and-conquer or parametric programming to discover the truncation region for conditional p-values. Leverage existing selective inference libraries, OT solvers, and LP routines (Loc et al., 17 Jan 2025, Loi et al., 2024).
For LLM-driven penalty-weighting, batch features judiciously to respect LLM input limits; calibrate penalty exponent and lambda by cross-validation (Zhang et al., 15 Feb 2025).
End-to-end deep models for multi-domain FS require coordination of sparsity-layer and multiple subnetworks (VAEs/classifiers), with careful hyperparameter tuning for balancing loss components (Salta et al., 2024).
Memory architectures in recommendation should limit the number of domain-sensitive features injected to control computational load, with linear attention variants preferred for online inference (Zhao et al., 2024).
When utilizing embedding-based sequence models for discrete-to-continuous FS, ensure that the joint training corpus includes diverse strategies and domains, and regularize the reconstructor to maintain subset diversity (Xiao et al., 2023).

6. Applications and Expanded Impact

Domain-aware feature selection now underpins a wide array of applied problems:

Transfer learning and adaptation across imaging modalities, omics platforms, or sensor types, where domain-related artifacts or shifts are prevalent.
Multi-domain recommender systems, where distinguishing domain-sensitive from globally shared features delivers improved personalization and resource efficiency.
Biomedical studies seeking robust biomarker identification acros multi-tissue, multi-cohort settings, and in the presence of severe sample imbalance (Salta et al., 2024, Massi et al., 2021).
Few-shot and meta-learning scenarios requiring per-task feature subset selection from heterogeneous pre-trained banks (Dvornik et al., 2020).
Enhanced interpretability and statistical reliability in high-stakes domains (clinical, finance), where domain shifts and selection bias can incur costly errors.

These advances collectively increase interpretability, trustworthiness, and generalization in machine learning systems operating at the interface of diverse data sources and evolving domains.