This repository provides the PyTorch implementation of RF-AE (Random Forest Autoencoders), a hybrid supervised dimensionality reduction framework that combines Random-Forest–guided manifold learning with neural autoencoders. RF-AE integrates the RF-PHATE embedder with a neural autoencoder to learn low-dimensional representations that preserve label-aware geometric structure from the input data.

The accompanying paper, accepted at NeurIPS 2025, is available at https://openreview.net/forum?id=NjxW4m6KdH. To fully reproduce the experiments and figures from the paper, including trained models and datasets, please refer to the following source code and data repository.

Requirements: Python >= 3.9

pip install git+https://github.com/JakeSRhodesLab/RF-AE.git

First, RF-AE trains a Random Forest to derive a supervised proximity matrix using RF-GAP, a Random-Forest-based similarity measure. Each sample is thereby represented as a node vector within a sparse, label-informed similarity graph defined by RF-GAP proximities. Next, an autoencoder is trained to reconstruct these RF-GAP node vectors, effectively encoding the underlying forest-based geometry. In parallel, RF-PHATE embeddings (obtained from a diffusion process applied to the same RF-GAP graph) are used to guide the training toward geometrically consistent low-dimensional representations.

Thus, RF-AE is optimized with a dual-objective loss function:

1. Reconstruction Loss: The autoencoder must be able to reconstruct the input RF-GAP proximity vectors accurately.
2. Geometric Loss: The autoencoder's "bottleneck" (the low-dimensional embedding) must be as close as possible (via MSE) to the target embedding generated by RFPHATE.

The final transform method uses this trained autoencoder to map new data into the learned low-dimensional space.

• Modality-agnostic: handles continuous, categorical, and mixed data types seamlessly.
• Task-flexible: supports both classification and regression settings.
• Robust to missing data: naturally accommodates missing feature values without requiring imputation.
• Label-free generalization: performs out-of-sample projections without needing labeled instances.

Here is a basic example of how to use the RFAE class.

import numpy as np
import torch
from rfae import RFAE

# Create dummy data
n_samples = 1000
n_features = 50
X_train = np.random.rand(n_samples, n_features)
y_train = np.random.randint(0, 3, n_samples)
X_test = np.random.rand(100, n_features)

# Automatically detect the best device (GPU or CPU)
if torch.cuda.is_available():
    device = 'cuda'
elif torch.backends.mps.is_available(): # For Apple Silicon
    device = 'mps'
else:
    device = 'cpu'

# Initialize the RFAE model
rfae = RFAE(
    n_components=2,          # Desired final embedding dimension
    epochs=100,              # Training epochs for the AE
    lam=0.01,                # Weight for reconstruction loss
    device=device,           # Use the detected device
    random_state=42
)

# Fit the model to the training data
# This runs RFPHATE and trains the autoencoder
rfae.fit(X_train, y_train)
# You can also project the training set after fitting using fit_transform(). Do not use transform(X_train) after fit(X_train, y_train), since transform() assumes out-of-sample instances
# X_train_embedding = rfae.fit_transform(X_train, y_train)

# Transform new data into the 2D embedding
# X_test -> Proximities -> AE Encoder -> 2D Embedding
X_test_embedding = rfae.transform(X_test)

print(f"Original test data shape: {X_test.shape}")
print(f"New embedding shape: {X_test_embedding.shape}")

The RFAE class is the main entry point for the model.

These are the arguments you can pass when creating an RFAE instance:

• n_components (int, default: 2): The target dimensionality for the final embedding (i.e., the size of the autoencoder's bottleneck).
• batch_size (int, default: 512): The batch size for training the autoencoder. RF-AE likes large batch sizes. Consider increasing this hyperparameter for faster training.
• lr (float, default: 1e-3): The learning rate for the AdamW optimizer used to train the autoencoder.
• weight_decay (float, default: 1e-5): The weight decay (L2 regularization) for the AdamW optimizer.
• random_state (int, default: None): A seed to ensure reproducible results.
• device (str, default: auto-detected cuda, mps or cpu): The device to run the neural network training on.
• epochs (int, default: 200): The number of epochs to train the autoencoder.
• hidden_dims (list[int], default: None): A list of integers defining the dimensions of the encoder hidden layers. The decoder is built as the mirror reverse of this. If None, hidden dimensions are dynamically set according to the RF-AE network input size: [0.4*input_shape, 0.2*input_shape, 0.05*input_shape]
• embedder_params (dict, default: None): A dictionary of parameters passed directly to the current rfphate.RFPHATE model. RF-PHATE embedding options such as n_components, t, n_landmark, kernel_symm, and verbose belong inside phate_params.
• lam (float, default: 1e-2): The weighting factor for the combined loss function: balanced_loss = lam * loss_recon + (1 - lam) * loss_emb.
  • lam=1.0: Only trains on reconstruction loss.
  • lam=0.0: Only trains on geometric (embedding) loss.
  • If transition-like structure is expected, consider decreasing it to lam=1e-3 or lam=1e-4 to better align with RF-PHATE. To emphasize on class internal structure and better separability, consider increasing to lam=1e-1.
  • Overall, values between lam=1e-2 and lam=1e-3 offer a good balance.
  • Avoid extreme values such as lam=0.0 and lam=1.0, which distort the embedding and over-compress it, respectively.
• dropout_prob (float, default: 0.0): The dropout probability to use in the autoencoder's hidden layers.
• recon_loss_type (str, default: 'jsd'): The loss function for the reconstruction task. Options are:
  • 'kl': Kullback-Leibler Divergence (stable, emphasizes on local proximity reconstruction).
  • 'jsd': Jensen-Shannon Divergence (better than KL for balancing local and global reconstruction due to the introduction of repulsive forces, but slightly less stable than KL).
  • 'potential': Euclidean distance between PHATE potential vectors, i.e. negative-log probability vectors.

When embedder_params=None, RF-AE initializes RFPHATE with defaults tuned for the current forestgeom-backed RF-PHATE API:

{
    "random_state": random_state,
    "n_jobs": -1,
    "proximity_params": {
        "weight_scheme": "gap",
    },
    "phate_params": {
        "n_components": n_components,
        "n_landmark": 2000,
        "kernel_symm": None,
        "verbose": 0,
    },
}

Common overrides:

rfae = RFAE(
    embedder_params={
        "model_type": "et",  # "rf" for Random Forest, "et" for Extra Trees, or "gbt" for Gradient Boosted Trees
        "n_jobs": -1,
        "forest_params": {
            "n_estimators": 500,
        },
        "proximity_params": {
            "weight_scheme": "gap",
        },
        "phate_params": {
            "n_components": 2,
            "t": "auto",
            "n_landmark": 1000,
        },
    }
)

adjust_diagonal and force_symmetric are fit-time RF-PHATE options. Pass them to fit() or fit_transform(), not to RFAE.__init__().

Fits the entire model. This is a two-stage process:

1. Fits the RFPHATE embedder on x and y, then reads the fitted graph from self.embedder.phate_op_.graph.
2. Reads the RF-GAP transition operator used as input to the RF-AE network.
3. Trains the internal PyTorch autoencoder to optimize the balanced reconstruction and geometric embedding loss.

• x (np.ndarray): The training data, shape (n_samples, n_features).
• y (np.ndarray): The training labels, shape (n_samples,).
• adjust_diagonal (bool, default: True): Passed to RFPHATE.fit().
• force_symmetric (bool, default: True): Passed to RFPHATE.fit().

Generates the low-dimensional embedding for new data.

1. Calculates the proximity of x to the training landmarks (from fit) with self.embedder.extend_to_data(x).
2. Passes these proximities through the trained encoder.

• x (np.ndarray): The new data to transform, shape (n_new_samples, n_features).
• Returns (np.ndarray): The low-dimensional embedding, shape (n_new_samples, n_components).

Important note: Calling transform() on the training set after fit() does not return the true training embeddings. This is because transform() first computes proximities using the extended out-of-bag RFGAP definition, which treats the input x as out-of-sample. To obtain the correct training embeddings, use fit_transform() directly.

A convenience method that calls fit(x, y) and returns the training embeddings.

• x (np.ndarray): The training data, shape (n_samples, n_features).
• y (np.ndarray): The training labels, shape (n_samples,).
• Returns (np.ndarray): The low-dimensional embedding for the training data, shape (n_samples, n_components).

Takes a low-dimensional embedding and passes it through the decoder to reconstruct the proximity matrix.

• x (np.ndarray): A low-dimensional embedding, shape (n_samples, n_components).
• Returns (np.ndarray): The reconstructed proximity matrix, shape (n_samples, n_landmark).

A convenience method that chains transform and inverse_transform. It maps raw data x to its low-D embedding and then decodes that embedding back into its reconstructed proximity matrix.

• x (np.ndarray): The raw data to reconstruct, shape (n_new_samples, n_features).
• Returns (np.ndarray): The reconstructed proximity matrix, shape (n_new_samples, n_landmark).