EnergyMap: Unraveling the Data Manifold with Energy-based Dimensionality Reduction

Learning meaningful low-dimensional representations of high-dimensional data remains a central challenge in machine learning. Despite their popularity, dimensionality reduction methods like UMAP and t-SNE rely on heuristic assumptions about the local structure of data that scale poorly to complex modalities like images. In this work, we introduce EnergyMap, a dimensionality reduction method that leverages the implicit structure of data captured by an Energy-based model (EBM). Instead of relying on local Euclidean distance heuristics, EnergyMap computes geometrically faithful distances between points by finding energy-minimizing geodesics that traverse dense regions of the data distribution, and produces a low-dimensional embedding that preserves these geodesic distances. In order to efficiently optimize geodesics we develop an efficient algorithm for computing geodesics that combines discrete path finding with continuous gradient-based refinement, substantially improving over naive gradient-based optimization. We provide a promising proof of concept demonstrating the potential for EnergyMap to compute embeddings that capture semantically meaningful concepts in high dimensional data.