Concrete Autoencoders for Feature Selection

• Concrete autoencoders are autoencoder architectures that embed the Concrete distribution into the encoder to perform end-to-end differentiable feature selection via Gumbel-Softmax sampling.
• They use a continuous relaxation of categorical selection to maintain differentiability during training, which improves reconstruction accuracy and facilitates task-specific predictions.
• Variants like CSAE and IP-CAE enhance stability and diversity of selected features, achieving superior performance on high-dimensional datasets such as gene expression and multi-omics.

A concrete autoencoder (CAE) is an autoencoder architecture that performs explicit feature selection by embedding the Gumbel-Softmax (Concrete) distribution into the encoder, enabling end-to-end differentiable selection of a fixed-size subset of the original input features. Unlike traditional approaches that perform selection by discrete or non-differentiable operations, CAEs and their variants use the Concrete distribution to maintain continuous, differentiable approximations to discrete selections during training. This paradigm extends to various supervised, unsupervised, and task-conditional settings, with demonstrated superiority for feature selection, reconstruction, interpretability, and downstream prediction in high-dimensional biological and general-purpose datasets (Abid et al., 2019, Avelar et al., 2022, Nilsson et al., 2024, Maddison et al., 2016).

1. Mathematical Foundations of the Concrete Layer

The foundation of CAEs is the Concrete distribution (Maddison et al., 2016), which offers a continuous relaxation of categorical variables, enabling the use of the reparameterization trick for gradient-based optimization. For a categorical selection among $d$ input features, each selector neuron $i$ is parameterized by logits $\alpha^{(i)} \in \mathbb{R}_{>0}^d$. During training, for each selector, the relaxed selection vector $m^{(i)} \in \Delta^{d-1}$ is sampled as

$$m_j^{(i)} = \frac{\exp\bigl((\log \alpha_j^{(i)} + g_j)/\tau\bigr)}{\sum_{\ell=1}^d \exp\bigl((\log \alpha_{\ell}^{(i)} + g_{\ell})/\tau\bigr)}, \quad g_j \sim \mathrm{Gumbel}(0,1)$$

where $\tau > 0$ is a temperature parameter annealed during training. As $\tau \to 0$, $m^{(i)}$ approaches a one-hot vector, effecting a near-discrete feature pick. This formulation permits backpropagation through the feature selection operation. At test time, feature selection is performed by $j^* = \arg\max_{j} \alpha^{(i)}_j$, so the $i^{th}$ neuron selects one input feature (Abid et al., 2019, Maddison et al., 2016).

2. CAE and Variants: Architectures

The canonical CAE consists of a selector (“encoder”) layer and a decoder. For input $i$0, the $i$1-dimensional pseudo-latent code is

$i$2

and the collection $i$3 is passed to a decoder $i$4 to reconstruct $i$5.

Supervised or task-conditional variants augment or replace the decoder with application-specific prediction heads (e.g., Cox proportional hazards layer for survival analysis in clinical omics (Avelar et al., 2022)). For instance, the Concrete Supervised Autoencoder (CSAE) combines:

• A differentiable Concrete feature-selector encoder $i$6 producing $i$7
• A reconstruction decoder $i$8 outputting $i$9
• A predictor head $\alpha^{(i)} \in \mathbb{R}_{>0}^d$0 (e.g., mapping $\alpha^{(i)} \in \mathbb{R}_{>0}^d$1 to log-hazard for survival outcomes).

Recent work identifies and remedies instability/fluctuations in classic CAE training, where selector neurons may collapse and redundantly select the same features. Indirectly Parameterized CAEs (IP-CAEs) address this by parameterizing selector logits as $\alpha^{(i)} \in \mathbb{R}_{>0}^d$2 with a learned low-dimensional embedding $\alpha^{(i)} \in \mathbb{R}_{>0}^d$3 and shared mapping $\alpha^{(i)} \in \mathbb{R}_{>0}^d$4, yielding smoother optimization and higher feature-selection diversity (Nilsson et al., 2024).

3. Training Objectives and Optimization

The standard CAE objective for unsupervised learning is squared-error reconstruction:

$\alpha^{(i)} \in \mathbb{R}_{>0}^d$5

where $\alpha^{(i)} \in \mathbb{R}_{>0}^d$6 is the $\alpha^{(i)} \in \mathbb{R}_{>0}^d$7 matrix stacking the $\alpha^{(i)} \in \mathbb{R}_{>0}^d$8 row-wise. Gradients flow through all parameters, including the selection logits via the reparameterization path.

For specialized settings, the loss is extended. The CSAE for survival analysis minimizes:

$\alpha^{(i)} \in \mathbb{R}_{>0}^d$9

with $m^{(i)} \in \Delta^{d-1}$0 the negative Cox partial-likelihood (for supervised survival), and $m^{(i)} \in \Delta^{d-1}$1 the standard weight decay, but no explicit entropy or sparsity regularization beyond the annealed temperature. During training, a temperature annealing schedule $m^{(i)} \in \Delta^{d-1}$2 is used, with $m^{(i)} \in \Delta^{d-1}$3 and $m^{(i)} \in \Delta^{d-1}$4 or $m^{(i)} \in \Delta^{d-1}$5 over $m^{(i)} \in \Delta^{d-1}$6 epochs (Avelar et al., 2022, Abid et al., 2019, Nilsson et al., 2024).

4. Feature-Selection Properties and Stability

CAEs offer explicit feature selection, yielding a discrete set of informative input variables used in both analysis and deployment. The stability of these selections has been quantitatively analyzed. In survival-based settings, repeated training with different initializations shows feature-selection histograms following a power-law: a small kernel of features is selected in nearly all runs, indicating robust association with outcome, while the majority of features are rarely selected (Avelar et al., 2022).

Classic CAEs may suffer from instability, including selection duplication and inconsistent convergence, especially when training with nonlinear decoders or with high-dimensional data. The IP-CAE formulation regularizes selector diversity by forcing logit embeddings onto a low-dimensional manifold, thereby increasing the percentage of unique selected features and reducing training variance. Ablation studies indicate that this mechanism is critical for stabilizing deeper or more nonlinear architectures (Nilsson et al., 2024).

5. Empirical Results and Practical Applications

Concrete autoencoders have been validated on both classical ML benchmarks (MNIST, Fashion-MNIST, COIL-20, ISOLET, activity/sensor datasets) and large-scale biological data (gene expression, multi-omics cancer cohorts). Across various settings, CAEs typically outperform or match state-of-the-art baselines for test-set reconstruction error and downstream prediction using the selected features (Abid et al., 2019, Avelar et al., 2022, Nilsson et al., 2024).

In high-dimensional omics, CAEs enable interpretable biomarker discovery with competitive performance to principal component analysis and variational autoencoders. In survival stratification, CSAE’s explicit feature subset enables clear clinical interpretability: in multi-omic TCGA cohorts, the most stably selected features were known clinical covariates, followed by gene-expression and methylation markers (Avelar et al., 2022).

Recent IP-CAE models yield substantial training speedups (e.g., over 20x on ISOLET) and retain unique selections near 100% across epochs. Quantitatively, CAEs and CSAEs match or improve upon alternatives (PCA, Laplacian score, AEFS, UDFS, MCFS, PFA, LassoNet, STG) on normalized MSE and classification accuracy, often with a smaller set of selected features (Abid et al., 2019, Nilsson et al., 2024).

6. Implementation and Practical Considerations

CAEs are implemented by augmenting standard autoencoders with a concrete selector layer in the encoder. Typical hyperparameters include selection cardinality $m^{(i)} \in \Delta^{d-1}$7, initial/final temperature, and decoder architecture (linear or MLP). Training proceeds by sampling Gumbel noise, performing selection, decoding, and updating parameters using standard optimizers (Adam), with or without auxiliary regularizers.

Inference replaces the relaxed $m^{(i)} \in \Delta^{d-1}$8 selection with hard argmax picks, extracting exactly $m^{(i)} \in \Delta^{d-1}$9 features per input for use in downstream tasks. This hard selection property permits the deployment of compact models or low-cost measurement regimes in resource-limited domains (e.g., gene panels) (Abid et al., 2019).

Empirical results indicate that early temperature warm-up and gradual annealing, combined with minor regularization (e.g., $$m_j^{(i)} = \frac{\exp\bigl((\log \alpha_j^{(i)} + g_j)/\tau\bigr)}{\sum_{\ell=1}^d \exp\bigl((\log \alpha_{\ell}^{(i)} + g_{\ell})/\tau\bigr)}, \quad g_j \sim \mathrm{Gumbel}(0,1)$$0 on weights), suffice for stable training. Models such as IP-CAE offer increased tolerance to deeper decoders and larger selection sizes by virtue of their coupled selector architecture (Nilsson et al., 2024).

7. Historical Development and Extensions

The introduction of the Concrete distribution (Maddison et al., 2016) enabled differentiable approximations to discrete sampling, which were then embedded in encoder architectures for feature selection in the original CAE (Abid et al., 2019). Extensions have included supervised variants (CSAE for survival), domain-adapted autoencoders for multi-omics integration, and indirect parameterization (IP-CAE) for enhanced selector diversity and stability (Avelar et al., 2022, Nilsson et al., 2024).

Comparative studies and downstream analyses have formalized CAE feature stability, kernel/long-tail effects in repeated selection, and the practical advantages over both classic unsupervised selection and embedded deep-learning approaches.

---

References:

• "The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables" (Maddison et al., 2016)
• "Concrete Autoencoders for Differentiable Feature Selection and Reconstruction" (Abid et al., 2019)
• "Multi-Omic Data Integration and Feature Selection for Survival-based Patient Stratification via Supervised Concrete Autoencoders" (Avelar et al., 2022)
• "Indirectly Parameterized Concrete Autoencoders" (Nilsson et al., 2024)