This paper introduces DiRe-JAX, a new dimensionality reduction toolkit designed to address the challenges of traditional methods like UMAP and tSNE, such as the loss of global data structure and computational inefficiency. Built on the JAX framework for modern hardware acceleration, DiRe-JAX employs a three-step process: it first captures the data's topology using a k-nearest neighbors (kNN) graph, creates an initial low-dimensional embedding (via random projection, PCA, or spectral methods), and then refines this layout using a force-directed optimization. The toolkit includes extensive quantitative metrics for evaluating embedding quality, focusing on the preservation of both local and global structures using techniques like persistence homology. Benchmarks against state-of-the-art UMAP and tSNE implementations show that DiRe-JAX demonstrates considerable promise, particularly in preserving global data relationships and context in various datasets, making it a scalable, efficient, and interpretable solution for machine learning, bioinformatics, and data science.
DiRe-JAX: A New Tool for Preserving Global Structure in Dimensionality Reduction
1 Statement of Need: Addressing UMAP/tSNE Limitations: DiRe-JAX is introduced as a solution to common problems in dimensionality reduction. While methods like UMAP are fast, they can lose the "big picture" global structure of the data, and methods like tSNE can be computationally expensive and difficult to tune.
2 Main methods: A Three-Step Embedding Process: DiRe-JAX reduces data dimensions through a structured workflow:
Topology Capture: It builds a k-nearest neighbors (kNN) graph of the high-dimensional data using the efficient FAISS library.
Initial Embedding: It creates a first-pass, low-dimensional version of the data using one of three methods: Random Projection, Principal Component Analysis (PCA), or Spectral Laplacian embedding.
Layout Optimization: It iteratively adjusts the positions of points in the low-dimensional space using a force-directed layout, where forces of attraction and repulsion are calculated to better reflect the neighborhood structures from the original data.
3 Quantitative measures: Advanced Quality Metrics: A key feature of DiRe-JAX is its suite of tools for quantitatively measuring the quality of an embedding. This moves beyond simple visual inspection.
Global Structure: Assessed using persistence homology (via persistence diagrams and Betti curves) to measure how well the large-scale topological features of the data are preserved.
Local Structure: Measured using embedding stress and a neighborhood preservation ratio to see how well local distances and neighbor relationships are maintained.
Context Preservation: Evaluates whether labeled clusters in the data remain distinct and separable after reduction, using SVM and kNN classifier accuracy as a proxy.
5.1-5.3 Benchmarks: Superior Global Preservation: In benchmark tests against GPU-accelerated versions of UMAP and tSNE, DiRe-JAX showed superior performance in preserving global structure on several datasets.
On a uniform disk dataset, DiRe-JAX was the only method that produced a uniformly distributed embedding, whereas UMAP and tSNE introduced artificial clumps and tears.
On the "Levine 32" bioinformatics dataset, DiRe-JAX was the only method to clearly separate two distinct cell populations from different patients, a critical global feature that other methods obscured.
5.4 Benchmarks: Better Context Preservation: When embedding a "half-moons" dataset, the original UMAP made the non-linearly separable clusters become separable by a straight line, losing the original context. DiRe-JAX's embedding correctly preserved their intertwined, non-linear relationship.
5.2 Benchmarks: Acknowledging Trade-offs: The paper shows that on the MNIST Digits dataset, tSNE performed better at preserving the persistence homology (a measure of global topology), indicating that the best choice of algorithm can be dataset-dependent.
6 Future work: Path to Improvement: The authors plan to add more quantitative measures, scale the algorithm beyond single-GPU memory limits, and develop a pure TPU version for even greater performance.