---
title: Differentiable Information Imbalance (DII)
url: https://www.emergentmind.com/topics/differentiable-information-imbalance-dii
type: topic
---

# Differentiable Information Imbalance (DII)

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 [2411.00851]. 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 [2508.15667], [2605.08381], [2602.00660].

## 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 $0$ and $1$: $\Delta=0$ corresponds to perfect preservation of nearest neighbors, whereas $\Delta\approx 1$ corresponds to random ranks in $B$ [2411.00851].

DII replaces the hard nearest-neighbor condition with a soft-nearest-neighbor weight,
$$
\mathrm{DII}(d^{(A)}(w)\rightarrow d^{(B)})=
\frac{2}{N^2}\sum_{i\neq j} c_{ij}(\lambda,d^{(A)}(w))\,r_{ij}^{(B)},
$$
with
$$
c_{ij}(\lambda,d^{(A)})=
\frac{\exp[-d_{ij}^{(A)}/\lambda]}
{\sum_{m\neq i}\exp[-d_{im}^{(A)}/\lambda]},
$$
where $\lambda>0$ is a softmax temperature. In the limit $\lambda\to 0$, one recovers the original information imbalance exactly [2411.00851].

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 $B$ are approximated by those in a reduced weighted space $A$. A weighted $A$-space with DII equal to $0$ perfectly preserves nearest-neighbor orders of $B$, whereas DII near $1$ is no better than random [2411.00851].

## 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 $w\in\mathbb{R}^{D_A}$,
$$
d_{ij}^{(A)}(w)=\|\,w\odot(X_i^{(A)}-X_j^{(A)})\,\|_2.
$$
Optimizing $w$ rescales each feature and therefore performs unit alignment and relative importance scaling simultaneously while preserving interpretability [2411.00851].

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$ and scalar target $Z$ through
$$
d_{ij}^{X}(w)=\|\,w\odot(x_i-x_j)\,\|_2,\qquad d_{ij}^{Z}=|z_i-z_j|,
$$
and a soft nearest-neighbor weight
$$
c_{ij}(w,\lambda)=\frac{\exp(-d_{ij}^{X}(w)^2/\lambda)}
{\sum_{m\neq i}\exp(-d_{im}^{X}(w)^2/\lambda)},
$$
with the same rank-based objective $\frac{2}{N^2}\sum_{i\neq j} c_{ij}(w,\lambda)\,r_{ij}^{Z}$ [2508.15667]. That work also introduces the Imbalance Gain,
$$
\mathrm{IG}=1-\frac{\min_w \mathrm{DII}(\mathrm{baseline}\rightarrow Z)}
{\min_w \mathrm{DII}(\mathrm{full}\rightarrow Z)},
$$
as a criterion for whether adding a candidate predictor improves preservation of target neighborhoods [2508.15667].

A different but closely related variant appears in work on unsupervised feature selection. There, the weighted squared Euclidean distance is
$$
d_{ij}(w)=\sum_{d=1}^D w_d(x_{i,d}-x_{j,d})^2,
$$
and DII is defined as the average Kullback–Leibler divergence between a reference soft-kNN distribution and the $w$-induced soft-kNN distribution:
$$
\mathrm{DII}(w)=\frac{1}{N}\sum_{i=1}^N\sum_{j\neq i}
p_{j|i}^{\mathrm{ref}}\ln\frac{p_{j|i}^{\mathrm{ref}}}{p_{j|i}(w)}.
$$
This formulation is nonnegative and vanishes only if the two neighbor distributions coincide [2602.00660].

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 $\lambda>0$, the DII loss is smooth and can be minimized by gradient-based optimization. In the original formulation, with $L(w)=\mathrm{DII}(d^{(A)}(w)\rightarrow d^{(B)})$, the gradient is given analytically by
$$
\frac{\partial L}{\partial w^\alpha}
=\frac{2}{\lambda N^2}w^\alpha
\sum_{i\neq j} c_{ij} r_{ij}^{(B)}
\left[
-\frac{(\Delta X_{ij}^{\alpha})^2}{\|w\odot \Delta X_{ij}\|}
+\sum_{m\neq i} c_{im}\frac{(\Delta X_{im}^{\alpha})^2}{\|w\odot \Delta X_{im}\|}
\right],
$$
and plain gradient descent updates $w\leftarrow w-\eta\nabla L(w)$ [2411.00851].

Two sparsity-inducing strategies are described in the original paper. The first is greedy backward selection: optimize $w$ using all $D_A$ features, zero out the smallest-magnitude $w^\alpha$, remove that feature, re-optimize, and repeat until only one feature remains. The second is $\ell_1$-regularized DII, minimizing
$$
L(w)+p\sum_\alpha |w^\alpha|,
$$
with a two-step “GD-clipping” update to enforce exact zeros [2411.00851]. As $p$ increases from $0$ 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 [2411.00851].

Later work adapts the optimization machinery to different settings. The causal-discovery study standardizes each predictor coordinate to unit variance, initializes $w\leftarrow 1_D$, uses mini-batches, defines an adaptive neighborhood scale $\lambda_i=\alpha\,(d_{i,j(k)}^{X}(w))^2$ with $\alpha\approx 0.1$, computes gradients via automatic differentiation, and updates weights with Adam and a cosine learning-rate schedule [2508.15667]. 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 $\sum_d w_d=1$ [2602.00660].

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 [2411.00851].

## 4. Algorithmic workflow, implementation, and computational cost

A canonical workflow begins by fixing a reference or ground-truth space $B$, 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 $R_{ij}^{(B)}$ from distances in $B$; initialize $w$ as the inverse standard deviation of $X^{(A)}$; at each epoch compute all pairwise distances in $A$, compute the soft weights $c_{ij}$, evaluate the loss, compute the analytic gradient, and update $w$ by gradient descent, with optional $\ell_1$ clipping or greedy elimination [2411.00851].

The naive cost of the original algorithm is $O(N^2D)$ per epoch because it requires all pairwise distances and soft weights. The same paper describes a subsampling “row-trick,” in which one restricts $i$ to $N_0\ll N$ rows, reducing the cost to $O(ND)$ or $O(N_0D)$ with minimal loss of accuracy [2411.00851]. In the causal-discovery implementation, each gradient step costs $O(B^2D)$ within a mini-batch, plus $O(B^2)$ for the softmax, so the total cost is $O(T\cdot(\#\mathrm{batches})\cdot B^2D)$; with $N=2{,}800$, $B=100$, $D\approx 35$, and $T=2{,}000$, the computation is reported as feasible on modern GPUs and CPUs [2508.15667].

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` [2411.00851]. In the photochemical workflow, optimization is also described as being carried out with standard gradient-based packages such as Adam in the DADApy library [2605.08381].

Implementation details differ with the application. In causal discovery, one precomputes $r_{ij}^{Z}$ once and recomputes $\lambda_i$ in each batch [2508.15667]. In unsupervised feature selection, one computes reference neighbor distributions from the full space and, if desired, repeats the optimization for different $\lambda$ or along a feature-elimination path to record $\mathrm{DII}(w_F)$ as a function of the number of retained features $F$ [2602.00660].

## 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 $\approx 1$, DII $\approx 0$; with mild $\ell_1$, perfect recovery of top-5 $w_{\rm GT}$; recovery of $8/10$ active monomials with cosine-overlap $0.99$ and lowest DII |
| Biomolecular collective variables | CLN025 hairpin, $N=1{,}429$ frames, candidate CVs $D_A=10$ | best 3-plet $\{ \mathrm{RGYR}, \mathrm{PC1}, \mathrm{PC2} \}$ with $w=\{1.0,3.5,4.7\}$ and validation cross-val DII $\approx 0.05$ |
| ML force fields | $\sim 350$ liquid $\mathrm{H_2O}$ atomic environments; 176 ACSFs vs 546 SOAP descriptors | $\sim 20$–$30$ ACSFs suffice for near-optimal DII; at 20 ACSFs, RMSE $<100\,\mathrm{meV}/\AA$; at 50 ACSFs, runtime reduced by $\sim 1/3$ with identical accuracy |
| Causal discovery | EUA returns with 34 candidate variables, $N=2{,}902$ | 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 $\approx 1$ and DII $\approx 0$. With mild $\ell_1$ 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, $\ell_1$-DII with $p=10^{-4}$ recovered $8/10$ active monomials with cosine-overlap $0.99$ and lowest DII, outperforming relief-based filters such as RReliefF and MultiSURF and also outperforming embedded tree-regressor importance [2411.00851].

For collective-variable discovery in a peptide free-energy landscape, the system was a $1.4\,\mu\mathrm{s}$ T-REMD simulation of the CLN025 hairpin with $N=1{,}429$ frames. Using all 4,278 pairwise heavy-atom distances as the ground-truth space and ten candidate CVs, exhaustive search over all $2^{10}-1$ subsets followed by DII optimization yielded the best 3-plet $\{\mathrm{RGYR},\mathrm{PC1},\mathrm{PC2}\}$ with weights $\{1.0,3.5,4.7\}$ and validation cross-val DII $\approx 0.05$. Clustering in this 3D space produced two dominant states, “$\beta$-pin” and “collapsed,” with populations matching full-space clustering with $87\%$ and $92\%$ purity; the single best CV, anti-$\beta$, attained only $45\%$ purity, and tree-regressor features gave $<65\%$ purity [2411.00851].

For machine-learning force fields, DII was used to select among 176 ACSF descriptors against 546 SOAP descriptors on $\sim 350$ atomic environments of liquid water. Both greedy-DII and $\ell_1$-DII found that $\sim 20$–$30$ ACSFs suffice to reach near-optimal DII. A Behler–Parrinello MLP trained on DII-selected ACSFs achieved RMSE $<100\,\mathrm{meV}/\AA$ at 20 ACSFs, matching the full 176-ACSF model, and at 50 ACSFs the runtime was reduced by $\sim 1/3$ with identical accuracy; random ACSF subsets performed significantly worse at small sizes [2411.00851].

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 $x^2\rightarrow z$ in a toy model where VAR(1) failed, and avoided a false positive $x^1\rightarrow z$ 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 $>0$ despite small F-statistics [2508.15667].

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 [2605.08381].

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 $U_N(F)=1-\langle m^4\rangle/(3\langle m^2\rangle^2)$ 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 [2602.00660].

## 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 $\ell_1$ regularization or greedy backward removal [2411.00851]. 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 [2508.15667].

The same sources also specify limitations. In the original formulation, DII requires a ground-truth metric unless one uses the $A\rightarrow A$ 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 $O(N^2D)$, although row-subsampling mitigates this [2411.00851]. 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 $\lambda$ empirically reached a good local minimum, with all synthetic and real-data tests converging within 2,000 epochs [2508.15667].

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 $\lesssim 0.1{-}0.4$ indicates strongly predictive or near-optimal coordinates, medium DII $\sim 0.4{-}0.7$ indicates moderate coupling, and high DII $\gtrsim 0.7$ indicates weak or no predictive relevance [2605.08381]. 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 [2411.00851], [2508.15667], [2602.00660]. 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.

Source: https://www.emergentmind.com/topics/differentiable-information-imbalance-dii