Coupled Transfer Distance, Interactive

← text notes · Gao & Chaudhari, An Information-Geometric Distance on the Space of Tasks (ICML 2021). Eight widgets. The numerical ones recompute, in your browser, the same quantities as the scripts in code/; the toy transfer problem is the one in transfer_toy.py.

1. The ruler: what one step costs

Two classifiers are close if they predict alike, not if their weights are close. For a binary classifier with logit $z$ and $p=\sigma(z)$, moving the logit by $\delta$ costs $2\,\mathrm{KL}(p_z\Vert p_{z+\delta})$, and for a small step this is $\delta^2\,p(1-p)$: the Fisher information of a coin. A confident classifier ($p$ near 0 or 1) barely changes its predictions when its logit moves; one at $p=\tfrac12$ changes the most.

exact $2\,\mathrm{KL}$quadratic $\delta^2p(1-p)$

2. Distances between predictions live on a sphere

For a three-class softmax the Fisher–Rao distance is $d(p,q)=2\arccos\sum_i\sqrt{p_iq_i}$ (the sphere of radius 2, seen through $u=\sqrt p$). Drag the two points. The blue curve is the geodesic; the dashed circle is every prediction at the same Fisher–Rao distance from the blue point, so it passes through the orange one. Near an edge the circle is squashed, and it stops at the edge: a small change in a rare class is expensive.

geodesicball through the orange point

3. Where the average over x sits

The Fisher matrix (3) averages the KL over the inputs inside the square root, $A=\int\sqrt{\mathbb E_x[2\mathrm{KL}_x]}$. Definition 2 averages the square root, $B=\mathbb E_x\int\sqrt{2\mathrm{KL}_x}$. Take a logistic model whose weights move in a straight line from $(1,0)$ to $(a_1,b_1)$, with inputs $x\sim N(0,\sigma_x^2)$. In the extreme where a fraction $f$ of the inputs feel the change and the rest are saturated and feel nothing, $B/A=\sqrt f$; in between it lies between $\sqrt f$ and 1 (Jensen).

density of $x$how much a logit step at $x$ changes the prediction, $\sqrt{p(1-p)}$ (final weights)

4. Moving a task, and the weights that follow it

A toy transfer problem in one dimension. The source has a class at $-2$ and a class at $+2$ (sd 1); the target is the same task translated by $s$. The classifier is logistic regression $\sigma(w_1x+w_0)$, and the weights follow the minimiser at each instant $\tau$ (the limit of eq. 10 for infinitely fast learning). Displacement moves each class by optimal transport; mixture draws from the source with probability $1-\tau$ and from the target otherwise (eq. 9). The pictures show the task at time $\tau$; the tables show the three lengths $A$, $B$ and $TV$ (see the notes) of the whole path.

displacementmixture
For $s\gtrsim1$ the mixture path is longer by $B$ and by $TV$, as Remark 1 says. By the Riemannian length $A$ it is longer only up to about $s=4$; for large $s$ it is shorter, because the mixture's minimiser goes through a much flatter classifier (its slope falls from 4 to 0.8 at $s=4$ and to 0.2 at $s=8$) whose predictions are cheap to move. Remark 1 is an expectation, not a theorem.

5. How fast the weights learn, and asymmetry

Eq. (10) follows the task with SGD; in the population limit that is $\dot w=-\kappa\nabla L_\tau(w)$, where $\kappa$ is the number of steps times the learning rate. Here the tasks move along a displacement path. For small $\kappa$ the weights barely move: the path is short and the fit is bad. The minimiser path is exactly reversible, so forward and backward lengths agree; the SGD path is not, and the two directions differ until $\kappa$ is very large.

SGD path (this $\kappa$)minimiser pathmidpoint of the class means

6. The empirical Fisher in eq. (21)

To turn the length into a change of loss, the proof replaces the model's own label distribution by the observed label (a Dirac delta). For one binary input, the true speed is $\sqrt{p(1-p)}\,|\dot z|$ and the empirical one is $(1-p_y)\,|\dot z|$, where $p_y$ is the probability the model gives to the observed label. Their ratio is $\sqrt{(1-p_y)/p_y}$.

7. Theorem 6 as printed, and with N

The proof ends with $\Pr\{\tfrac1K\sum_k\text{gap}_k>\epsilon\}\le\exp\{-\tfrac{2K}{M^2}(\epsilon-2\bar R)^2\}$ (15), where $\bar R$ is the average Rademacher complexity. Here $N$ never appears in the exponent; it enters only through $\bar R$. Applying McDiarmid's inequality instead of Hoeffding's lemma to $\varphi_k$ puts a factor $N$ there. The last row shows what the printed $\lambda=4K(\epsilon-2\bar R)/M^2$ does to the exponent.

printed (15)with McDiarmid, $\exp\{-2NK(\epsilon-2\bar R)^2/M^2\}$

8. Re-running the Mantel tests of Figures 2 and 3

The distance matrices printed in the paper's Figures 2 and 3, compared with the fine-tuning matrix. With four tasks there are only $4!=24$ ways to relabel them, so no Mantel test can return a $p$-value below $1/24\approx0.042$; with five tasks the floor is $1/120$. The bars are the correlation $r$ for every relabelling; the line is the correlation for the true labelling.