DR-SNE: Density-Regularized Stochastic Neighbor Embedding

Technical Summary: DR-SNE (Density-Regularized Stochastic Neighbor Embedding)

Problem Statement
Standard dimensionality reduction methods, particularly t-SNE and its variants, are formulated as geometric problems aimed at preserving local neighborhood relationships. While effective at maintaining conditional structure (local similarities), these methods implicitly ignore the distribution of probability mass across the dataset. Consequently, they often introduce severe distortions in data density, expanding sparse regions and compressing dense ones. This leads to embeddings that misrepresent the underlying data distribution, producing misleading cluster sizes and spatial relationships. Existing density-aware approaches, such as DensMAP and DenSNE, attempt to mitigate this by enforcing consistency in local neighborhood scales (e.g., radii or bandwidths). However, the paper argues that these scale-based proxies are only indirectly related to local sampling density and can be sensitive to embedding scale and neighborhood definitions.

Methodology
The authors propose DR-SNE (Density-Regularized Stochastic Neighbor Embedding), a reformulation of dimensionality reduction as the joint alignment of two components:

1. Conditional Structure: Capturing local relationships (standard in SNE/t-SNE).
2. Relative Density Structure: Capturing the variation of probability mass via local density statistics.

Theoretically, this is grounded in the decomposition of the Kullback–Leibler (KL) divergence between joint distributions into conditional and marginal components. Standard SNE optimizes only the conditional term by fixing marginal distributions to be uniform. DR-SNE introduces a complementary objective to align the marginal density.

Algorithmic Implementation:

• Density Estimation: Local densities in both high-dimensional ($X$) and low-dimensional ($Z$) spaces are estimated using $k$-nearest neighbor ($k$-NN) distances.
• Normalization: These density estimates ($\rho$) are normalized to unit mean to ensure scale invariance.
• Objective Function: The total loss $L$ combines the standard sparse KL divergence ($L_{KL}$) with a density regularization term ($L_{dens}$):
  $$L = L_{KL} + \lambda L_{dens}$$
  where $L_{dens}$ minimizes the squared difference between normalized log-density estimates:
  $$L_{dens} = \frac{1}{n} \sum_{i=1}^{n} \left( \log \tilde{\rho}^{high}_i - \log \tilde{\rho}^{low}_i \right)^2$$
• Optimization: The method uses a warm-up phase with early exaggeration followed by optimization via Adam with gradient clipping. The parameter $\lambda$ controls the trade-off between neighborhood fidelity and density preservation.

Key Contributions

1. Theoretical Reformulation: The paper explicitly frames dimensionality reduction as the joint alignment of conditional structure and marginal density, identifying the neglect of the marginal term as the root cause of density distortion in standard methods.
2. Direct Density Alignment: Unlike prior work that relies on scale-based proxies, DR-SNE directly aligns normalized log-density estimates. This provides a scale-invariant mechanism for preserving relative density variations.
3. New Evaluation Metric: The authors introduce Density Correlation (DC), defined as the correlation between normalized log-density estimates in the original and embedding spaces, to quantitatively assess density preservation.
4. Empirical Validation: Extensive experiments demonstrate that DR-SNE achieves superior density preservation compared to t-SNE, UMAP, PaCMAP, DensMAP, and DenSNE, while maintaining competitive neighborhood fidelity.

Results

• Density Preservation: Across diverse datasets (Digits, Fashion-MNIST, MNIST, PBMC, Shuttle, and synthetic spirals), DR-SNE consistently achieves the highest Density Correlation scores. For instance, on the MNIST dataset, DR-SNE achieved a DC of 0.991 compared to 0.987 for DensMAP and 0.118 for standard t-SNE.
• Topology vs. Density Trade-off: Ablation studies reveal a clear Pareto frontier. Increasing $\lambda$ improves density correlation monotonically but gradually degrades trustworthiness (local structure). Notably, small non-zero values of $\lambda$ often improve both density preservation and local structure, acting as a stabilizer.
• Downstream Performance: In anomaly detection tasks (using kNN, LOF, Isolation Forest, and centroid-based detectors), DR-SNE yields gains, particularly in density-driven regimes (e.g., Shuttle, Thyroid) where anomalies correspond to low-density regions. It shows robustness across different detectors compared to methods that excel only with specific scoring functions.
• Computational Cost: DR-SNE retains the $O(n^2)$ time complexity per iteration of standard t-SNE due to dense pairwise interactions, though the authors note that existing acceleration techniques (e.g., Barnes-Hut) could be applied.

Significance and Claims
The paper claims that DR-SNE offers a "simple and principled extension" to the stochastic neighbor framework that addresses a fundamental limitation of widely used visualization tools. By explicitly incorporating density information, DR-SNE provides a more faithful representation of the underlying data distribution, particularly in hybrid datasets containing both discrete clusters and continuous variations (e.g., single-cell RNA-seq data).

The authors emphasize that their work suggests dimensionality reduction should be viewed not merely as a geometric problem, but as the alignment of complementary components of the data distribution. They argue that while geometry-focused methods remain preferable for purely clustered or manifold-dominated settings, density-preserving embeddings are critical for statistical analysis, anomaly detection, and tasks where the relative density of data points carries meaningful information. The method provides a controllable trade-off, allowing practitioners to tune the embedding based on the specific structural regime of their data.