\documentclass[11pt]{article}
\usepackage[margin=1.2in]{geometry}
\usepackage{amsmath, amssymb, amsthm}
\usepackage{enumitem}
\usepackage{xcolor}
\usepackage{framed}
\usepackage{tcolorbox}
\usepackage{booktabs}
\usepackage{natbib}

\newtheorem{definition}{Definition}
\newtheorem{lemma}{Lemma}
\newtheorem{theorem}{Theorem}

\definecolor{keyinsight}{RGB}{230, 245, 255}
\definecolor{intuition}{RGB}{255, 245, 230}
\definecolor{example}{RGB}{240, 255, 240}
\definecolor{warning}{RGB}{255, 240, 240}

\newtcolorbox{keyinsightbox}{
  colback=keyinsight,
  colframe=blue!50!black,
  title=Key Insight,
  fonttitle=\bfseries
}

\newtcolorbox{intuitionbox}{
  colback=intuition,
  colframe=orange!50!black,
  title=Intuition,
  fonttitle=\bfseries
}

\newtcolorbox{examplebox}{
  colback=example,
  colframe=green!50!black,
  title=Concrete Example,
  fonttitle=\bfseries
}

\newtcolorbox{warningbox}{
  colback=warning,
  colframe=red!50!black,
  title=Caution,
  fonttitle=\bfseries
}

\title{Notes on susceptibilities}
\author{Andy Arditi \& Claude 4.5 Opus}
\date{January 30, 2026}

\begin{document}
\maketitle

\begin{abstract}
The \emph{susceptibility} of an observable to a data perturbation measures how that observable---a function of the model weights---responds when we shift the data distribution. If we trained on slightly different data, how would this quantity change? \citet{baker2025susceptibilities} show that we can answer this question without retraining: the susceptibility equals a covariance, computed by sampling from the posterior induced by the \emph{original} data.

One natural application: for each model component (attention head, MLP layer), define an observable measuring that component's contribution to loss. Computing susceptibilities for all components yields a vector per datapoint---a profile of which components matter for that prediction, and by how much. Crucially, datapoints that the model processes similarly will have similar profiles. Clustering datapoints by these profiles therefore groups them by how the model computes on them, revealing both the patterns the model has learned to distinguish and the internal structure it uses to do so.
\end{abstract}

\tableofcontents
\newpage

%==============================================================================
\section{Background: the Bayesian view}
%==============================================================================

To understand susceptibilities, we need to think about things from a Bayesian perspective.

\subsection{Setup and notation}

\begin{itemize}[nosep]
    \item \textbf{Data:} $n$ training examples $D_n = \{(x_i, y_i)\}_{i=1}^n$ drawn from distribution $q(x, y)$
    \item \textbf{Model:} Probability $p(y|x, w)$ parameterized by weights $w$
    \item \textbf{Losses:}
    \begin{align*}
    L_n(w) &= -\frac{1}{n}\sum_{i=1}^n \log p(y_i | x_i, w) \quad \text{(empirical)} \\
    L(w) &= \mathbb{E}_{(x,y) \sim q}\left[-\log p(y | x, w)\right] \quad \text{(population)}
    \end{align*}
\end{itemize}

\subsection{The Bayesian view: a distribution over weights}

Normally, we think of training as finding a single $w^*$ that minimizes loss. The Bayesian perspective is different: our training data induces a \emph{probability distribution} over weights, where low-loss configurations are more probable.

\subsection{The posterior distribution}

From Bayes' rule:
\[
\underbrace{p(w | D_n)}_{\text{posterior}} \;\propto\; \underbrace{p(D_n | w)}_{\text{likelihood}} \cdot \underbrace{\varphi(w)}_{\text{prior}}.
\]
Since $p(D_n | w) = \prod_i p(y_i | x_i, w) = e^{-n L_n(w)}$, the posterior is
\[
p_n(w) = \frac{1}{Z_n} \cdot e^{-n \cdot L_n(w)} \cdot \varphi(w),
\]
where $\varphi(w)$ is the prior and $Z_n = \int e^{-n L_n(w)} \varphi(w) \, dw$ is the normalizing constant (also called the \emph{marginal likelihood} or \emph{evidence}).

\begin{intuitionbox}
Low loss $\Rightarrow$ high probability. With lots of data (large $n$), the distribution becomes sharply peaked around the best-fitting weights.
\end{intuitionbox}

\subsection{Adding temperature}

The \emph{tempered posterior} adds a parameter $\beta \geq 0$:
\[
p^\beta_n(w) = \frac{1}{Z^\beta_n} e^{-n\beta \cdot L_n(w)} \cdot \varphi(w), \qquad Z^\beta_n = \int e^{-n\beta L_n(w)} \varphi(w) \, dw.
\]
The parameter $\beta$ is called the \emph{inverse temperature}. It controls how much the data constrains our beliefs: $\beta = 0$ ignores the data entirely (giving the prior), $\beta = 1$ is standard Bayesian inference, and as $\beta \to \infty$ the distribution concentrates at the loss minimizer.

%==============================================================================
\section{Observables and susceptibilities}
%==============================================================================

\subsection{What is an observable?}

An \textbf{observable} $\phi(w)$ is any measurable function of the weights. Examples include:
\begin{itemize}[nosep]
    \item The total loss: $\phi(w) = L(w)$
    \item A single weight value: $\phi(w) = w_{17}$
    \item The norm of a layer: $\phi(w) = \|W_{\text{layer}}\|$
    \item How much loss increases when we perturb a specific component (see Section~\ref{sec:application})
\end{itemize}

We will denote the expected value of an observable under the posterior as $\langle \phi \rangle$:
\[
\langle \phi \rangle := \mathbb{E}_{w \sim p^\beta_n}\left[\phi(w)\right] = \int \phi(w) \cdot p^\beta_n(w) \, dw.
\]

\subsection{Perturbing the data distribution}

Let $q$ be the original data distribution and let $q'$ be some other data distribution of interest ($q'$ is sometimes referred to as a \textbf{probe} distribution). For example, $q'$ might be a distribution over GitHub code, or a single example.

We define a mix of these two distributions, parameterized by $h \in [0,1]$:
\[
q_h = (1-h) \cdot q + h \cdot q'.
\]

As we vary $h$, a chain of dependencies unfolds:
\[
q_h \;\xrightarrow{\text{changes}}\; L^h(w) \;\xrightarrow{\text{changes}}\; p^\beta_n(w|h) \;\xrightarrow{\text{changes}}\; \langle\phi\rangle_h
\]

\begin{definition}[Susceptibility]
The \emph{susceptibility of $\phi$ to the perturbation $q \to q_h$} is defined as
\[
\chi := \frac{1}{n\beta} \frac{\partial}{\partial h} \langle \phi \rangle_h \bigg|_{h=0}.
\]
It measures how fast the expected value of $\phi$ changes as we perturb the data toward $q'$.
\end{definition}

%==============================================================================
\section{Computing susceptibilities}
%==============================================================================

How do we actually compute $\chi$ in practice? The following lemma gives a feasible way to do so:

\begin{lemma}[Susceptibility as covariance]
\[
\chi = -\mathrm{Cov}_{w \sim p^\beta_n}\left[\phi(w), \; \Delta L(w)\right],
\]
where $\Delta L(w) = L^{q'}(w) - L^q(w)$ is the difference in loss between the probe and original distributions, and the covariance is over the posterior $p^\beta_n$.
\end{lemma}

This is useful because it means we can estimate $\chi$ from posterior samples: draw samples $w_1, \ldots, w_m$ from the posterior, evaluate $\phi(w_i)$ and $\Delta L(w_i)$ for each, and compute the sample covariance.

\textbf{Intuition:} Shifting data toward $q'$ reweights the posterior: weights that perform well on $q'$ (low $\Delta L$) gain probability. Whether $\langle\phi\rangle$ increases or decreases depends on whether high-$\phi$ weights are winners or losers under this reweighting. For example, if high-$\phi$ weights tend to have low $\Delta L$ (e.g., if there's negative covariance between $\phi$ and $\Delta L$), they gain probability, so $\langle\phi\rangle$ increases.

See Appendix~\ref{sec:proof} for the proof.

%==============================================================================
\section{Application: component-wise susceptibilities}
\label{sec:application}
%==============================================================================

The general theory above applies to any observable. A particularly useful application, developed in \cite{baker2025susceptibilities, wang2025embryology, gordon2026spectroscopy}, is to study how individual model components (attention heads, MLP layers) respond to data perturbations.

\subsection{The component-wise loss observable}

For neural network interpretability, we want to understand what each component does. A natural question: \emph{does component $C$ help predict pattern $P$ more than average?}

To answer this, we use a specific observable. Consider component $C$ (say, attention head 1:3). Write the weights as $w = (u, v)$ where $v$ are $C$'s weights and $u$ are everything else. Define
\[
\phi_C(w) = L(w) - L(w^*).
\]
This measures how much loss increases when we perturb away from the trained weights $w^*$.

To make this specific to component $C$, we \textbf{restrict the sampling}: only vary $C$'s weights while keeping everything else frozen at $u^*$. The expected value becomes
\[
\langle \phi_C \rangle = \int \phi_C(u^*, v) \cdot p^\beta_C(v) \, dv,
\]
where $p^\beta_C(v) \propto e^{-n\beta L(u^*, v)} \varphi(v)$ is a posterior over just $C$'s weights.

\textbf{Why this is useful:} If $\langle \phi_C \rangle$ increases when we shift data toward pattern $P$, it means perturbations to $C$ hurt more on $P$ than on average data. In other words: $C$ is disproportionately important for $P$. This lets us build a ``job description'' for each component---which patterns it helps with, which it ignores, and which it actively suppresses.

\paragraph{A note on notation.} The original papers write $\phi_C(w) = \delta(u - u^*)[L(w) - L(w^*)]$ with a Dirac delta. This formally expresses ``integrate only over $C$'s weights.'' Operationally, the delta isn't literally evaluated---we restrict sampling to $C$'s subspace.

\subsection{Per-token susceptibility}

A particularly useful special case: the susceptibility of component $C$ to a single example $(x, y)$. Taking the probe to be $q' = \delta_{(x,y)}$ and writing $\ell_{(x,y)}(w) = -\log p(y|x,w)$:
\[
\chi^C_{(x,y)} = -\mathrm{Cov}\left[\phi_C(w), \;\; \ell_{(x,y)}(w) - L(w)\right].
\]
\textbf{Interpretation:} When we perturb $C$, does $(x \to y)$ suffer more than average, less, or about average?

\begin{center}
\begin{tabular}{@{}cll@{}}
\toprule
\textbf{Sign} & \textbf{Term} & \textbf{Meaning} \\
\midrule
$\chi < 0$ & \textbf{Expression} & $C$ helps this prediction \textbf{more} than average \\
$\chi = 0$ & Neutral & $C$ helps about average \\
$\chi > 0$ & \textbf{Suppression} & $C$ helps \textbf{less} than average \\
\bottomrule
\end{tabular}
\end{center}

\begin{examplebox}
For ``The cat sat on the mat. The cat sat on the \underline{\hspace{0.5cm}}'', an induction head should have $\chi < 0$ (strong expression) for predicting ``mat.''
\end{examplebox}

Note: ``Suppression'' doesn't necessarily mean $C$ is hurting the prediction---it could just be irrelevant, or actively predicting something else.

\subsection{The response matrix}

Compute susceptibilities for many (component, probe) pairs to get the \textbf{response matrix}:
\[
X_{ij} = \text{susceptibility of component } C_j \text{ to probe } q_i.
\]
\textbf{What patterns reveal:}
\begin{itemize}[nosep]
    \item \textbf{Similar columns} $\Rightarrow$ components with similar functions (same ``circuit'')
    \item \textbf{Similar rows} $\Rightarrow$ probes involving similar patterns
    \item \textbf{PCA} reveals latent modes organizing both data and components
\end{itemize}

In a 2-layer transformer with 16 heads, PCA revealed:
\begin{center}
\begin{tabular}{@{}clp{7cm}@{}}
\toprule
\textbf{PC} & \textbf{Var.} & \textbf{What it captures} \\
\midrule
PC1 & 89.6\% & Universal patterns (word boundaries)---all heads load similarly \\
PC2 & 3.8\% & Separates \textbf{multigram heads} from \textbf{induction circuit} \\
PC3 & 1.3\% & Separates two induction heads: code vs. natural language \\
\bottomrule
\end{tabular}
\end{center}

\begin{keyinsightbox}
PCA on susceptibilities \emph{automatically rediscovered} known circuits without being told they exist---purely from response patterns.
\end{keyinsightbox}

%==============================================================================
\section{Summary}
%==============================================================================

\subsection{The method in one paragraph}

To understand what a neural network component does, we ask: ``If the training data shifted toward some probe (e.g., more code), how would this component's contribution change?'' This is the susceptibility---mathematically, the derivative of an expected observable with respect to a data perturbation. A key lemma shows this derivative equals a covariance that can be estimated from samples drawn from the posterior induced by the \emph{original} data. The resulting susceptibilities reveal whether components help certain predictions more or less than average. By computing susceptibilities across many (component, probe) pairs and analyzing the response matrix, we automatically discover functional structure.

\subsection{The five key ideas}

\begin{enumerate}
    \item \textbf{Bayesian framing:} Think of a trained model as inducing a distribution over weights, not just a single point. This lets us reason about ``what if the data were different?''

    \item \textbf{The causal chain:} Changing data changes the loss landscape, which changes which weights are optimal, which changes observables.

    \item \textbf{The key lemma:} The effect of a data perturbation equals a covariance that can be estimated from samples drawn from the posterior induced by the \emph{original} data. This makes the method practical.

    \item \textbf{Expression vs. suppression:} The sign of a component's susceptibility to a datapoint reveals whether that component helps that prediction more or less than average.

    \item \textbf{Structure discovery:} Patterns in the response matrix (found via PCA or clustering) automatically reveal functional organization---circuits, specialization, competition.
\end{enumerate}

\subsection{The one-liner}

\begin{center}
\fbox{\parbox{0.9\textwidth}{\centering
\textbf{The susceptibility of a component to a datapoint measures how much that component \\disproportionately helps (or doesn't help) that prediction.}
}}
\end{center}

%==============================================================================
\section{Where to go next}
%==============================================================================

This document covered the core ideas. Topics for further study:

\begin{itemize}[nosep]
    \item \textbf{Connection to generalization:} Susceptibilities relate to the ``local learning coefficient'' from singular learning theory, which directly connects to Bayesian generalization error.

    \item \textbf{Modes:} The data distribution can be decomposed into ``modes'' (via SVD), and susceptibilities measure the model's response to these modes.

    \item \textbf{Scaling:} Promising results up to Pythia-1.4B \cite{gordon2026spectroscopy}.

    \item \textbf{Connection to singular learning theory:} See \cite{watanabe2009algebraic} for the mathematical foundations.
\end{itemize}

\textbf{Further reading:} \cite{baker2025susceptibilities, wang2025embryology, gordon2026spectroscopy}.

\bibliographystyle{plainnat}
\bibliography{references}

%==============================================================================
\newpage
\appendix
\section{Proof of the covariance formula}
\label{sec:proof}
%==============================================================================

We prove that $\chi = -\mathrm{Cov}[\phi, \Delta L]$.

\paragraph{Setup.} The expected value of $\phi$ under the perturbed posterior is
\[
\langle \phi \rangle_h = \frac{N_h}{Z_h},
\]
where we define:
\begin{align*}
N_h &= \int \phi(w) \, e^{-n\beta L^h(w)} \varphi(w) \, dw, \\[4pt]
Z_h &= \int e^{-n\beta L^h(w)} \varphi(w) \, dw.
\end{align*}

\paragraph{Differentiating.} We want to compute $\frac{\partial}{\partial h}\langle \phi \rangle_h$. By the quotient rule:
\[
\frac{\partial}{\partial h}\langle \phi \rangle_h = \frac{1}{Z_h}\frac{\partial N_h}{\partial h} - \frac{N_h}{Z_h^2}\frac{\partial Z_h}{\partial h}.
\]

To compute $\frac{\partial N_h}{\partial h}$, we interchange the derivative and integral (see remark below):
\begin{align*}
\frac{\partial N_h}{\partial h}
&= \frac{\partial}{\partial h} \int \phi(w) \, e^{-n\beta L^h(w)} \varphi(w) \, dw \\[6pt]
&= \int \phi(w) \cdot \frac{\partial}{\partial h}\left[ e^{-n\beta L^h(w)} \right] \varphi(w) \, dw \\[6pt]
&= \int \phi(w) \cdot \left(-n\beta \frac{\partial L^h}{\partial h}\right) e^{-n\beta L^h(w)} \varphi(w) \, dw \\[6pt]
&= -n\beta \int \phi(w) \frac{\partial L^h}{\partial h} \, e^{-n\beta L^h(w)} \varphi(w) \, dw.
\end{align*}

Similarly for $Z_h$:
\begin{align*}
\frac{\partial Z_h}{\partial h}
&= \int \left(-n\beta \frac{\partial L^h}{\partial h}\right) e^{-n\beta L^h(w)} \varphi(w) \, dw \\[6pt]
&= -n\beta \int \frac{\partial L^h}{\partial h} \, e^{-n\beta L^h(w)} \varphi(w) \, dw.
\end{align*}

\paragraph{Combining terms.} Substituting into the quotient rule:
\begin{align*}
\frac{\partial}{\partial h}\langle \phi \rangle_h
&= \frac{1}{Z_h} \left( -n\beta \int \phi \frac{\partial L^h}{\partial h} \, e^{-n\beta L^h} \varphi \, dw \right) - \frac{N_h}{Z_h^2} \left( -n\beta \int \frac{\partial L^h}{\partial h} \, e^{-n\beta L^h} \varphi \, dw \right) \\[8pt]
&= -n\beta \cdot \frac{1}{Z_h} \int \phi \frac{\partial L^h}{\partial h} \, e^{-n\beta L^h} \varphi \, dw + n\beta \cdot \frac{N_h}{Z_h} \cdot \frac{1}{Z_h} \int \frac{\partial L^h}{\partial h} \, e^{-n\beta L^h} \varphi \, dw \\[8pt]
&= -n\beta \left\langle \phi \cdot \frac{\partial L^h}{\partial h} \right\rangle_h + n\beta \langle \phi \rangle_h \left\langle \frac{\partial L^h}{\partial h} \right\rangle_h \\[8pt]
&= -n\beta \left( \left\langle \phi \cdot \frac{\partial L^h}{\partial h} \right\rangle_h - \langle \phi \rangle_h \left\langle \frac{\partial L^h}{\partial h} \right\rangle_h \right) \\[8pt]
&= -n\beta \cdot \mathrm{Cov}_h\left[\phi, \frac{\partial L^h}{\partial h}\right].
\end{align*}

\paragraph{Evaluating at $h=0$.} Since $L^h(w) = (1-h)L^q(w) + h L^{q'}(w)$, we have:
\begin{align*}
\frac{\partial L^h}{\partial h} &= \frac{\partial}{\partial h} \left[ (1-h)L^q(w) + h L^{q'}(w) \right] \\[4pt]
&= -L^q(w) + L^{q'}(w) \\[4pt]
&= L^{q'}(w) - L^q(w) \\[4pt]
&= \Delta L(w).
\end{align*}

Therefore, evaluating at $h=0$:
\[
\frac{\partial}{\partial h}\langle \phi \rangle_h \bigg|_{h=0} = -n\beta \cdot \mathrm{Cov}_{h=0}[\phi, \Delta L].
\]

Since $\chi = \frac{1}{n\beta} \frac{\partial}{\partial h}\langle \phi \rangle_h \big|_{h=0}$ by definition, we conclude:
\[
\chi = -\mathrm{Cov}[\phi, \Delta L],
\]
where the covariance is under the unperturbed posterior ($h=0$). \hfill $\square$

\paragraph{Remark on regularity.} We freely interchanged derivative and integral above. This is justified by the Leibniz integral rule under mild smoothness and integrability conditions, which we do not verify here.

%==============================================================================
\section{How to sample from the posterior: SGLD}
%==============================================================================

The susceptibility formula $\chi = -\mathrm{Cov}[\phi, \Delta L]$ requires samples from the posterior. But neural network posteriors are high-dimensional and intractable. How do we actually get samples?

\subsection{The problem with the global posterior}

The full Bayesian posterior $p(w|D) \propto e^{-nL_n(w)}\varphi(w)$ integrates over \emph{all} weight configurations. For a neural network with millions of parameters, this is computationally infeasible. We can't enumerate all possible weights.

Moreover, we usually care about a \emph{specific trained model} $w^*$, not the global posterior. We want to understand what this particular checkpoint does.

\subsection{The local posterior}

The solution is to \textbf{localize} the posterior around $w^*$. Replace the prior $\varphi(w)$ with a Gaussian centered at the trained weights:
\[
p_{\text{local}}(w) \propto \exp\left\{ -n\beta L_n(w) - \frac{\gamma}{2}\|w - w^*\|^2 \right\}.
\]
The term $\frac{\gamma}{2}\|w - w^*\|^2$ acts like a ``leash'' that keeps samples close to $w^*$. The parameter $\gamma$ controls the leash strength:
\begin{itemize}[nosep]
    \item Large $\gamma$: Samples stay very close to $w^*$
    \item Small $\gamma$: Samples can wander further
\end{itemize}

This local posterior asks: ``In the neighborhood of $w^*$, which weight configurations are probable?''

\subsection{Stochastic gradient Langevin dynamics (SGLD)}

SGLD \cite{welling2011sgld} is a method to sample from the local posterior. The idea: gradient descent finds modes (minima), but if we add noise, we can \emph{sample} from the distribution instead.

\textbf{The SGLD update rule:}
\[
w_{t+1} = w_t - \varepsilon \nabla_w \left[ n\beta L_n(w_t) + \frac{\gamma}{2}\|w_t - w^*\|^2 \right] + \sqrt{2\varepsilon} \cdot \eta_t,
\]
where:
\begin{itemize}[nosep]
    \item $\varepsilon$ is the step size (learning rate)
    \item $\eta_t \sim \mathcal{N}(0, I)$ is Gaussian noise
    \item The gradient term pulls toward high-probability regions
    \item The noise term enables exploration (prevents collapse to a single point)
\end{itemize}

\begin{intuitionbox}
\textbf{SGLD = SGD + noise.} Without noise, gradient descent converges to a point. With calibrated noise, we get samples from the posterior instead.
\end{intuitionbox}

\subsection{In practice}

Typical hyperparameters from \cite{baker2025susceptibilities}:
\begin{itemize}[nosep]
    \item Localization strength: $\gamma = 300$
    \item Inverse temperature: $n\beta = 30$
    \item Step size: $\varepsilon = 0.001$
    \item Number of samples: 100--200 draws per chain
    \item Multiple chains (e.g., 4) for better coverage
\end{itemize}

For component-specific susceptibilities, SGLD is run only on the component's weights, with all other weights frozen at $w^*$.

\subsection{Why this works}

Under mild conditions, as $\varepsilon \to 0$ and the number of steps $\to \infty$, SGLD samples converge to the true posterior distribution. In practice, we use finite $\varepsilon$ and finite steps, so we get \emph{approximate} samples. The quality of the approximation depends on hyperparameter tuning.

The key insight: we don't need perfect samples. We just need samples good enough that the covariance estimate $\hat\chi = -\widehat{\mathrm{Cov}}[\phi, \Delta L]$ is reliable. Empirically, SGLD with the above hyperparameters works well for small transformers.

\end{document}
