Partition First, Embed Later: Laplacian-Based Feature Partitioning for Refined Embedding and Visualization of High-Dimensional Data

Partition First, Embed Later: Laplacian-Based Feature Partitioning for Refined Embedding and Visualization of High-Dimensional Data

Speakers: Erez Peterfreund, Postdoc, Applied Math Department, Yale University; Ofir Lindenbaum; Yuval Kluger; Boris Landa

Conference: ICML 2025

YouTube: https://slideslive.com/39044062

Overview

In this compelling talk from ICML 2025, Erez Peterfreund, a postdoc at Yale University, presented novel research on Laplacian-based feature partitioning for enhancing the embedding and visualization of complex high-dimensional data. Titled "Partition First, Embed Later," the work addresses a critical challenge in unsupervised data analysis: how to generate interpretable embeddings when the observed data is influenced by multiple, distinct latent variables. This is particularly prevalent in fields like genomics and medical imaging, where samples (e.g., cells, images) are characterized by a vast number of features (e.g., genes, pixels), yet these features ultimately reflect a much smaller set of underlying biological or physical processes.

The core problem arises because traditional embedding techniques, such as UMAP or t-SNE, typically attempt to capture all underlying structures within a single, unified embedding. When multiple latent processes are simultaneously at play—for instance, both linear and cyclic dynamics in biological data—these methods often produce complex, convoluted visualizations that mix all latent information, making it exceedingly difficult for practitioners to discern and interpret individual processes. Peterfreund and his co-authors propose a paradigm shift: instead of embedding all features at once, they advocate for partitioning the features into distinct groups first, and then generating separate embeddings for each group. This approach aims to isolate and reveal the simpler, individual latent variables, thereby yielding multiple, more easily understandable visualizations.

This research is highly significant for the ML and scientific communities. By providing a principled method to disentangle mixed latent structures, it promises to unlock deeper insights from high-dimensional datasets that are currently obscured by the limitations of conventional visualization tools. The ability to generate refined, process-specific embeddings can accelerate discovery in areas where understanding complex biological or physical phenomena is paramount, offering a powerful tool for researchers and data scientists grappling with the interpretability of high-dimensional data representations.

Background

▶ Watch: Introduction and challenge of good high-dimensional embedding (0:00)

The analysis of high-dimensional data is a cornerstone of modern scientific discovery, particularly in domains such as genomics, proteomics, and medical imaging. In these fields, datasets commonly consist of numerous samples, each characterized by thousands or even millions of features. For example, in single-cell genomics, a dataset might contain thousands of cells (samples), with each cell's state described by the expression levels of tens of thousands of genes (features). The inherent assumption in much of this analysis is that while the observed feature space is vast, the true underlying complexity is governed by a smaller number of latent variables or processes. These latent variables represent the fundamental biological states, developmental trajectories, or disease mechanisms that drive the observed data variations.

A primary goal of unsupervised data analysis in this context is dimensionality reduction and visualization, which seek to project the high-dimensional data into a lower-dimensional space (typically 2D or 3D) while preserving the intrinsic structure. Techniques like t-SNE (t-distributed Stochastic Neighbor Embedding) and UMAP (Uniform Manifold Approximation and Projection) are widely adopted for this purpose. They operate by constructing an affinity matrix (often denoted as W) that quantifies the similarity or proximity between pairs of samples in the high-dimensional feature space. This matrix effectively encodes the underlying geometric or topological structure of the data. Subsequently, an embedding is generated from this affinity matrix, either through gradient-based optimization or eigenvector-based solutions, aiming to represent these relationships faithfully in a low-dimensional plot.

However, a fundamental challenge arises when the observed data is influenced by multiple distinct latent variables or processes operating simultaneously. Consider, for instance, a population of cells undergoing both a linear differentiation process (e.g., stem cell maturation) and a cyclic cell division process (e.g., the cell cycle). If a standard embedding algorithm processes all genes (features) at once, it attempts to capture both the linear and cyclic structures within a single embedding. The result is often a complex, interwoven manifold where the distinct processes are blended together, making it nearly impossible for a human practitioner to disentangle and interpret each individual latent variable. The speaker illustrates this with a toy example of features reflecting a cyclic process (blue features) and a linear process (orange features); a single embedding of all features yields a "quite complex" visualization that mixes these distinct structures. This lack of interpretability undermines the utility of these powerful visualization tools, hindering scientific insight and downstream analysis. The problem exists because current methods often implicitly assume a single dominant underlying manifold or struggle to decompose the data when multiple orthogonal or partially independent manifolds coexist.

Key Findings

▶ Watch: Limitations of traditional 'embed all features' approach (2:00)

The central discovery and contribution of this work is the "Partition First, Embed Later" paradigm, which fundamentally rethinks how high-dimensional data with multiple latent variables should be analyzed and visualized. Instead of generating a single, potentially convoluted embedding from all features, the proposed method partitions the features themselves into K distinct groups. Each feature group is hypothesized to correspond to a specific latent variable or process. Subsequently, separate embeddings are generated for the samples based on each of these partitioned feature groups. This approach yields multiple, simpler, and more interpretable embeddings, with each one reflecting a distinct underlying latent variable.

The key findings and contributions can be summarized as follows:

1. Feature Partitioning for Clarity: The core idea demonstrates that partitioning features, rather than samples, can effectively disentangle mixed latent structures. This leads to multiple, simplified embeddings, each much easier to understand because it reflects only one latent variable.
2. Optimization Framework: The authors formulate an optimization problem that aims to simultaneously find optimal feature partitions and their corresponding sample affinity matrices. This problem minimizes a sum of "smoothness scores," where each score considers only the features within a specific partition and its associated affinity matrix. This formulation naturally extends standard affinity matrix estimation.
3. Closed-Form Solutions and Algorithm: The research provides closed-form formulas characterizing the optimal affinity matrices and optimal feature partitions for the proposed objective function. These analytical solutions enable the construction of an efficient alternating minimization procedure, which iteratively refines the feature partitions and affinity matrices until convergence, leading to the optimal or near-optimal results.
4. Empirical Validation on Genomic Data: The method's effectiveness is strongly demonstrated on real-world genomic data. When applied to cells undergoing both cell differentiation (a linear process) and cell cycle (a cyclic process), the algorithm successfully partitions genes into two groups. The resulting two embeddings distinctly capture the cell differentiation trajectory and the cell cycle phase, respectively, a separation that was not achievable with a single, all-feature embedding. Crucially, these gene partitions correspond to known biological marker genes, reinforcing the method's biological relevance.
5. Quantitative Performance: An evaluation metric based on the average K-nearest neighbor (KNN) across embeddings, paired with distinct biological processes, showed that the "Partition First, Embed Later" approach best reveals the two biological processes compared to other single or two-embedding algorithms.
6. Theoretical Guarantees: The work provides theoretical guarantees in an asymptotic regime. For a modified objective function with an entropy constraint, it is proven that for K=2 feature subsets that are partially independent, the minimum of the objective function occurs when the partitions match the true underlying feature subsets. This theoretical alignment between the optimization problem and the desired outcome is a significant finding.
7. Generalizability: The approach is not limited to embedding and visualization but can be applied to other downstream tasks requiring disentangled feature representations, highlighting its broad potential utility.

Technical Deep Dive

▶ Watch: Background: Affinity matrices in common embedding techniques (3:00)

The technical foundation of "Partition First, Embed Later" lies in a sophisticated optimization framework that extends the principles of manifold learning and graph theory. Traditional embedding techniques, such as t-SNE or diffusion maps, begin by constructing an affinity matrix W. This matrix quantifies the similarity between any pair of samples based on their high-dimensional feature values. A common form for these affinities is a normalized Gaussian kernel, where higher scores indicate greater similarity. Once W is established, the embedding is generated either via gradient-based methods (e.g., t-SNE) or eigenvector-based solutions (e.g., diffusion maps, Laplacian eigenmaps). The crucial point is that this W is typically derived from all available features, leading to a single, aggregated representation of the data's underlying structure.

The authors reframe the problem of constructing this affinity matrix as an optimization task. Specifically, they consider an objective function that minimizes a smoothness score over Euclidean distances, taking into account all features. This optimization problem aims to find an affinity matrix W~ that minimizes this score, subject to certain constraints. It's a known result that the affinity matrix W~ minimizing such a smoothness score often takes the same normalized Gaussian form as those used in standard embedding techniques.

The innovation begins by extending this optimization problem to accommodate feature partitioning. The core idea is to partition the total set of features into K mutually exclusive and collectively exhaustive groups: Ω₁ , Ω₂, ..., Ω_K. For each feature partition Ω_k, a separate affinity matrix W_k is sought, which captures the underlying structure of the samples only considering the features within Ω_k.

Mathematically, the problem is formulated as minimizing the sum of individual smoothness scores:

$$

\min_{\{\Omega_k\}_{k=1}^K, \{W_k\}_{k=1}^K} \sum_{k=1}^K \text{SmoothnessScore}(W_k, \text{Features in } \Omega_k)

$$

where each SmoothnessScore term quantifies the smoothness of the embedding with respect to the affinity matrix W_k and the subset of features in Ω_k. The key insight is that if two feature partitions use the same affinity matrix, the sum of their smoothness scores is equivalent to the smoothness score calculated without any separation, indicating a natural extension.

The authors derive closed-form formulas for the optimal parameters of this extended problem:

1. Optimal Affinity Matrix (W_k): For a given feature partition Ω_k, the optimal affinity matrix W_k is shown to attain a normalized Gaussian form. Crucially, this Gaussian kernel is computed only using the features belonging to Ω_k. This mechanism ensures that each W_k specifically captures the manifold structure implied by its corresponding feature subset, effectively isolating the latent information.

1. Optimal Feature Partition (Ω_k): For a given set of affinity matrices W_k, a feature S is assigned to partition Ω_k if its corresponding affinity matrix W_k minimizes a specific smoothness score, which is calculated considering only the S-th coordinate across all samples. This criterion ensures that features are grouped together based on how well they align with the manifold structure encoded by a particular affinity matrix.

The existence of these closed-form solutions is critical because it enables the construction of an efficient iterative algorithm: the alternating minimization procedure. This algorithm proceeds in alternating steps:

• Step 1: Optimize Affinity Matrices: Given the current feature partitions {Ω_k}, update each affinity matrix W_k using its closed-form solution, considering only the features in Ω_k.
• Step 2: Optimize Feature Partitions: Given the updated affinity matrices {W_k}, reassign each individual feature S to the partition Ω_k whose W_k minimizes the feature-specific smoothness score, as described by its closed-form solution.

These two steps are repeated until convergence, at which point the algorithm yields the optimal (or a locally optimal) set of feature partitions and their corresponding affinity matrices. Once these K affinity matrices are obtained, any standard embedding technique (e.g., UMAP, t-SNE) can be applied to each W_k to generate K distinct, lower-dimensional embeddings, each reflecting a specific latent variable.

The term "Laplacian-based" in the title implicitly refers to the underlying mathematical framework of these smoothness scores and affinity matrices. In manifold learning, these scores are often related to Laplacian operators (e.g., graph Laplacians or Laplace-Beltrami operators), which are fundamental for spectral clustering and dimensionality reduction methods like Laplacian Eigenmaps or Diffusion Maps. While the transcript doesn't explicitly detail the derivation from a Laplacian, the minimization of a smoothness score over distances is a hallmark of such techniques, where the Laplacian matrix plays a central role in encoding the data's intrinsic geometry.

Experimental Setup & Results

▶ Watch: Extended optimization for feature partitioning and optimal parameters (5:00)

To validate the "Partition First, Embed Later" approach, the authors designed a series of experiments using both synthetic and real-world high-dimensional datasets. The primary focus of the presentation was on a real-world genomic data set, which serves as an excellent illustration of the challenges and benefits of the proposed method.

Dataset and Latent Processes:

The experimental dataset consisted of genomic data where:

• Samples corresponded to individual cells.
• Features corresponded to genes.
• Each entry in the data matrix represented the gene expression level within a specific cell.

Crucially, these cells were known to undergo two distinct biological processes simultaneously:

1. Cell Differentiation: A linear process where cells transition from one state to another, much like the "orange process" in the illustrative toy example.
2. Cell Cycle Phase: A cyclic process reflecting the stages of cell division, analogous to the "blue process" in the toy example.

Baseline Experiment:

The initial step involved applying a standard embedding technique (e.g., UMAP or t-SNE, though UMAP was shown) to the samples using all the genes. The resulting embedding was then colored by both the cell differentiation stage and the cell cycle phase.

• Result: This single embedding primarily captured the cell differentiation stage, appearing as a coherent trajectory. However, it largely failed to reveal the underlying cell cycle phase, which remained obscured or mixed within the differentiation structure. This confirmed the premise that traditional methods struggle to disentangle multiple latent variables.

Proposed Method Application:

The "Partition First, Embed Later" algorithm was then applied to this genomic data. The parameter K was set to 2, reflecting the two known biological processes. The algorithm extracted two distinct gene partitions (feature partitions). Subsequently, two separate embeddings were generated, each based on the genes within one of these partitions.

Key Results:

1. Disentangled Embeddings:

• The embedding based on genes from Partition 1 clearly reflected only the cell differentiation stage, showing a smooth linear progression.
• The embedding based on genes from Partition 2 distinctly revealed the cell cycle phase, exhibiting a characteristic cyclic structure.
• This dramatic separation demonstrated the algorithm's ability to isolate and visualize distinct biological processes that were previously conflated.

1. Biological Validation of Gene Partitions: The genes assigned to each partition were found to correspond to known biological marker genes. For instance, genes like Lef1 (a known marker for cell differentiation) were found in Partition 1, while Mki67 (a well-established marker for cell proliferation and cell cycle progression) was found in Partition 2. This alignment with established biological knowledge strongly validates the meaningfulness and accuracy of the feature partitioning.

Quantitative Evaluation:

To quantitatively assess the effectiveness of the proposed embeddings, the authors constructed a novel evaluation metric: the average K-nearest neighbor (KNN) across embeddings. This metric works by pairing each embedding with its corresponding distinct biological process. For algorithms that produce only a single embedding, this embedding was duplicated for comparison. The metric quantifies how well the local neighborhoods in the embeddings preserve the relationships defined by the ground truth biological processes.

• Comparison: The proposed approach (marked by black markers in the presentation) was compared against other embedding algorithms, including standard t-SNE and other methods designed to generate single or multiple embeddings.
• Headline Result: The "Partition First, Embed Later" approach consistently demonstrated the best performance, indicating that it most effectively revealed and captured the two distinct biological processes compared to all other algorithms tested.

Theoretical Results:

Beyond empirical validation, the talk highlighted theoretical guarantees. In an asymptotic regime, where features are composed of K partially independent subsets, and for a slightly modified objective function incorporating an entropy constraint for tractability, the authors proved significant results:

• They derived the limit of the objective function in terms of the feature partitions.
• Crucially, for the specific case of K = 2, they proved that the minimum of this modified objective function occurs precisely when the learned feature partitions match the true underlying feature subsets. This theoretical finding provides strong foundational support for the algorithm's ability to correctly identify and separate distinct latent structures.

Practical Implications

▶ Watch: Empirical results: Genomic data example (cell differentiation, cell cycle) (6:30)

The "Partition First, Embed Later" framework carries significant practical implications for a wide range of practitioners, infrastructure teams, model builders, and deployers working with high-dimensional data. Its core benefit lies in enhancing the interpretability of machine learning models and data visualizations, a critical aspect often overlooked in the pursuit of predictive accuracy.

For Practitioners and Model Builders:

• Clearer Discovery: Researchers in fields like genomics, proteomics, medical imaging, and materials science can gain unprecedented clarity into complex, multi-faceted biological or physical processes. Instead of struggling to decipher a single, entangled embedding, they can now analyze separate visualizations, each reflecting a distinct underlying mechanism. This can accelerate hypothesis generation and scientific discovery.
• Targeted Analysis: The ability to disentangle features allows for more targeted downstream analysis. For instance, in genomics, identifying specific gene sets associated with cell differentiation versus cell cycle enables researchers to focus on relevant pathways and regulatory networks for each process individually, leading to more precise biological insights.
• Improved Model Building: By understanding which features contribute to which latent variable, practitioners can build more robust and interpretable predictive models. Feature engineering can be guided by these partitions, potentially leading to models that are not only accurate but also explainable in terms of underlying phenomena.
• Diagnostic Utility: In clinical settings, disentangling disease subtypes or progression stages that are driven by different molecular pathways could lead to more accurate diagnoses and personalized treatment strategies.

For Infrastructure Teams and Deployers:

• Scalability: While the talk doesn't delve deeply into computational infrastructure, the alternating minimization procedure is generally amenable to parallelization, especially in the step of computing affinity matrices for different partitions. This suggests potential for efficient deployment on modern GPU/TPU architectures for large datasets.
• Framework Compatibility: The method outputs affinity matrices, which are standard inputs for many existing dimensionality reduction libraries (e.g., scikit-learn, UMAP, t-SNE implementations). This means it can be readily integrated into existing ML pipelines and visualization tools without requiring a complete overhaul of the ecosystem.

Tradeoffs and Limitations:

• Pre-specification of K: A primary limitation, as highlighted in the Q&A, is the need to pre-specify the number of partitions (K). While in biological contexts, domain knowledge might suggest the number of processes, this is not always available. However, the speaker offered a practical solution: observing the decay in the objective function's score. As K increases, the objective score will decrease. A significant drop followed by a flattening of the curve (an "elbow" point) can indicate the optimal K, similar to methods used in K-means clustering.
• Expert Interpretation: Although the method provides clearer embeddings, the ultimate interpretation of what each partition represents still requires expert domain knowledge. The algorithm identifies groups of features and their corresponding latent structures, but assigning biological or physical meaning to these structures remains a human task. This is a common aspect of unsupervised learning and not unique to this method.
• Computational Cost: For extremely high-dimensional data and a large number of samples, the iterative nature of the alternating minimization, particularly the computation of affinity matrices, could be computationally intensive. Further research might focus on approximation strategies or more efficient solvers.
• Partial Independence Assumption: The theoretical guarantees rely on an asymptotic regime where feature subsets are "partially independent." While this is a reasonable assumption for many real-world scenarios, its exact implications for arbitrarily complex or highly correlated latent structures might warrant further investigation.

Despite these considerations, the "Partition First, Embed Later" approach offers a powerful paradigm shift, moving beyond the limitations of single-embedding visualizations to provide a more nuanced, interpretable, and biologically or scientifically meaningful view of high-dimensional data. Its broader applicability to other downstream tasks further solidifies its potential impact.

Key Takeaways

• Standard high-dimensional data embedding techniques (e.g., UMAP, t-SNE) often produce complex, uninterpretable visualizations when multiple distinct latent variables are present simultaneously.
• The "Partition First, Embed Later" approach proposes partitioning features, not samples, into K groups, and then generating separate, simpler embeddings for each feature group.
• This method is underpinned by an optimization problem that finds optimal feature partitions and corresponding affinity matrices, minimizing a sum of smoothness scores based on feature subsets.
• An efficient alternating minimization procedure, driven by closed-form solutions for optimal parameters, iteratively refines feature assignments and affinity matrices.
• Empirical validation on genomic data successfully disentangled cell differentiation (linear) and cell cycle (cyclic) processes, with feature partitions corresponding to known biological marker genes (e.g., Lef1, Mki67).
• Quantitative evaluation demonstrated superior performance in revealing distinct biological processes compared to other embedding algorithms, supported by theoretical guarantees for specific conditions.
• While K must be specified upfront, its selection can be guided by observing the decay in the objective function's score, and expert domain knowledge remains crucial for interpreting the resulting embeddings.

About the Speaker(s)

The primary speaker for this presentation was Erez Peterfreund, who is a postdoc in the Applied Math Department at Yale University. He presented this work as part of his recent research. His co-authors on this project include Ofir Lindenbaum, Yuval Kluger, and Boris Landa, contributing to the collaborative nature of this advanced research in high-dimensional data analysis.