Differentiable Information Imbalance (DII)

• DII is a differentiable framework for automatic feature ranking and selection, quantifying how well reduced representations preserve neighborhood structures.
• It employs soft-nearest-neighbor weighting, gradient-based optimization, and ℓ1 regularization to align feature scales and enhance interpretability.
• DII has proven effective in applications such as molecular systems, nonlinear causal discovery, and unsupervised feature selection while managing computational costs.

Searching arXiv for the cited DII papers and closely related work. Differentiable Information Imbalance (DII) is a differentiable framework for ranking, weighting, and selecting features by measuring how well neighborhood relations in one representation of a dataset preserve those of another representation. It was introduced for automatic feature selection and weighting in molecular systems, where the central problem is to identify a low-dimensional subset of features that retains essential information from a full or ground-truth space while also aligning units and learning relative feature importance (Wild et al., 2024). In the subsequent literature, DII has been used for nonlinear causal discovery, unsupervised mode extraction from nonadiabatic molecular dynamics, and unsupervised feature-selection analyses with a statistical-physics interpretation (Salvagnin et al., 21 Aug 2025, Banerjee et al., 8 May 2026, Fiorentino et al., 31 Jan 2026).

1. Information imbalance and the differentiable relaxation

The underlying object is the Information Imbalance, denoted $\Delta$, defined between two representations of the same $N$ data points. In the 2024 formulation, if $$X^{(A)}=\{X_i^{(A)}\in\mathbb{R}^{D_A}\}_{i=1\dots N}$$ and $$X^{(B)}=\{X_i^{(B)}\in\mathbb{R}^{D_B}\}_{i=1\dots N}$$ are two spaces with distances $d_{ij}^{(A)}$ and $d_{ij}^{(B)}$, and if $r_{ij}^{(A)}$ is the rank of $j$ among all points ordered by increasing $d_{i\cdot}^{(A)}$, then

$$\Delta(d^{(A)}\rightarrow d^{(B)})= \frac{2}{N^2}\sum_{i=1}^N\sum_{j:r_{ij}^{(A)}=1} r_{ij}^{(B)}.$$

This quantity lies between $N$0 and $N$1: $N$2 corresponds to perfect preservation of nearest neighbors, whereas $N$3 corresponds to random ranks in $N$4 (Wild et al., 2024).

DII replaces the hard nearest-neighbor condition with a soft-nearest-neighbor weight,

$N$5

with

$N$6

where $N$7 is a softmax temperature. In the limit $N$8, one recovers the original information imbalance exactly (Wild et al., 2024).

The interpretation given in the original work is that DII measures how much information is lost when neighborhood relations in a full or ground-truth space $N$9 are approximated by those in a reduced weighted space $$X^{(A)}=\{X_i^{(A)}\in\mathbb{R}^{D_A}\}_{i=1\dots N}$$0. A weighted $$X^{(A)}=\{X_i^{(A)}\in\mathbb{R}^{D_A}\}_{i=1\dots N}$$1-space with DII equal to $$X^{(A)}=\{X_i^{(A)}\in\mathbb{R}^{D_A}\}_{i=1\dots N}$$2 perfectly preserves nearest-neighbor orders of $$X^{(A)}=\{X_i^{(A)}\in\mathbb{R}^{D_A}\}_{i=1\dots N}$$3, whereas DII near $$X^{(A)}=\{X_i^{(A)}\in\mathbb{R}^{D_A}\}_{i=1\dots N}$$4 is no better than random (Wild et al., 2024).

2. Weighted distances and alternative formulations in the literature

In the original molecular feature-selection setting, the reduced space is parameterized by a positive weight vector $$X^{(A)}=\{X_i^{(A)}\in\mathbb{R}^{D_A}\}_{i=1\dots N}$$5,

$$X^{(A)}=\{X_i^{(A)}\in\mathbb{R}^{D_A}\}_{i=1\dots N}$$6

Optimizing $$X^{(A)}=\{X_i^{(A)}\in\mathbb{R}^{D_A}\}_{i=1\dots N}$$7 rescales each feature and therefore performs unit alignment and relative importance scaling simultaneously while preserving interpretability (Wild et al., 2024).

Subsequent work retains the same core idea but introduces related formulations. In nonlinear causal discovery for European Union Allowances returns, DII is defined for predictor space $$X^{(A)}=\{X_i^{(A)}\in\mathbb{R}^{D_A}\}_{i=1\dots N}$$8 and scalar target $$X^{(A)}=\{X_i^{(A)}\in\mathbb{R}^{D_A}\}_{i=1\dots N}$$9 through

$$X^{(B)}=\{X_i^{(B)}\in\mathbb{R}^{D_B}\}_{i=1\dots N}$$0

and a soft nearest-neighbor weight

$$X^{(B)}=\{X_i^{(B)}\in\mathbb{R}^{D_B}\}_{i=1\dots N}$$1

with the same rank-based objective $$X^{(B)}=\{X_i^{(B)}\in\mathbb{R}^{D_B}\}_{i=1\dots N}$$2 (Salvagnin et al., 21 Aug 2025). That work also introduces the Imbalance Gain,

$$X^{(B)}=\{X_i^{(B)}\in\mathbb{R}^{D_B}\}_{i=1\dots N}$$3

as a criterion for whether adding a candidate predictor improves preservation of target neighborhoods (Salvagnin et al., 21 Aug 2025).

A different but closely related variant appears in work on unsupervised feature selection. There, the weighted squared Euclidean distance is

$$X^{(B)}=\{X_i^{(B)}\in\mathbb{R}^{D_B}\}_{i=1\dots N}$$4

and DII is defined as the average Kullback–Leibler divergence between a reference soft-kNN distribution and the $$X^{(B)}=\{X_i^{(B)}\in\mathbb{R}^{D_B}\}_{i=1\dots N}$$5-induced soft-kNN distribution:

$$X^{(B)}=\{X_i^{(B)}\in\mathbb{R}^{D_B}\}_{i=1\dots N}$$6

This formulation is nonnegative and vanishes only if the two neighbor distributions coincide (Fiorentino et al., 31 Jan 2026).

These formulations share the same operational aim: optimize a differentiable neighborhood-preservation objective over feature weights. The literature therefore uses the term DII for a family of differentiable relaxations of Information Imbalance rather than a single fixed equation. This suggests that the defining feature of DII is the differentiable optimization of neighborhood preservation, not one unique parametrization.

3. Optimization, regularization, and sparsity

For $$X^{(B)}=\{X_i^{(B)}\in\mathbb{R}^{D_B}\}_{i=1\dots N}$$7, the DII loss is smooth and can be minimized by gradient-based optimization. In the original formulation, with $$X^{(B)}=\{X_i^{(B)}\in\mathbb{R}^{D_B}\}_{i=1\dots N}$$8, the gradient is given analytically by

$$X^{(B)}=\{X_i^{(B)}\in\mathbb{R}^{D_B}\}_{i=1\dots N}$$9

and plain gradient descent updates $d_{ij}^{(A)}$0 (Wild et al., 2024).

Two sparsity-inducing strategies are described in the original paper. The first is greedy backward selection: optimize $d_{ij}^{(A)}$1 using all $d_{ij}^{(A)}$2 features, zero out the smallest-magnitude $d_{ij}^{(A)}$3, remove that feature, re-optimize, and repeat until only one feature remains. The second is $d_{ij}^{(A)}$4-regularized DII, minimizing

$d_{ij}^{(A)}$5

with a two-step “GD-clipping” update to enforce exact zeros (Wild et al., 2024). As $d_{ij}^{(A)}$6 increases from $d_{ij}^{(A)}$7 to large values, the number of nonzero features decreases, and the optimal sparsity can be chosen by inspecting the “elbow” or minimal DII across cardinalities (Wild et al., 2024).

Later work adapts the optimization machinery to different settings. The causal-discovery study standardizes each predictor coordinate to unit variance, initializes $d_{ij}^{(A)}$8, uses mini-batches, defines an adaptive neighborhood scale $d_{ij}^{(A)}$9 with $d_{ij}^{(B)}$0, computes gradients via automatic differentiation, and updates weights with Adam and a cosine learning-rate schedule (Salvagnin et al., 21 Aug 2025). The unsupervised feature-selection formulation with KL divergence likewise uses projected gradient descent or an Adam-type optimizer, together with nonnegativity and optionally a simplex constraint $d_{ij}^{(B)}$1 (Fiorentino et al., 31 Jan 2026).

Across these variants, a consistent point is that the learned weights have a dual role: they indicate feature relevance and rescale coordinates. In the terminology of the original paper, DII therefore performs automatic unit alignment and relative weighting at the same time (Wild et al., 2024).

4. Algorithmic workflow, implementation, and computational cost

A canonical workflow begins by fixing a reference or ground-truth space $d_{ij}^{(B)}$2, precomputing its ranks or soft neighborhood structure, initializing feature weights, and iteratively minimizing the DII loss. In the original molecular formulation, the high-level pseudocode is: precompute the rank matrix $d_{ij}^{(B)}$3 from distances in $d_{ij}^{(B)}$4; initialize $d_{ij}^{(B)}$5 as the inverse standard deviation of $d_{ij}^{(B)}$6; at each epoch compute all pairwise distances in $d_{ij}^{(B)}$7, compute the soft weights $d_{ij}^{(B)}$8, evaluate the loss, compute the analytic gradient, and update $d_{ij}^{(B)}$9 by gradient descent, with optional $r_{ij}^{(A)}$0 clipping or greedy elimination (Wild et al., 2024).

The naive cost of the original algorithm is $r_{ij}^{(A)}$1 per epoch because it requires all pairwise distances and soft weights. The same paper describes a subsampling “row-trick,” in which one restricts $r_{ij}^{(A)}$2 to $r_{ij}^{(A)}$3 rows, reducing the cost to $r_{ij}^{(A)}$4 or $r_{ij}^{(A)}$5 with minimal loss of accuracy (Wild et al., 2024). In the causal-discovery implementation, each gradient step costs $r_{ij}^{(A)}$6 within a mini-batch, plus $r_{ij}^{(A)}$7 for the softmax, so the total cost is $r_{ij}^{(A)}$8; with $r_{ij}^{(A)}$9, $j$0, $j$1, and $j$2, the computation is reported as feasible on modern GPUs and CPUs (Salvagnin et al., 21 Aug 2025).

The original method is available in the Python library DADApy, including a function return_weights_optimize_dii for feature weighting and selection, and a GPU-accelerated variant return_weights_optimize_dii_jax in dadapy.jax_feature_weighting (Wild et al., 2024). In the photochemical workflow, optimization is also described as being carried out with standard gradient-based packages such as Adam in the DADApy library (Banerjee et al., 8 May 2026).

Implementation details differ with the application. In causal discovery, one precomputes $j$3 once and recomputes $j$4 in each batch (Salvagnin et al., 21 Aug 2025). In unsupervised feature selection, one computes reference neighbor distributions from the full space and, if desired, repeats the optimization for different $j$5 or along a feature-elimination path to record $j$6 as a function of the number of retained features $j$7 (Fiorentino et al., 31 Jan 2026).

5. Representative applications and empirical findings

DII has been applied in benchmark molecular tasks, nonlinear time-series analysis, photochemical mechanism extraction, and unsupervised protein feature selection.

Domain	Setting	Reported outcome
Molecular benchmarks	Ten Gaussian features; 285 monomials of 10 Gaussians	cosine overlap $j$8, DII $j$9; with mild $d_{i\cdot}^{(A)}$0, perfect recovery of top-5 $d_{i\cdot}^{(A)}$1; recovery of $d_{i\cdot}^{(A)}$2 active monomials with cosine-overlap $d_{i\cdot}^{(A)}$3 and lowest DII
Biomolecular collective variables	CLN025 hairpin, $d_{i\cdot}^{(A)}$4 frames, candidate CVs $d_{i\cdot}^{(A)}$5	best 3-plet $d_{i\cdot}^{(A)}$6 with $d_{i\cdot}^{(A)}$7 and validation cross-val DII $d_{i\cdot}^{(A)}$8
ML force fields	$d_{i\cdot}^{(A)}$9 liquid $$\Delta(d^{(A)}\rightarrow d^{(B)})= \frac{2}{N^2}\sum_{i=1}^N\sum_{j:r_{ij}^{(A)}=1} r_{ij}^{(B)}.$$0 atomic environments; 176 ACSFs vs 546 SOAP descriptors	$$\Delta(d^{(A)}\rightarrow d^{(B)})= \frac{2}{N^2}\sum_{i=1}^N\sum_{j:r_{ij}^{(A)}=1} r_{ij}^{(B)}.$$1–$$\Delta(d^{(A)}\rightarrow d^{(B)})= \frac{2}{N^2}\sum_{i=1}^N\sum_{j:r_{ij}^{(A)}=1} r_{ij}^{(B)}.$$2 ACSFs suffice for near-optimal DII; at 20 ACSFs, RMSE $$\Delta(d^{(A)}\rightarrow d^{(B)})= \frac{2}{N^2}\sum_{i=1}^N\sum_{j:r_{ij}^{(A)}=1} r_{ij}^{(B)}.$$3; at 50 ACSFs, runtime reduced by $$\Delta(d^{(A)}\rightarrow d^{(B)})= \frac{2}{N^2}\sum_{i=1}^N\sum_{j:r_{ij}^{(A)}=1} r_{ij}^{(B)}.$$4 with identical accuracy
Causal discovery	EUA returns with 34 candidate variables, $$\Delta(d^{(A)}\rightarrow d^{(B)})= \frac{2}{N^2}\sum_{i=1}^N\sum_{j:r_{ij}^{(A)}=1} r_{ij}^{(B)}.$$5	top IGs highlight IBEX35 index and coal futures; DII detects additional nonlinear contributors missed by VAR
Photochemical decay	Methaniminium, furan, L-glutamine, L-pyroglutamine-ammonium, molecular motor	known mechanistic coordinates recovered; energy gaps often governed by a small number of localized coordinates, oscillator strengths by more collective rearrangements
Unsupervised protein feature selection	physico-chemical and structural descriptors	critical feature number from DII coincides with saturation of downstream binary classification performance for physico-chemical descriptors

In the original feature-selection benchmarks, DII without regularization recovered known ground-truth weights in synthetic Gaussian data with cosine overlap $$\Delta(d^{(A)}\rightarrow d^{(B)})= \frac{2}{N^2}\sum_{i=1}^N\sum_{j:r_{ij}^{(A)}=1} r_{ij}^{(B)}.$$6 and DII $$\Delta(d^{(A)}\rightarrow d^{(B)})= \frac{2}{N^2}\sum_{i=1}^N\sum_{j:r_{ij}^{(A)}=1} r_{ij}^{(B)}.$$7. With mild $$\Delta(d^{(A)}\rightarrow d^{(B)})= \frac{2}{N^2}\sum_{i=1}^N\sum_{j:r_{ij}^{(A)}=1} r_{ij}^{(B)}.$$8 regularization, only five active features remained and the top five ground-truth weights were perfectly recovered. In a harder benchmark involving 285 monomials of 10 Gaussians with true support 10, $$\Delta(d^{(A)}\rightarrow d^{(B)})= \frac{2}{N^2}\sum_{i=1}^N\sum_{j:r_{ij}^{(A)}=1} r_{ij}^{(B)}.$$9-DII with $N$00 recovered $N$01 active monomials with cosine-overlap $N$02 and lowest DII, outperforming relief-based filters such as RReliefF and MultiSURF and also outperforming embedded tree-regressor importance (Wild et al., 2024).

For collective-variable discovery in a peptide free-energy landscape, the system was a $N$03 T-REMD simulation of the CLN025 hairpin with $N$04 frames. Using all 4,278 pairwise heavy-atom distances as the ground-truth space and ten candidate CVs, exhaustive search over all $N$05 subsets followed by DII optimization yielded the best 3-plet $N$06 with weights $N$07 and validation cross-val DII $N$08. Clustering in this 3D space produced two dominant states, “$N$09-pin” and “collapsed,” with populations matching full-space clustering with $N$10 and $N$11 purity; the single best CV, anti-$N$12, attained only $N$13 purity, and tree-regressor features gave $N$14 purity (Wild et al., 2024).

For machine-learning force fields, DII was used to select among 176 ACSF descriptors against 546 SOAP descriptors on $N$15 atomic environments of liquid water. Both greedy-DII and $N$16-DII found that $N$17–$N$18 ACSFs suffice to reach near-optimal DII. A Behler–Parrinello MLP trained on DII-selected ACSFs achieved RMSE $N$19 at 20 ACSFs, matching the full 176-ACSF model, and at 50 ACSFs the runtime was reduced by $N$20 with identical accuracy; random ACSF subsets performed significantly worse at small sizes (Wild et al., 2024).

In nonlinear causal discovery for EUA returns, DII was compared with multivariate Granger causality on data from January 2013 to April 2024. Synthetic experiments showed that DII plus Imbalance Gain detected a nonlinear causal link $N$21 in a toy model where VAR(1) failed, and avoided a false positive $N$22 in a common-driver scenario where VAR(1) spuriously reported it. In the empirical EU Allowances study, both VAR and DII identified IBEX35 and coal futures as the strongest predictors, but DII also highlighted additional nonlinear contributors such as Silver futures and VSTOXX volatility with IG $N$23 despite small F-statistics (Salvagnin et al., 21 Aug 2025).

In photochemical applications, DII was embedded in a multi-step workflow that first reduces Coulomb-matrix descriptors by Jensen–Shannon divergence hotspots, then filters those hotspots with DII, maps them to interpretable internal coordinates through a second DII step, and finally ranks the internal coordinates by DII against observables such as energy gaps and oscillator strengths. Across methaniminium cation, furan, L-glutamine, L-pyroglutamine-ammonium, and an overcrowded alkene molecular motor, the method recovered known mechanistic coordinates and showed a systematic trend: energy gaps are often governed by a small number of localized coordinates, whereas oscillator strengths require more and more distributed modes (Banerjee et al., 8 May 2026).

In unsupervised protein feature selection, DII was treated as an order parameter over many random subsamples. For physico-chemical descriptors, the empirical distribution of optimal DII values became bimodal at small feature counts and the Binder cumulant $N$24 exhibited pronounced minima, indicating a glass-like to liquid-like transition as a function of retained features. The extrapolated critical number of features coincided with the saturation of downstream binary classification AUROC; for structural descriptors, by contrast, the transition was less sharp and the AUROC grew steadily with no clear plateau (Fiorentino et al., 31 Jan 2026).

6. Scope, limitations, and points of interpretation

The original paper lists several advantages of DII: it is a fully differentiable filter method, it does not require exponential subset enumeration, it handles multi-dimensional continuous ground-truth spaces, it simultaneously learns unit scaling and feature importance, and it induces sparsity through either $N$25 regularization or greedy backward removal (Wild et al., 2024). In the causal-discovery setting, further stated advantages are that DII is nonlinear, model-free, and able to test whether adding a variable improves prediction of a target beyond what other variables already provide (Salvagnin et al., 21 Aug 2025).

The same sources also specify limitations. In the original formulation, DII requires a ground-truth metric unless one uses the $N$26 unsupervised mode, in which case the choice of ground-truth weights is ambiguous; nominal or binary ground-truth features yield degenerate rank matrices; and the naive cost is $N$27, although row-subsampling mitigates this (Wild et al., 2024). In the causal-discovery implementation, optimization is explicitly nonconvex; the claim made there is not a formal global convergence guarantee but rather that Adam plus mini-batching and neighborhood-adaptive $N$28 empirically reached a good local minimum, with all synthetic and real-data tests converging within 2,000 epochs (Salvagnin et al., 21 Aug 2025).

A recurring interpretive point in the literature is that low DII indicates strong predictive relevance and high DII indicates weak relevance, but the operational threshold depends on the application. In the photochemical study, the final DII scores are interpreted as follows: low DII $N$29 indicates strongly predictive or near-optimal coordinates, medium DII $N$30 indicates moderate coupling, and high DII $N$31 indicates weak or no predictive relevance (Banerjee et al., 8 May 2026). This suggests that DII is best used comparatively, ranking coordinates or feature sets within a fixed problem rather than relying on a universal cross-domain threshold.

Another important point is that the literature does not present DII as a single immutable metric. Rank-based soft-nearest-neighbor objectives, adaptive batchwise versions for causal discovery, and KL-divergence formulations over soft-kNN distributions all appear under the same label (Wild et al., 2024, Salvagnin et al., 21 Aug 2025, Fiorentino et al., 31 Jan 2026). A plausible implication is that DII is most coherently understood as a differentiable neighborhood-preservation principle for feature weighting and selection. Under that interpretation, its various instantiations form a common methodology for constructing compact, interpretable representations that preserve the local informational geometry of a chosen reference space.