IDMap Framework: Feature-Aligned Diffusion Maps

• IDMap is a geometric data analysis method that iteratively refines embeddings by deforming the data manifold with locally adaptive anisotropic kernels.
• It employs local linear regression to estimate Jacobians and tangent spaces, aligning the Riemannian metric with directions where the feature map varies most.
• Empirical results show that IDMap robustly outperforms isotropic diffusion maps, effectively collapsing irrelevant dimensions in noisy, high-dimensional data.

The Iterated Diffusion Map (IDMap) framework is a geometric data analysis method designed to identify, extract, and emphasize features of interest in high-dimensional data lying on manifolds. By iteratively deforming the intrinsic geometry of the data through adaptive, anisotropic kernels guided by the feature map’s local Jacobian, IDMap produces low-dimensional embeddings that reflect target features while removing irrelevant directions. This approach generalizes classical diffusion maps by leveraging local covariance structures that align manifold geometry to the feature map, and provides rigorous tools for tangent space estimation, intrinsic dimension selection, and robust manifold learning, especially in settings involving product manifolds or degenerate mappings.

1. Anisotropic Local Kernel Construction

Let $\mathcal{M} \subset \mathbb{R}^m$ be a $d$-dimensional manifold and $$\mathcal{H}: \mathcal{M} \to \mathcal{N} \subset \mathbb{R}^n$$ a feature map. The central innovation of IDMap is the use of local kernels:

$$K_\epsilon(x, y) = \exp\left(-\tfrac{1}{2} (x-y)^\top C(x)^{-1} (x-y) / \epsilon\right)$$

with data-dependent, anisotropic covariance $C(x) \in \mathbb{R}^{m \times m}$. The local geometry is modulated in the directions along which $\mathcal{H}$ varies most strongly. The covariance is chosen to satisfy:

$$c(x)^{-1} = \mathcal{I}(x) C(x)^{-1} \mathcal{I}(x)^\top = D\mathcal{H}(x)^\top D\mathcal{H}(x)$$

where $\mathcal{I}(x)$ projects onto the tangent space $T_x \mathcal{M}$, and $D\mathcal{H}(x)$ is the Jacobian of $\mathcal{H}$. For computational stability, a convex combination is used:

$$C_{\mathcal{H}^{(0)}}(x) = [(1-\tau)I_m + \tau\, D\mathcal{H}(x)^\top D\mathcal{H}(x)]^{-1},\quad 0<\tau \ll 1$$

These kernels induce, in the continuum limit, a Riemannian metric on $\mathcal{M}$:

$$\tilde{g}(u, v) = g_\mathcal{M}(c(x)^{-1/2}u,\ c(x)^{-1/2}v)$$

such that diffusion distances and resulting embeddings become increasingly sensitive to the feature directions of $\mathcal{H}$.

2. Estimation of Local Jacobian and Tangent Spaces

Given sample data $\{x_i\}$ and observed features $\{y_i = \mathcal{H}(x_i)\}$, IDMap estimates $D\mathcal{H}(x)$ at each $x$ using local linear regression. For each $x$, its $k$ nearest neighbors $x_{I(j)}$ are selected, and weighted centered differences are constructed:

$$dx_j = \exp\left[-\|x_{I(j)}-x\|^2/(4\epsilon)\right](x_{I(j)}-x)$$

$$dy_j = \exp\left[-\|x_{I(j)}-x\|^2/(4\epsilon)\right](y_{I(j)}-y)$$

These are assembled into matrices $X \in \mathbb{R}^{m \times k}$ and $Y \in \mathbb{R}^{n \times k}$. The regression

$$\widehat{D\mathcal{H}}(x) = Y X^\top (X X^\top)^{-1}$$

yields a first-order estimate of the projected Jacobian with $O(\epsilon)$ error under regularity conditions. The SVD of $X$ at each point provides information about the local tangent space and rates of change associated with $\mathcal{H}$.

3. Dimension Selection and Bandwidth Scaling Laws

Analysis of the SVD of $X$ over a neighborhood provides both intrinsic local dimension and optimal bandwidth selection. For $d$ tangent directions, singular values scale as $\sigma \sim \sqrt{\epsilon}$, while the remaining scale as $O(\epsilon)$. This scaling is quantified by exponents:

$$\alpha_\ell = \frac{d \ln \sigma_\ell}{d \ln \epsilon}, \qquad \alpha \approx \frac{1}{2} \text{ in tangent directions.}$$

Aggregated scaling exponents give determinant-based dimension estimates, e.g.,

$$d_2(\epsilon) = 2\sum_{\ell=1}^{\lfloor d_1\rfloor} \alpha_\ell + 2 (d_1-\lfloor d_1\rfloor)\, \alpha_{\lfloor d_1\rfloor+1}$$

Bandwidth $\epsilon$ is selected locally by maximizing $d_1(\epsilon)$ (density-based) or minimizing a consistency criterion:

$$M(\epsilon) = \left| \frac{d_1 - d_2}{(d_1 + d_2)/2} \right| + \left| \frac{d\ln d_1}{d\ln\epsilon} \right| + \left| \frac{d\ln d_2}{d\ln\epsilon} \right|$$

This provides a principled approach to local geometry adaptation and regularization.

4. Iterated Diffusion Map Algorithm and Geometric Flow Interpretation

The core IDMap procedure iteratively reshapes data geometry to emphasize the designated feature:

1. At each iteration $\ell$:
   • For each data point, estimate $D\mathcal{H}^{(\ell)}(x)$.
   • Construct an adaptive kernel using updated $C_{\mathcal{H}^{(\ell)}}(x)$.
   • Build the sparse kernel matrix, normalize to a graph Laplacian.
   • Compute the top $M$ diffusion map coordinates.
   • Rescale embeddings:

$$x_i^{(\ell+1)} = (2\pi)^{d/4}(4s)^{d/4+1/2} \left(e^{s \ln \xi_1}\varphi_1(x_i^{(\ell)}),\ldots, e^{s \ln \xi_M}\varphi_M(x_i^{(\ell)})\right)^\top$$

1. Repeat for $T$ iterations.

The process induces a discrete approximation to a geometric flow:

$$\dot{g} = -g + \tfrac{1}{2}\left(D\mathcal{H}^\top D\mathcal{H}\,g + g\,D\mathcal{H}^\top D\mathcal{H}\right)$$

which collapses directions orthogonal to $\ker D\mathcal{H}$, isolating the submanifold that encodes the feature of interest.

Pseudocode Summary:

Input: {x_i}, {y_i = H(x_i)}, T, τ Initialize x_i^(0) = x_i For ℓ = 0,...,T-1: For each x_i^(ℓ): Find k neighbors, estimate d_i, q_i, DĤ_i Build C_i = [(1−τ) I + τ DĤ_i^T DĤ_i]^{-1} Build kernel K_ij = exp[−(x_j−x_i)^T C_i^{-1}(x_j−x_i)/(2ε)] Normalize to graph Laplacian Compute top M eigenpairs (ξ_r,φ_r) Set x_i^(ℓ+1) = diffusion map embedding using (ξ_r,φ_r) Output: Ψ(x_i) = x_i^(T)

5. Convergence Properties and Theoretical Guarantees

Suppose $\mathcal{M} = \mathcal{N} \times \mathcal{P}$, a product manifold where $\mathcal{H}(y, z) = y$. Locally, $D\mathcal{H} = \text{diag}(I_{d_\mathcal{N}}, 0)$. The geometric flow satisfies:

$$\dot{g}_\mathcal{N} = 0, \quad \dot{g}_\mathcal{P} = -g_\mathcal{P}$$

As $t \to \infty$, the irrelevant factor $\mathcal{P}$ is collapsed, and IDMap produces an embedding isometric to $\mathcal{N}$, the feature manifold. In settings where features are not exact quotients, IDMap empirically yields a lower-dimensional embedding where neighborhoods track the feature geometry. Theoretical results demonstrate convergence to the quotient manifold for product structures and near-isometric recoveries up to linear alignment of coordinates.

6. Empirical Results and Robustness

Empirical evaluation spans several canonical manifolds:

• Annulus ($r, \theta$): Target feature = radius $r$. After $T=4$ iterations (τ=0.65, M=250, k=500), the angular component is suppressed and the embedding recovers $r$. Nearest neighbor tests show convergence of feature neighborhoods.
• Torus ($\mathbb{T}^2$): Target feature = either circle factor. The irrelevant circle is collapsed (τ=0.4 or 0.65), yielding 1D embeddings aligned to the chosen circle, again after 4 iterations.
• Sphere ($S^2$): Target feature = $x$-coordinate or more complex functions. IDMap (τ ≈ 0.6–0.7) contracts the sphere to a 1D embedding aligned with the feature.
• Performance metrics include neighborhood-recovery accuracy (fraction of true feature neighbors captured among $k$ nearest neighbors), embedding distortion, and SVD scree plots.
• Robustness: The method is tolerant to moderate ambient Gaussian noise and high curvature. Satisfactory operation depends on proper tuning of local bandwidth $\epsilon$ and pseudotime $\tau$.

In all cases, IDMap outperforms isotropic diffusion maps, which fail to suppress irrelevant factors, thus demonstrating the claim that local anisotropic adaptation aligned with the feature is essential for targeted manifold learning.

Table: Key Algorithmic Elements of IDMap

Component	Mathematical Expression / Procedure	Purpose
Anisotropic Kernel	$K_\epsilon(x, y)$ as above	Induces feature-aligned metric
Jacobian Estimation	$\widehat{D\mathcal{H}}(x)$ via weighted regression	Adapts local geometry
Iterative Update	$x_i^{(\ell+1)}=$ rescaled diffusion coordinates	Geometric flow toward feature
Dimension Selection	SVD scaling laws, maximizing $d_1(\epsilon)$ or minimizing $M(\epsilon)$	Adapt bandwidth, select $d$

7. Context, Implications, and Related Frameworks

IDMap is situated in the lineage of spectral manifold learning methods, extending classical diffusion maps [Coifman & Lafon] by adaptively deforming geometry using supervised feature information. It unifies concepts from kernelized Laplacian approximation, data-driven Riemannian metrics, and geometric flows, and builds on precedents such as local kernels [LK], vector diffusion maps [Singer & Wu], and neighborhood regression.

Key implications:

• IDMap enables feature-focused embedding and manifold quotienting where only a subset of variables are relevant.
• The local scaling and bandwidth adaptation procedures provide robust practical tools for high-curvature, noisy, or high-codimension datasets.
• A plausible implication is that this methodology can be extended to unsupervised feature selection by iteratively aligning kernels to data-driven discriminants.

Potential limitations include sensitivity to proper kernel bandwidth and the elevation of noise if feature signals are weak compared to irrelevant coordinates. Further paper of convergence rates, global alignment, and scaling in massive data regimes remains an open direction.