Random Forest Autoencoders for Guided Representation Learning

Thank you for your highly relevant remarks and overall positive review.

Re W1. The term expert labels refers to annotations provided by experts and assigned to each data point in the training set. The training set used in Section 4 is labeled, meaning that each instance is associated with a specific class. Labels are necessary for supervised learning, as they enable us to train a Random Forest and capture task-specific structure in 2D through our RF-AE framework. I.e., the expert labels are used to train the random forest, from which the RF-GAP proximities are extracted. Previous work has extensively shown that the RF-GAP proximities encode the relationships between points, emphasizing the variables that are relevant for the supervised task (e.g. Rhodes et al. [1]). Thus all of the supervised methods in Section 4 were trained using using expert labels, while unsupervised methods operated without labels. For a detailed description of the data sets, including the number of classes in each, please refer to Appendix F.

Re W2. Thank you for pointing this out. We agree that one of UMAP’s key strengths is its speed and scalability, particularly in unsupervised visualization tasks. However, RF-AE is designed for a different goal: producing class-aware out-of-sample embeddings that preserve supervised geometric structure. This naturally involves additional computation, such as the generation of random forest-based proximity matrices.

Since UMAP does not use label information during training, it addresses a fundamentally different problem. As such, runtime comparisons between UMAP and RF-AE are not meaningful, as the learning objectives and computational pathways are not aligned. A more appropriate comparison should be made with supervised methods such as CE, SSNP, NCA, and supervised kernel extensions of PHATE and UMAP.

To provide context, we include training and inference time estimates (in seconds) across datasets in the attached tables. Based on average runtime over 10 repetitions on benchmark datasets such as Sign MNIST (Table R4.1) and OrganC MNIST (Table R4.2), RF-AE falls within the same order of magnitude (OoM, i.e. power of ten) as P-SUMAP, NCA, and RF-PHATE during training. It is approximately one OoM slower than CE and SSNP, and about two OoM faster than RF-PHATE’s default kernel extension. At inference time, RF-AE is one OoM slower than P-SUMAP but significantly faster than RF-PHATE.

RF-AE remains scalable across large datasets. Memory usage was not a limiting factor in any of our experiments, and the method benefits from prototype selection and batch-wise processing to keep both computation and memory demand manageable.

In the current implementation, the most time-consuming step is the computation of the RF-GAP proximity matrix. We are actively working on accelerating this step using GPU, which we expect will substantially reduce training time.

Table R4.1: Average training and test times on Sign MNIST, in seconds, across 10 repetitions.

Table R4.2: Average training and test times on OrganC MNIST, in seconds, across 10 repetitions.

Re M1. In our work, we define new data as data points that were not part of the original training set. In practice, these points may come from the same distribution as the training data or from a different one, provided they belong to the same domain. Labels are not required for new data, as all trained models in our comparison—including RF-AE—can generate out-of-sample embeddings without supervision. In our evaluation framework (Section 4), we artificially introduced new data through test splits. We used predefined training/test splits when available; otherwise, we created them using a random 80/20 stratified split to account for class imbalance.

Re M2. While our method belongs to the broader class of dimensionality reduction techniques, our primary focus is on projecting data into 2D or 3D to enable interpretable insights. We believe the term data visualization more accurately reflects this goal, as the purpose of the embedding is to facilitate direct understanding and subsequent annotation of the data. This wording is also typically used in similar papers that perform dimensionality reduction for data visualization.

References

[1] Rhodes et al. Geometry- and accuracy-preserving random forest proximities. IEEE Transactions on Pattern Analysis and Machine Intelligence, pages 1–13, 2023.

We would like to thank you for the positive review and insightful comments. Here are our detailed answers to your questions.

Re Q1. Thank you for this valuable observation. It is true that in our current implementation, the full random forest proximity matrix is computed before selecting prototypes. However, the matrix is not stored in memory permanently. It is held temporarily for processing and then released, so memory usage remains controlled. In practice, memory was not a limiting factor in any of our experiments.

We agree that for very large datasets, the pairwise proximity computation can still be time-consuming. To improve scalability, we are actively working on accelerating this step through GPU-based implementation. This optimization is currently in progress, and we are confident that our GPU-based implementation will be completed by the time the camera-ready submission is due. Unfortunately, per NeurIPS guidelines, external links (e.g., to updated code) are not allowed during the rebuttal period.

To provide a clearer view of computational cost, we include training and test time comparisons in Table R3.1 (Sign MNIST) and Table R3.2 (OrganC MNIST), in seconds, across 10 repetitions. RF-AE falls within the same order of magnitude (OoM, i.e. power of ten) as RF-PHATE, P-SUMAP, and NCA during training, and is approximately one OoM slower than CE and SSNP. At inference time, RF-AE is roughly two OoM faster than RF-PHATE’s default kernel extension and about one OoM slower than P-SUMAP.

Table R3.1: Average training and test times on Sign MNIST, in seconds, across 10 repetitions.

Table R3.2: Average training and test times on OrganC MNIST, in seconds, across 10 repetitions.

Re Q2. While RF-AE can incorporate any geometric constraint, RF-PHATE is the core focus of our paper, as it ensures that visualizations are aligned with the underlying supervised data geometry. This choice is motivated by Rhodes & Aumon et al. [1], who empirically demonstrated that RF-PHATE outperformed alternatives.

However, to demonstrate the generality of geometry-regularized RF-AE network, we provide in Table R3.3 ablation results on Sign MNIST and OrganC MNIST using three alternative geometric constraints: UMAP, Supervised UMAP (SUMAP), and RF-UMAP (i.e., UMAP applied to RF-GAP proximities). We set our default hyperparameters $\lambda=0.01,N^*=0.1$. RF-AE (RF-UMAP) achieved the best overall performance, followed by RF-AE with SUMAP and UMAP constraints. While RF-AE (RF-UMAP) was competitive with RF-AE (RF-PHATE) on OrganC MNIST, it largely underperformed on Sign MNIST. These results suggest that RF-AE is especially effective when paired with RF-based kernel methods—particularly RF-PHATE—which already capture the underlying RF geometry. In such cases, the geometric and reconstruction objectives are well-aligned, enabling more effective multi-task learning. Qualitatively, RF-AE (RF-UMAP) still fragments same-class clusters on the Sign MNIST dataset, though less severely than UMAP and SUMAP. This supports the idea that RF-PHATE more effectively captures both local and global supervised structure through diffusion.

In summary, RF-PHATE is a strong default regularizer for RF-AE overall, though alternative RF-based regularizers like RF-UMAP may offer valuable refinements in specific scenarios. We will include these additional quantitative and qualitative experiments in the Appendix.

Table R3.3: SIA and kNN accuracies of RF-AE using three different geometric constraints on the Sign MNIST and OrganC MNIST datasets. Scores are averaged across 10 repetitions (mean ± std).

Re Q3. Rhodes et al. [2] demonstrated that the geometry captured by RF-GAP proximities remains stable across varying numbers of trees and minimum node sizes—outperforming other RF-derived kernels in robustness. Moreover, RF-PHATE itself is largely insensitive to its primary hyperparameter, the diffusion time $t$, as shown by Moon et al. [3]. Thus, we did not consider it necessary to perform ablation studies on these secondary hyperparameters, but we are open to doing so if requested.

References

[1] Rhodes, Aumon, et al. Gaining biological insights through supervised data visualization. bioRxiv, 2024.

[2] Rhodes et al. Geometry- and accuracy-preserving random forest proximities. IEEE Transactions on Pattern Analysis and Machine Intelligence, pages 1–13, 2023.

[3] Moon et al. Visualizing structure and transitions in high-dimensional biological data. Nat. Biotechnol., 37(12):1482–1492, Dec 2019.

We appreciate the reviewer highlighting this important line of work. While the out-of-sample extension problem has indeed been extensively studied—including in classical kernel methods such as Nyström and geometric harmonics—our work addresses a distinct and underexplored setting: supervised out-of-sample visualization, where the goal is to preserve both the local structure of the data and class-aware separability in the embedding space. The classic work of Bengio et al. [1] laid important groundwork for out-of-sample extensions using linear mappings for LLE, Isomap, MDS, and spectral methods. However, these approaches typically rely on unconstrained least-squares minimization and are restricted to linear or kernel-based mappings, making them sensitive to the quality of the training data and often inadequate for capturing complex manifold structure in a way that generalizes effectively to new inputs [2,3].

The manifold extension problem remains an active research area, with recent developments such as parametric $t$-SNE [4], parametric UMAP [5], and several geometry-regularized or geometry-aware autoencoders [6--9]. These methods offer more flexible parametric mappings, but most focus either on unsupervised structure or on optimizing embeddings solely for predictive performance, without preserving label-informed geometry that generalizes to new data points.

RF-AE is designed to address this gap by learning an embedding function that preserves both local structure and class separability, based on Random Forest proximities. Unlike supervised methods such as RF-PHATE, which capture label-aware structure but do not offer a learned mapping, RF-AE provides a parametric, interpretable, and generalizable embedding function. Neural methods such as autoencoders can learn mappings, but often lack a strong supervision signal tied to relational structure between points.

As shown in Table 1 (Section 4), RF-AE consistently ranks highest across 20 datasets on both structure-preserving alignment scores and classification accuracy. These results support our claim that RF-AE improves supervised visualization for out-of-sample data by balancing geometric fidelity and label-based information in a way that prior methods do not fully achieve.

We will revise the introduction to more clearly distinguish our motivation and contributions, emphasizing that RF-AE targets interpretable, supervised embeddings that generalize effectively to new data points.

References

[1] Bengio et al. Out-of-sample extensions for lle, isomap, mds, eigenmaps, and spectral clustering. NeurIPS, 16, 2003.

[2] Rudi et al. Less is more: Nyström computational regularization. In Advances in Neural Information Processing Systems, volume 28. Curran Associates, Inc., 2015.

[3] Quispe et al. Extreme learning machine for out-of-sample extension in laplacian eigenmaps. Pattern Recognition Letters, 74:68–73, April 2016.

[4] van der Maaten et al. Learning a parametric embedding by preserving local structure. In Proceedings of the Twelfth International Conference on Artificial Intelligence and Statistics, volume 5 of Proceedings of Machine Learning Research, pages 384–391, Hilton Clearwater Beach Resort, Clearwater Beach, Florida USA, 16–18 Apr 2009. PMLR.

[5] Sainburg et al. Parametric umap embeddings for representation and semisupervised learning. Neural Computation, 33(11):2881–2907, 10 2021.

[6] A.F. Duque, S. Morin, G. Wolf, and K.R. Moon. Geometry regularized autoencoders. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(6):7381–7394, 2022.

[7] Sun et al. Geometry-aware autoencoders for metric learning and generative modeling on data manifolds. ICML 2024 Workshop on Geometry-grounded Representation Learning and Generative Modeling. 2024.

[8] Lee et al. Neighborhood reconstructing autoencoders. In Advances in Neural Information Processing Systems, volume 34, page 536–546. Curran Associates, Inc., 2021.

[9] Nazari et al. Geometric autoencoders - what you see is what you decode. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 25834–25857. PMLR, 23–29 Jul 2023.

Thank you for the valuable feedback and overall positive review. Below, we provide a point-by-point response to the concerns you raised.

Re Q1. In the revised version, we will include a training and inference time comparison in the Appendix. While RF-AE incurs additional training time due to the computation of the random forest proximity, we note that it significantly improves out-of-sample supervised embeddings (see quantitative results in Table 1), which we believe makes the additional cost worthwhile in many practical scenarios.

All experiments were conducted on a shared computing cluster, where compute nodes are allocated dynamically. As a result, runtimes may vary depending on system load and available resources. To give a rough estimate of RF-AE’s runtime relative to other baselines, we report the training and test times on the Sign MNIST (Table R1.1) and OrganC MNIST (Table R1.2) datasets. RF-AE falls within the same order of magnitude (OoM) as RF-PHATE, P-SUMAP, and NCA during training, and one OoM slower than CE and SSNP. At inference time, RF-AE is approximately 2 OoM faster than RF-PHATE, and one OoM slower than P-SUMAP. Memory usage did not pose a constraint in any of our experiments. RF-AE scales well across datasets without requiring extensive memory resources, and our design choices, including our prototype selection, help keep computations and memory demand manageable.

Furthermore, we are actively optimizing our RF-GAP implementation by migrating proximity matrix computation to the GPU (in progress) to improve efficiency. These updates will be reflected in our GitHub repository once completed.

Table R1.1: Average training and test times on Sign MNIST, in seconds, across 10 repetitions.

Table R1.2: Average training and test times on OrganC MNIST, in seconds, across 10 repetitions.

Re Q2. We analyzed the SIA performance of all models under different classification importance strategies, as detailed in Appendix I. In addition to our ensemble approach, we derived feature importances from a standalone $k$-NN classifier, as well as from an aggregated strategy where importances were computed separately from each model ($k$-NN, SVM, MLP) and then combined. We observed that RF-AE consistently ranked at the top across all metrics, regardless of the feature importance strategy. This suggests that RF-AE is able to preserve meaningful structure in a more general sense—not just from the perspective of a single classifier. As for a sensitivity analysis of SIA under different perturbation strategies, we are not aware of other correlation-aware schemes beyond our own, which is inspired by Kaneko [1]. That said, we are open to exploring additional strategies in future work. Finally, although not explicitly stated in the manuscript, feature importances were averaged over 10 random sampling procedures to ensure robustness. We will clarify these points in the final paper.

Re Q3. Yes: RF-AE supports partially labeled datasets, making it a compelling approach in semi-supervised settings compared to strictly class-conditional methods. As noted in Rhodes & Aumon et al. [2]: “ (...) a trained random forest can define similarities to out-of-sample observations, supporting semi-supervised learning (...) ” This is further elaborated in our manuscript (L155): “Note that this definition (of RF-GAP) naturally extends to OOS observations (...) which can be treated as out-of-bag for all trees.” Thus, on the one hand, assuming $N_L$ labeled points and $N_U$ unlabeled points, for a total training size of $N = N_L + N_U$, we treat unlabeled training samples as “out-of-sample” and compute $N$ proximity vectors of size $N_L$, which are used as input to train the RF-AE network. On the other hand, to generate RF-PHATE embeddings for the full training set, we can rely on the Landmark PHATE algorithm proposed by Moon et al. [3]. We first construct the RF-GAP kernel matrix of size $N_{L} \times N_{L}$ between labeled landmarks and compute their embeddings with PHATE. Then, we project the $N_{U}$ remaining unlabeled points with the linear landmark extension using their RF-GAP proximities to the labeled points (the landmarks), as in L118-L122. This provides all the necessary components to train our regularized RF-AE using information from both labeled and unlabeled data. We will add this discussion to the final version of the paper.

Re Limitations: While technically involved, the core idea of RF-AE—embedding data in 2D—is intuitive. Autoencoders learn low-dimensional representations that capture intrinsic structure. The RF-GAP kernel (Section 3.1) replaces standard Euclidean measures with similarities based on discriminative features, effectively downweighting irrelevant ones. RF-PHATE regularization further denoises the embedding and captures long-range relationships. Combined, these elements let RF-AE organize data based on local and global RF-GAP similarities: nearby points share key features, while distant ones differ meaningfully. We validated this using interpretable metrics (Section 4.1; $k$-NN, SIA metrics) and qualitatively through 2D visualizations showing stronger class-aware structure than competing methods.Thus we do not believe that the results of RF-AE will be substantially more difficult to interpret than other, similar methods.

Still, we agree that supervised 2D visualizations should be interpreted with caution due to potential biases in label assignment. When label assignments reflect social or demographic biases, supervised methods are prone to embedding structural biases, as they aim to discriminate between classes. Another source of label bias arises in highly imbalanced scenarios. In such cases, the Random Forest is biased toward majority classes. As a result, minority classes may appear closer to or farther from other classes than expected, since the features that characterize them may not be adequately captured in the RF-GAP proximities. We note that these concerns are not specific to RF-AE but apply broadly to supervised methods. Additionally, RF-AE helps mitigate these effects by avoiding the exaggerated separations often introduced by purely class-conditional methods. These and the other ethical considerations raised by the reviewer will be discussed in a dedicated section of the manuscript.

[1] Kaneko. Cross-validated permutation feature importance considering correlation between features. Analytical Science Advances, 2022

[2] Rhodes, Aumon, et al. Gaining biological insights through supervised data visualization. bioRxiv, 2024

[3] Moon et al. Visualizing structure and transitions in high-dimensional biological data. Nat. Biotechnol., Dec 2019