Step 14: tool results

read_arxiv_source

{
  "file": "main.tex",
  "offset": 20000,
  "arxiv_id": "2502.16681",
  "max_chars": 20000
}

{"ok":true,"arxiv_id":"2502.16681","file":"main.tex","kind":"text","text":"l}$, we choose $k \\ll W$ and find indices $\\mathcal{I}$ as follows:\n% \\begin{align}\\label{eqn:select_k_indices}\n% \\mathcal{I} = \\ArgTopK_{i \\in \\{W\\}}\\left | \\frac{1}{n_1} \\sum_{\\substack{j \\in \\{n_1\\}}} SAE(X^l_{-1})[j, i]  - \\frac{1}{n_0} \\sum_{\\substack{j \\in \\{n_0\\}}} SAE(X^l_{-1})[j, i]  \\right |\n% \\end{align}\n\\begin{align}\\label{eqn:select_k_indices}\n\\mathcal{I} = \\ArgTopK_{i \\in \\{W\\}} \\left| \\frac{1}{|T_1|} \\sum_{j \\in T_1} Z_{j,i} - \\frac{1}{|T_0|} \\sum_{j \\in T_0} Z_{j,i} \\right|\n\\end{align}\n\nWe then train a probe $p_\\textrm{SAE}$ to map from $Z[:, \\mathcal{I}] \\mapsto t$.\n\nNote that previous work in SAE probing has aggregated SAE latents across tokens to provide a richer input space for SAE probes \\cite{neuronpediaSAEBenchComprehensive, features_as_classifiers}. However, most probing studies operate on the final token, which we choose to emulate. In \\cref{sec:why-didnt-this-work}, we present results considering probes on multiple tokens.\n% , more closely emulating an ML scientist's realistic probing workflow, as well as to eliminate the confounding variable of aggregation technique (SAE activations are naturally easier to aggregate across tokens, but stronger baselines for activations . \n\nThroughout this study, we use the area under the ROC curve (AUC) to evaluate the quality of a probe. AUC is the likelihood that the classifier assigns a higher probability of $\\mathrm{target}=1$ to a $\\mathrm{target}=1$ example than to a $\\mathrm{target}=0$ example. Using AUC allows us to comprehensively assess probe performance agnostic to classification thresholds. For additional details on AUC, see \\cref{sec:app_auc_explainer}.\n % providing a comprehensive evaluation of its discriminatory capability. \n \nOften, a probe $p$ has hyperparameters $h_{p}$ we would like to optimize. \n% To mimic real life probing situations, \nWe select $h_p$ that has the maximal validation AUC using the cross-validation strategy described in \\cref{tab:crossval}. We then test $p$ with optimal $h_p$ on a held out test set to calculate $\\mathrm{AUC}_p^{\\mathrm{test}}$. All datasets have at least $100$ testing examples, with most having more (the average test set size is $1945$). \n\n\n\\subsection{Probing Methods}\n\\label{sec:probing_methods}\nWe use the following 5 probing methods, with hyperparameters in parenthesis:\n\\begin{enumerate}\n[topsep=1pt,itemsep=0.2pt,parsep=0pt,leftmargin=10pt]\n\n    \\item Logistic Regression ($\\mathrm{L}_1$ regularization for SAE probes, $\\mathrm{L}_2$ otherwise).\n    \\item PCA Regression (\\# of PCA components to reduce $X_{-1}^l$ to before using unregularized logistic regression).\n    \\item K-Nearest Neighbors (KNN) (\\# of nearest neighbors).\n    \\item XGBoost \\cite{xgboost} ($\\mathrm{n\\_estimators}$, $\\mathrm{max\\_depth}$, $\\mathrm{learning\\_rate}$, $\\mathrm{subsample}$, $\\mathrm{colsample\\_bytree}$, $\\mathrm{reg\\_alpha}$, $\\mathrm{reg\\_lambda}$, and $\\mathrm{min\\_child\\_weight}$).\n    \\item Multilayer Perceptron (MLP) (network depth, hidden state width, learning rate, and weight decay).\n\\end{enumerate}\n\nWe evaluate 10 hyperparameter values $h_p$ for each probing method $p$. For the first three probing methods, we use grid search, while for MLPs and XGBoost, we randomly sample from a hyperparameter grid to manage the larger search space. For additional details on hyperparameter ranges, see \\cref{sec:app_probe_hyperparam}. We use all five probing methods as baselines, but for SAE probes we only use logistic regression.\n\n\\subsection{Experimental Setup: Quiver of Arrows}\n\\label{sec:quiver}\nTo evaluate whether SAE probes provide an advantage over baselines, we use an evaluation metric we call the ``Quiver of Arrows'': we ask whether adding SAE probes (a new ``arrow'') to the set of existing probing methods available to a practitioner (the ``quiver'') increases performance compared to a practitioner \\textit{without} access to SAE probes. In other words, over a collection of methods we choose the best method by validation AUC and then report the test AUC of that method. We can then compare the marginal improvement of adding SAE probes to a practitioner's toolkit by comparing the Quiver of Arrows AUC of baeline methods with and without SAE probes. We describe the Quiver of Arrows more formally in \\cref{app:quiver}.\n\nThe quiver of arrows approach is designed as a robustness check to fairly evaluate the benefit of adding SAE probes to a practioner's toolkit. If we instead constructed probes from a large set of SAEs and chose the best one based on test AUC, we could be tricked into thinking SAEs outperform baselines because we accessed the held-out test set. With the quiver of arrows approach, we emulate the information a real practioner has access to, allowing us to make a robust recommendation on the usefulness of SAE probes.\n\n\n\\begin{figure}\n    \\centering\n    \\includegraphics[width=\\linewidth]{figures/Normal/sae_l0_k_normal_pareto_16k.pdf}\n    \\vspace{-0.6cm}\n    \\caption{For a given width, using higher L0 and constructing probes with a larger basis of latents ($k$) is more performant.}\n    \\label{fig:sae_normal_l0_k_pareto}\n\\end{figure}\n\n\n\n% While e.g. comparing test accuracies might be simpler, it is unrealistic, as in the regimes we consider a practitioner does not have access to a large ground truth test set to choose the best method. \n% CV: Leave 2 out for < 12. k = 6 CV < 128. otherwise 20\\% val. \n\n\n% \\subsection{Experiment Setup: Noisy Labels}\n% In this setting, we assume that the practitioner suspects that there will be some sort of distribution shift, or that their labels are noisy, and so cannot trust the validation data to select between methods. Thus, we just show the difference between a \\textit{global} choice of the 16k width SAE with highest L0 versus logistic regression.\n\n\n\\section{Comparing Probing Techniques in Different Regimes}\n\n\\subsection{Standard Conditions}\nWe initially assess probes under standard conditions, characterized by sufficient data for probe convergence and balanced classes. We find that baseline methods perform the best on layer 20 (see \\cref{fig:app-baseline-layer}, \\cref{sec:app_baselines}), so we run experiments with this layer. We train baselines on 1024 data points (or the maximum number of points in the dataset). \n% This simulates where a practitioner without SAEs would probe, establishing a proper counterfactual for when we add SAEs to the quiver.\n\n% Since Gemma Scope contains 34 SAEs of $\\saesmall$ to $\\saelarge$ with multiple $\\lzero$ values, w\nWe first conduct a preliminary investigation to cut down the large space of Gemma Scope SAEs. We train probes for all SAEs using $k$ latents (for logspaced values of $k$) on all datasets using 1024 training examples. We calculate the average test AUC for each SAE probe across all datasets. We find that SAE $\\mathrm{width}$ is relatively unimportant to probe success, while larger $\\lzero$s lead to more performant probes (\\cref{sec:app_sae_normal}). Additionally, as we might expect, using a larger value of $k$ leads to better probe performance, as shown in \\cref{fig:sae_normal_l0_k_pareto}.\n\n\n\\begin{figure}\n    \\centering\n    \\includegraphics[width=\\linewidth]{figures/Normal/sae_vs_nosae_normal_comparison.pdf}\n    \\vspace{-0.6cm}\n    \\caption{In standard conditions, when SAE probes are added to the quiver, we find a slight decrease in performance.}\n    \\label{fig:sae_normal}\n\\end{figure}\n\n% In the appendix (\\cref{fig:all_datasets_vs_k}), we plot the test AUC vs. $k$ for every individual dataset.\n\n\nThus, throughout the paper we train probes using the largest $\\lzero$ for SAEs $\\saesmall, \\saemid$, and $\\saelarge$. We use $k = 16$ to construct easily interpretable probes that potentially overfit less and use $k = 128$ for performance. \n\n\nUsing this set of probes, we consider our quiver of arrows approach in standard conditions. SAEs are chosen as the ``arrow'' for 14/113 tasks, however, we see a slight \\textit{decrease} in performance when they are added to the quiver in \\cref{fig:sae_normal}.\n\n\nIt is unsurprising that SAE probes underperform baselines in standard conditions, as their inductive bias is marginal given a large, balanced dataset. Thus, we now investigate more difficult settings to test if the inductive bias of SAEs translates into a competitive advantage for probing.\n\n\n\n\\begin{figure*}[h]\n    \\centering\n    \\includegraphics[width=\\linewidth]{figures/combined/target_datasets_quiver_combined_new.pdf}\n    \\vspace{-0.5cm}\n    \\caption{For three of the datasets in \\cref{tab:example_data}, we visualize the performance when SAE probes are in the quiver (dashed) versus when they are not (solid) for the regimes of data scarcity (left), class imbalance (middle), and label noise (right). In all three regimes, we see that on average (bottom row), SAEs do not help.}\n    \\label{fig:combined-target-datasets}\n\\end{figure*}\n\n\\begin{figure*}[htbp]\n    \\centering\n    % First figure\n    \\includegraphics[width=\\linewidth]{figures/combined/quiver_improve_regimes.pdf}\n    \\vspace{-0.6cm}\n    \\caption{For the regimes of data scarcity, class imbalance, and label noise, we see no average improvement across datasets when adding SAEs to the quiver. Shading represents 95\\% confidence intervals.}\n    \\label{fig:combined-improvement}\n\\end{figure*}\n\n\n\n\n\\subsection{Data Scarcity, Class Imbalance, and Label Noise}\n\n\nIn this section, we consider more difficult probing regimes: limiting the training data (Data Scarcity), changing the relative frequency of $\\mathrm{target}=1$ examples (Class Imbalance), and randomly flipping a fraction of the $\\mathrm{targets}$ (Label Noise). All of these settings are realistic workflows - for example, a researcher observing a rare model phenomena would like a probe that generalizes with few total examples or few examples of the desired class, while a researcher working with collected user data wants a generalizable probe even if some fraction of users respond randomly.\n\nEach regime is characterized by a parameter which we vary to compare SAEs and baselines.\n\\begin{itemize}\n[topsep=0.3pt,itemsep=0.2pt,parsep=0pt,leftmargin=10pt]\n    \\item Data Scarcity - We use 20 values of $n$ logspaced in $[2,1024]$, where $n$ is the number of training examples. We use a standard test set.\n    \\item Class Imbalance - We use 19 values of $\\mathrm{ratio}$ linearly spaced between $[0.05, 0.95]$, where $\\mathrm{ratio} = \\frac{n_1}{n}$ and $n_1$ is the number of $\\mathrm{target}=1$ training examples. The test set uses the same $\\mathrm{ratio}$.\n    \\item Label Noise - We use 11 values of $\\mathrm{fraction}$ linearly spaced between $[0,0.5]$, where $\\mathrm{fraction} = \\frac{n_{\\mathrm{corrupted}}}{n}$. We use a standard test set (uncorrupted).\n\\end{itemize}\n\nWe evaluate the first two regimes with the quiver of arrows method. However, the quiver of arrows method relies on the validation data being representative of the test data. This assumption fails for the label noise setting since the validation data is corrupted. Thus, a hypothetical practitioner would likely deploy their most performant method from other settings instead of allowing corrupted validation AUC to arbitrarily choose a method. To emulate this, we compare logistic regression and the $\\saesmall$, $k = 128$ SAE probe head-to-head in the setting of label noise. \n% These methods represent the most performant baseline and SAE probe in standard settings (\\cref{sec:app_baselines}, \\ref{sec:app_sae_normal}).\n\nIn \\cref{fig:combined-target-datasets}, we visualize the SAE versus non-SAE quivers for a sample of datasets listed in \\cref{tab:example_data} across the three regimes. Additionally, we visualize the average performance difference between the SAE and non-SAE quivers for each regime across all datasets in \\cref{fig:combined-improvement}. At all parameter values in each regime, SAEs show no meaningful improvement over baselines. This is not because SAEs are not chosen from the quiver; in \\cref{sec:app-regimes} we see that SAEs are chosen for up to 40 datasets in each regime. SAEs simply underperform baselines in these settings. See \\cref{sec:app-regimes} for results for individual datasets. \n\n\n\n\\subsection{Covariate Shift}\n\nWe now investigate if SAE probes are more resilient to distribution shifts in prompts. To model this covariate shift, we create or use 8 out-of-distribution (OOD) datasets: two pre-existing GLUE-X datasets which are designed as ``extreme'' versions of tasks 87 and 90 to test grammaticality and logical entailment respectively; three datasets (tasks 66, 67, and 73) which we alter the language of; and three datasets (tasks 5, 6, and 7) where we use syntactical alterations to names or use cartoon characters instead of historical figures. We train probes on these datasets in standard settings and evaluate on 300 covariate shifted test examples.\n\nLike the label noise setting, our validation data is unfaithful to our test data. Thus, we compare logistic regression to a single SAE probe. We construct SAE probes from the $\\saeneuronpedia$ SAE, as latent descriptions are available for this SAE on Neuronpedia \\citep{neuronpedia}. Our results (\\cref{fig:ood}) show that baselines outperform SAE probes when generalizing to covariate shifted data.\n\n\n\n\\section{Interpretability}\nWhile we find that SAE probes are not helpful in traditional probing settings, one intrinsic advantage of SAE probes is that their input basis is interpretable. To leverage this, we use the technique of automatic interpretation, or autointerp, to create natural language descriptions of top latents for each dataset  (see \\citet{bills2023language}).\n% by providing a language model with each latent's maximally activating examples. \nWe investigate three applications of labeled latents:\n\\begin{enumerate}[topsep=1pt,itemsep=0pt,parsep=0pt,leftmargin=10pt]\n\\item \\textbf{Probe Interpretability}: In \\cref{sec:probe_interp}, we investigate why SAE probes fail to perform well by pruning latents that o1 \\cite{o1} ranks as spurious.\n\\item \\textbf{Latent Interpretability}: In \\cref{sec:latent_interp}, we generate autointerp explanations of each dataset's top latent to find spurious latents and latents that fit well to our probes in a way not explained by the autointerp label.\n\\item \\textbf{Detecting Dataset Quality Issues}: In \\cref{sec:dataset_interp}, we invesigate spurious dataset features and label errors using insights from the top latent descriptions and firing patterns.\n\\end{enumerate}\n\n\\begin{figure}\n    \\centering\n    \\includegraphics[width=\\linewidth]{figures/OOD/ood_task_comp.pdf}\n    \\vspace{-0.8cm}\n    \\caption{SAE probes often generalize worse than baseline logistic regression with covariate shift.}\n    \\label{fig:ood}\n    \\vspace{-0.3cm}\n\\end{figure}\n\n\n% By generating a description of probe latents, we can determine if a given latent represents a spurious correlation. We can use this information to construct probes with fewer spurious correlations, which we expect will generalize better, or we can easily identify datasets with potential errors. We consider each of these situations in turn.\n\n\n\n\\subsection{Probe Intepretability: Pruning and Latent Generalization}\n\\label{sec:probe_interp}\nWe first use autointerp to investigate why SAE probes fail to generalize to covariate shift. We have two initial hypotheses: 1) the latents SAE probes rely on leverage spurious correlations in the training data and 2) the latents are not spurious, but they are not robust to the distribution shift.\n\nFirst, we investigate hypothesis 1. Specifically, we attempt to improve SAE probes OOD performance by pruning latents deemed spurious by their autointerp description. We focus on three OOD tasks that SAE probes perform poorly on: 7\\_hist\\_fit\\_ispolitician, 66\\_living-room, and 90\\_glue\\_qnli (\\cref{fig:ood}). We use the $\\saeneuronpedia$ SAE.\n\nFor each task, we take the $k = 8$ top latents by mean difference and use Claude-3.5-Sonnet \\cite{sonnet} to generate latent descriptions with autointerp. We then use OpenAI's o1 model to rank the relevance of each latent's description to the task. We also prompt o1 to downrank spurious latents given the OOD transformation. An example of this procedure for the task 66\\_living-room, which identifies if the phrase ``living room'' is in an English sentence, is shown in \\cref{tab:latents-description}, \\cref{sec:app-pruning}. For this task, o1 ranked latent 12274, which identifies ``mentions of living rooms'' first, while ranking latent 51330, which identifies ``objects and materials related to scientific experiments,'' last. \n\nWe then construct a probe with the top $k$ latents using o1's relevance rank for $k \\in [1,8]$. If there are spurious latents in the $k = 8$ probe, we expect a probe with $k<8$ to have better OOD generalization. We visualize each task's performance by $k$ in \\cref{fig:ood-pruning}, \\cref{sec:app-pruning}. For two of the datasets, 66\\_living-room and 90\\_glue\\_qnli, we see that pruning works, with a 0.024 and 0.052 increase in AUC between the $k=8$ and $k=1$ probe for each dataset. However, for 7\\_hist\\_fig\\_ispolitician, OOD test AUC \\textit{increases} by 0.077 after pruning. This indicates that the task 7\\_hist\\_fig\\_ispolitician requires $k=8$ latents to represent.\n\n\nOur preliminary results show that pruning helps SAE probes generalize. However, the drop between in-distribution (ID) and OOD performance is much larger than the modest improvement from pruning. This indicates that spurious correlations are not the primary reason SAE probes fail to generalize. Thus, we consider the second hypothesis - that the underlying SAE latents do not generalize well OOD.\n\nOn the task 66\\_living-room, latent 122774 has an ID test AUC of 0.99 while only having an OOD test AUC of 0.64. Since the OOD transformation involves translating the prompts to French, we hypothesize that this latent is active on the phrase ``living room'' in English but not other languages. We use GPT-4o to translate ``living room'' to 15 languages, and in \\cref{fig:ood-living-room} we indeed observe that latent 122774 is more active on the English translation than all other languages, and it is not active on the French translation at all. Thus, latent 122774 does not generalize OOD, which explains why the SAE probe also failed to generalize. \n% This highlights a failing of robustness in the Gemma Scope SAEs.\n\n\\begin{figure}\n    \\centering\n    \\includegraphics[width=\\linewidth]{figures/OOD/living_room_latent_translations.pdf}\n    \\vspace{-0.8cm}\n    \\caption{Latent 122774 is most active on the English translation of ``living room,'' and does not fire on French.}\n    \\label{fig:ood-living-room}\n\\end{figure}\n\n\n\n\\subsection{Latent Interpretability}\n\\label{sec:latent_interp}\nWe next generate autointerp descriptions for the most determinative latent for each dataset, allowing us to identify interesting categories of latents. Specifically, we consider each of the top 128 latents for each dataset for the $\\saeneuronpedia$ SAE, and evaluate each latent's $k=1$ SAE probe in standard settings. We then generate an autointerp explanation with GPT-4o for the best latent. \n\n\\cref{tab:latents-description} (\\cref{sec:app-latents}) contains a breakdown of a selection of interesting top latents that we find. We see that some latents' descriptions fit their tasks, like latent 81210 for 5\\_hist\\_fig\\_ismale, which activates on ``references to female individuals.'' Another interesting category is incorrect autointerp labels (e.g. for 125\\_wold\\_country\\_Italy, latent 50817 has 0.989 AUC, implying (correctly) that it fires on Italian concepts; however, the description reads ``names of researchers and scientific authors''). However, some latents appear to be spurious and unrelated to the semantics of their dataset. For example, 22\\_headline\\_isobama, which targets headlines published during the Obama administration, is classified with $0.782$ AUC by latent 10555, which activates on ``strings of numbers, often with mathematical notations.'' \n% F","chars":20000,"total_chars":104886,"offset":20000,"truncated":true,"next_offset":40000}