Ablation studies, knockout studies, lesion studies

1 Experimental design in ML ablations

Suppose our method has \(p\) components, each of which we can switch on or off, and that we are not sure which ones provided the secret sauce to make the method go.

A configuration is one setting of all \(p\) switches, which we identify with the set \(S \subseteq \{1, \dots, p\}\) of components switched on; write \(x_j = 1\) if \(j \in S\) and \(x_j = 0\) otherwise. A design is the list of configurations we choose to run. Each run retrains the network from scratch at its configuration² and returns a score \(y = f(S) + \varepsilon\), where \(f(S) = \mathbb{E}[y \mid S]\) is the expected score and \(\varepsilon\) is run-to-run noise from seeds, data order and so on. So a design requires of us one full training run per configuration in it; each run yields one noisy score \(y\), and we estimate \(f\) from all of them together. That is to say, we are thinking about a model for the \(\{0,1\}^p \to \mathbb{R}\) map from configurations \(S\) to expected validation score \(f(S)\). The details of the method’s fabulously complex interior are irrelevant for this purpose.

When we experiment on a method, the implied model is that \(f\) is linear in the indicators and their products,

\[f(S) = f(\emptyset) + \sum_{j} \beta_j x_j + \sum_{j<k} \beta_{jk} x_j x_k + \sum_{j<k<l} \beta_{jkl} x_j x_k x_l + \dots,\]

which is to say, an ordinary linear regression of the score on the indicators and their products. Taken all the way up to the \(p\)-way product \(x_1 x_2 \cdots x_p\), this has \(2^p\) coefficients for \(2^p\) configurations, so it can represent any \(f\) whatsoever; the assumptions enter when we truncate it. As in the physical sciences, most sane people truncate after the pairwise terms; we do too, and call what is left the pairwise model.

The quantities a design might target are all built from the gain in expected score from switching on \(j\) when the components in \(S\) are already on,

\[\Delta_j(S) = f(S \cup \{j\}) - f(S), \qquad j \notin S,\]

which experimental designers call a simple effect. Under the pairwise model, \(\Delta_j(S) = \beta_j + \sum_{k \in S} \beta_{jk}\). So \(\beta_j = \Delta_j(\emptyset)\) is the gain from switching on \(j\) alone, and the interaction \(\beta_{jk}\) is the extra gain from having \(j\) and \(k\) on together, over and above \(\beta_j + \beta_k\). The main effect \(m_j\) is the simple effect averaged over all configurations of the other components, i.e. with each other component on half the time; under the pairwise model, \(m_j = \beta_j + \tfrac{1}{2}\sum_{k \neq j} \beta_{jk}\). A regression table would label \(\beta_j\) the main effect of \(j\); the two coincide only when \(j\) interacts with nothing³.

A neural net ablation study typically removes one component at a time from the “best” configuration (by which the hopeful author means my configuration, their fingers crossed as the NeurIPS paper deadline looms large) and reports how much some score drops for each. In experimental design, this is called one-factor-at-a-time, and has been recommended against since Fisher (1926)⁴, for the following reasons:

1. It cannot detect interactions, and when there are interactions there is no canonical “effect of component \(i\)” to measure. Leave-one-out reflects a choice of convenience rather than principle. Under the pairwise model, the true gain from including \(i\) depends on what is already included, \(\Delta_i(S) = \beta_i + \sum_{j \in S} \beta_{ij}\), so each choice of baseline configuration \(S\) defines a different “effect”. Adding \(i\) to the bare configuration targets \(\Delta_i(\emptyset) = \beta_i\). Contrariwise, removing it from the full configuration targets \(\beta_i\) plus all of its interactions. The Shapley value targets \(\beta_i\) plus half of each. These all agree only when there are no interactions. Because a one-at-a-time study varies each component from our preferred baseline only, it reports that baseline’s answer. We are left to speculate about how different the answer would be from another baseline, or averaged over all of them. All of this argument applies to pairwise interactions but becomes increasingly funky as we add higher-order interactions.
2. For less diligent experimenters, at least, this introduces a confounding with the hyperparameter tuning budget. That is, having just completed an exceptionally expensive hyperparameter optimization at my configuration, I am not generally inclined to repeat the process for competing hypotheses that do not incorporate my sophisticated features, and in any case, might make the opposition look better in comparison. Melis, Dyer, and Blunsom (2017), for example, found that several claimed architecture improvements disappeared once every model received the same tuning budget.
3. It has no error bar.

A common excuse for this is that anything more sophisticated than these meagre studies is computationally prohibitive. Fitting models, is slow and expensive.

But, man, it’s not that pricey, and the reward is that it actually delivers what was claimed on the tin. And, you guys, ML is not the only field that has expensive experiments; we get no sympathy when we complain about the cost when other sciences have multi-million dollar wet labs and rare transgenic mice or whatever. Real scientists, with wet labs etc, have also faced prohibitive experiments and they have spent a century figuring out how to be as efficient as possible while actually obtaining the results we need. We can use their strategies.

A fractional factorial design varies all the components at once and estimates each main effect from every run. For twelve components, a thirty-two-run design (a “folded-over Plackett-Burman design”) gives main effects that are not contaminated by pairwise interactions, with runs to spare for estimating the noise. By contrast, leave-one-out needs thirteen runs, which is only 40% as many, sure, but its estimates mean nothing in particular unless an oracle wandered in and delivered stone tablets on which is graven LINEAR WITH NO INTERACTIONS.

The Shapley value is also a nice thing to target: it splits each interaction equally among its participants, so the shares add up to exactly the total. Under the pairwise model, the main effect \(m_i = \beta_i + \tfrac{1}{2}\sum_{j \neq i} \beta_{ij}\) is identical to the Shapley value of \(i\), so the thirty-two-run design delivers the Shapley attributions directly. Once three-way interactions enter, the two drift apart: the main effect credits \(i\) with a quarter of each \(\beta_{ijk}\), the Shapley value with a third. But when higher-order interactions matter, Shapley values can be directly estimated by sampling the orders in which components are added (Lundberg and Lee 2017).

2 Global sensitivity analysis might do the right thing actually

Global sensitivity analysis is not strictly ablation, but it can do the job of determining what is important rather well. Moreover, it is often free. If we are picking the model components as hyperparameters the right way, i.e. by Bayesian optimization, then we have already fit a surrogate model while selecting those hyperparameters. Ax, for example, can estimate variance-based (Sobol) sensitivity indices on that surrogate, i.e. the share of the variance in the surrogate’s predicted score attributable to each component (Oakley and O’Hagan 2004). Because the surrogate’s indices average over the hyperparameters rather than fixing them where the full model was tuned, they sidestep the tuning-budget confound above. As a bonus, this sometimes allows us to estimate sensitivities for continuous and discrete hyperparameters jointly.

3 References