Step 14: tool results

read_arxiv_source

{
  "file": "main.tex",
  "offset": 60000,
  "arxiv_id": "2501.17148",
  "max_chars": 18000
}

{"ok":true,"arxiv_id":"2501.17148","file":"main.tex","kind":"text","text":"itive or negative example. The detection and steering setup is then identical to DiffMean.\n% where we also experiment with L1 or L2 regularizations on the learned weights.\n% Excluding regularization loss on TopK activations allow selected latents to be highly activated when concepts is presented in the inputs. To summary, the objective contains a binary cross-entropy loss with regularization terms to ensure effective and sparse activations.\n% To evaluate, we then get the token-level concept score predictions as:\n% \\[\n% \\Psi_{\\text{Detect}}^{\\text{Linear}}(\\mathbf{h}) = \\text{ReLU}(\\mathbf{h} \\times \\mathbf{w}_{\\text{Probe}})\n% \\]\n% We then use max-pooling as described in \\cref{sec:concept-detection} to get sequence-level predictions. For evaluating model steering with learned $\\mathbf{w}_{\\mathrm{Probe}}$, we follow the same setup as DiffMean.\n\n\n\\paragraph{\\concepttag{}\\steeringtag{} Supervised steering vector (SSV).} The supervised steering vector method directly learns an intervention that maximises the language-modelling probability of the positive responses. For a sequence of token representations \\(\\mathbf{h}\\), we apply an intervention to each token representation:\n\\begin{equation}\n\\Phi^{\\text{SSV}}(h_i) = h_i + \\mathbf{w}_{\\text{SSV}}\n\\end{equation}\nwhere \\(\\mathbf{w}_{\\text{SSV}} \\in \\mathbb{R}^{d \\times 1}\\) is a learned vector. As described in \\cref{sec:model-steering}, we backpropagate gradients by training with the language modeling loss, similar to supervised fine-tuning (SFT):\n\\begin{equation}\n\\min_{\\mathbf{w}_{\\text{SSV}}}\\left\\{\\sum_{t=1}^n \\log P_{\\mathrm{LM}}\\bigl(y_t \\mid y_{<t}, \\mathbf{x}; \\mathbf{h} \\leftarrow \\Phi^{\\text{SSV}}(\\mathbf{h})\\bigr)\\right\\}\n\\end{equation}\nwhere \\(y_i\\) is the \\(i\\)-th output token, \\(y_{<i}\\) are the preceding tokens, and \\(\\mathbf{x}\\) is the prompt. For evaluating concept detection and model steering SSV follows the same setup as DiffMean. We apply $\\mathrm{ReLU}$ to get the detection scores.\n\n\\paragraph{\\concepttag{}\\steeringtag{} Rank-1 representation finetuning (ReFT-r1).} We introduce a novel method based on ReFT~\\cite{wu2024reft} which jointly learns concept detection and steering on supervised data by combining the training objectives of linear probing and supervised steering.\n\nWe compute latents for concept detection as:\n\\begin{equation}\n\\Psi_{\\text{Detect}}^{\\text{ReFT-r1}}(h_i) = \\mathrm{ReLU}(h_i \\cdot \\mathbf{w}_{\\text{ReFT-r1}})\n\\end{equation}\nDuring training we perform a representation-level intervention on each $h_i$ based on the latents of the sequence $\\mathbf{h}$:\n\\begin{equation}\n\\Phi^{\\text{ReFT-r1}}(h_i) = h_i + \\left(\\frac{1}{k}{\\left\\lVert \\topk(\\Psi_{\\text{Detect}}^{\\text{ReFT-r1}}(\\mathbf{h})) \\right\\rVert_1}\\right) \\mathbf{w}_{\\text{ReFT-r1}}\n\\end{equation}\n% \\begin{equation}\n% \\Phi^{\\text{ReFT-r1}}(h_i) = h + \\frac{1}{|\\mathbf{h}|}\\sum_{h_i \\in \\mathbf{h}}{\\topk_{h_i \\in \\mathbf{h}}(\\Psi_{\\text{Detect}}^{\\text{ReFT-r1}}(h_i))}\n% \\end{equation}\nwhere \\(\\mathbf{w}_{\\text{ReFT-r1}} \\in \\mathbb{R}^{d \\times 1}\\) is a learned vector. Finally, the training objective combines language modelling loss subject to this intervention, along with L1 regularisation on the non-top-$k$ latents:\n\\begin{equation}\n\\min_{\\mathbf{w}_{\\text{ReFT-r1}}}\\left\\{\n-\\sum_{t=1}^n{\\log P_{\\mathrm{LM}}^{\\Phi^{\\text{ReFT-r1}}}(y_t \\mid y_{<t}, \\mathbf{x})}\n+ \\lambda\\sum_{\\mathclap{a_i \\notin \\topk\\left(\\Psi(\\mathbf{h})\\right)}}{\\lVert a_i \\rVert_1}\n\\right\\}\n\\end{equation}\nDetection and steering is identical to DiffMean.\n\n% For a sequence of token representations \\(\\mathbf{h}\\), ReFT-r1 modifies the activation function of the linear probe by getting the mean of top-\\(k\\) activations over tokens in two steps:\n% \\begin{equation} \\label{reftr1-eq1}\n% \\begin{split}\n% % \\Phi_{\\text{Linear}}(\\mathbf{h}_i)\n% \\textbf{a} & = \\text{ReLU}(\\mathbf{h} \\times \\mathbf{w}_{\\text{ReFT-r1}}) \\\\\n% \\overline{a} & = \\textsc{Mean}\\bigl(\\textsc{TopK}(\\textbf{a})\\bigr)\n% \\end{split}\n% \\end{equation}\n% where $\\textbf{a}$ is the raw latent activations and $\\overline{a}$ is the mean of top-\\(k\\) activations.\n% Then, ReFT-r1 uses $\\mathbf{w}_{\\text{ReFT-r1}}$ as the steering vector:\n% \\begin{equation} \\label{reftr1-eq2}\n% \\begin{split}\n% \\Phi_{\\text{ReFT-r1}}(\\mathbf{h}) & = \\mathbf{h} + \\overline{a} \\cdot \\mathbf{w}_{\\text{ReFT-r1}}\n% \\end{split}\n% \\end{equation}\n% where $\\overline{a}$ is used as the steering factor during training. This ensures that only the most salient features in the sequence contribute to the steering effect. The training objective combines language modeling and sparsity:\n\n% \\begin{align*}\n% \\mathcal{L} = & - \\sum_{i=1}^n \\log P_{\\text{LM}}\\bigl(\n%     y_t \\mid y_{<t}, \\mathbf{x}; \\mathbf{h} \\leftarrow \n%     \\Phi_{\\text{ReFT-r1}}(\\mathbf{h})\\bigr) \\\\\n%     & + \\lambda_{\\text{ReFT-r1}} \n%     \\sum_{\\mathclap{a_i \\notin \\text{TopK}(\\mathbf{a})}} \\|a_i\\|_1\n% \\end{align*}\n% Similar to SSV, ReFT-r1 push the activation valyes down for non-TopK activations. We also experiment with L1 or L2 regularizations on the learned weights. For evaluating concept detection and model steering with $\\mathbf{w}_{\\text{ReFT-r1}}$, we follow the same setup as for sparse linear probe.\n\n\\paragraph{\\concepttag{}\\steeringtag{} Sparse autoencoders (SAE).} Sparse autoencoders are a self-supervised dictionary learning method (see \\cref{sec:related}). We use pretrained SAEs from \\texttt{GemmaScope}, which are the best available SAEs for \\texttt{Gemma}-family LLMs~\\cite{lieberum-etal-2024-gemma}.\\footnote{\\texttt{GemmaScope} releases a set of SAEs for \\texttt{Gemma-2-27B}, but the concept list is not publicly released, which makes the SAEs for \\texttt{Gemma-2-9B} the largest ones available for evaluations.}\n% It uses JumpReLU SAEs~\\cite{rajamanoharan2024jumping} and often shows a lower reconstruction error than other SAEs such as TopK SAEs~\\cite{gao2024scalingevaluatingsparseautoencoders}, Gated SAEs~\\cite{rajamanoharan2024improving} or vanilla ReLU SAE~\\cite{huben2024sparse}.\nThe SAEs we used are trained to learn two dictionary matrices, $\\{\\mathbf{W}_{\\mathrm{enc}}, \\mathbf{W}_{\\mathrm{dec}}\\} \\in \\R^{d\\times z}$ where $z$ is the number of latents. For our evaluating concept $\\mathbf{c}$, we use $\\{\\mathbf{w}_{\\mathrm{enc}}, \\mathbf{w}_{\\mathrm{dec}}\\} \\in \\R^{d \\times 1}$ as the detection and steering representations, respectively:\n\\[\n\\Psi_{\\text{Detect}}^{\\text{SAE}}(h_i) = \\sigma\\left(h_i \\cdot \\mathbf{w}_{ \n \\textrm{enc}} + b_{\\textrm{enc}} \\right)\n\\]\nwhere $\\sigma$ is an activation function (in our case, JumpReLU) and $b_{\\text{enc}}$ is a learned bias.\\footnote{Note that this parameterisation cannot apply to TopK \\cite{gao2024scalingevaluatingsparseautoencoders} and BatchTopK SAEs, which require loading in the entire encoder matrix to compute latents.}\nFor steering, we use activation addition as DiffMean.\n% $$\n% \\Phi_{\\text{Steer}}^{\\text{SAE}}(h_i) = h_i + \\alpha\\mathbf{w}_{\\textrm{dec}}\n% $$\nNote that \\citet{templeton2024scaling} use activation clamping; we report ablations in \\cref{sec:abaltions}.\n\n\\paragraph{\\concepttag{}\\steeringtag{} SAEs with AUROC selection (SAE-A).} Given that other methods have access to a training dataset, to enable fair comparison we attempt to use our training dataset for SAE feature selection. For each feature, we compute its max-pooled activations per \\cref{eq:maxpool-detect} over each training example, compute AUROC over the dataset given true labels, and select the highest-scoring feature by this metric.\n\n% Note that for concept detection, we rely on learned $\\textbf{w}_{\\text{enc}}$. For model steering, we adapt the existing approach as in the Golden Gate Bridge \\texttt{Claude}~\\cite{templeton2024scaling} example:\n\n% \\begin{align}\\label{eqn:sae-steer}\n% \\mathbf{a}^{\\text{clamp}} &= \\sigma\\bigl(\n%     \\mathbf{h} \\times \\textbf{W}_{\\text{enc}} \n%     + \\textbf{b}_{\\text{enc}} \\bigr)\\,\\odot\\, \n%     (\\mathbf{1} - \\mathbf{e}_{\\mathbf{c}}) \\notag \\\\\n% &\\quad + \\alpha \\cdot \\mathbf{e}_{\\mathbf{c}} \\\\\n% \\Phi_{\\text{Steer}}^{\\text{SAE}}(\\mathbf{h}) &= \\underbrace{\\mathbf{h} - \\hat{\\mathbf{h}}}_{\\text{reconst.}} \n%     +\\,\\mathbf{a}^{\\text{clamp}} \\cdot \\mathbf{W}_{\\text{dec}}\n% \\end{align}\n\n% where $\\mathbf{e}_{\\mathbf{c}}$ is an one-hot vector representing the corresponding concept, and $\\mathbf{a}^{\\text{clamp}}$ modifies the latent activation by clamping the activation corresponding to the concept to a higher value to amplify the contribution of the specific concept in the output representation. In this paper, we focus on a newer variant of SAEs with the JumpReLU activation function~\\cite{rajamanoharan2024jumping}. Details about JumpReLU and how it is trained can be found in its original paper.\n\n% \\paragraph{\\steeringtag\\ Low-rank adaptation (LoRA).} LoRA~\\cite{hu2022lora} is a parameter-efficient approach that injects rank-$r$ trainable matrices into select weight layers, keeping the original weights~\\(\\theta_{\\text{LM}}\\) frozen.  \n% We then train these low-rank parameters on the supervised steering objective (similar to \\(\\Phi_{\\text{SSV}}\\)) to produce concept-steered outputs. As a result, the negative training set is discarded.\n% We use LoRA as one of our baselines for model steering.\n\n\n% \\paragraph{\\steeringtag\\ Low-rank representation finetuning (LoReFT).} LoReFT~\\cite{wu2024reft} is a variant of ReFT that targets orthogonal linear subspaces through learning interventions composed of rank$r$ trainable matrices. Unlike LoRA which learns a weight transformation, LoReFT learns representation edits to adapt the LM to the training task. Similarly to LoRA, we train LoReFT on the supervised steering objective. The negative training set is discarded.\n% We use LoReFT as one of our baselines for model steering.\n\n% \\paragraph{\\steeringtag\\ Full finetuning (FFT).} Full finetuning (FFT) updates \\emph{all} model parameters \\(\\theta_{\\text{LM}}\\) using the same supervised steering objective as LoRA or LoReFT. The negative training set is discarded.\n% We use FFT to establish a strong baseline for model steering, as it trades compute resource for performance.\n\n\\paragraph{\\concepttag{} Bag-of-Words (BoW).} For the BoW baseline, we first construct a featurizer that tokenizes text by whitespace and counts word frequencies. The vocabulary for this featurizer is derived from the training dataset. We then train a logistic regression classifier to predict class probabilities, framing the task as binary classification. To mitigate overfitting, we incorporate a regularization term. This BoW approach leverages statistical biases inherent in LLM-generated data.\n\n\\paragraph{\\concepttag{} Gradient-based baselines.} We test two gradient-based attribution methods, which are applicable only to concept detection: Input $\\times$ gradients (\\textbf{I$\\times$G}) and Integrated gradients (\\textbf{IG}; \\citealp{integratedgradient}). For both, we train a classification head on the hidden representations of some layer and apply the methods to produce token-level attribution scores $\\Psi_{\\text{Detect}}(h_i)$. Implementation details are in \\cref{sec:gradient-baselines}.\n\n\\paragraph{\\concepttag{}\\steeringtag{} Prompting baseline.} For concept detection, we use the same LLM judge as described in \\cref{sec:model-steering} to rate the presence of a concept on a scale of 0 to 2.\nFor model steering, we use an LLM to \\emph{engineer} a prompt \n% with few-shot in-context demonstrations \ngiven a concept, which we use to steer our local model by prepending it to the actual instruction. We provide prompt templates and examples in \\cref{sec:prompt_templates} and \\cref{sec:model_generations}.\n% In the current paper, prompting baseline for concept detection is not included.\n% \\TODO{\\textbf{The prompt baseline is not evaluated for concept detection.} This is not true anymore, we will add in prompt baseline, but prompt baseline will be have AUC since it is a per-sequence class label (0, 1). We should consider to report the classic \\texttt{Recall@N} at the end in the main text? }\n\n\\paragraph{\\steeringtag{} Finetuning baselines.} We test full-parameter supervised finetuning (\\textbf{SFT}) and two parameter-efficient finetuning methods: Low-rank adaptation (\\textbf{LoRA}; \\citealp{hu2022lora}) and low-rank representation finetuning (\\textbf{LoReFT}; \\citealp{wu2024reft}). In all cases, we finetune to minimise the language-modelling loss on the responses in the positive split of the dataset; the negative training split is discarded. We then use the finetuned models as baselines for steering.\n\nFor all of our SDLs except SSV, we constrain any learned subspace to have a unit norm, following the same setup as SAEs. With a unit-norm constraint, we find that SSV is hard to use for steering models. For prompting and finetuning baselines, we randomly score one generation on the testing instruction set (since the factor is not a parameter for those methods), resulting in the same number of observations for those methods.\n\n% \\paragraph{Supervised probes.} \n\n% \\paragraph{Unsupervised interventions.} \n\n% \\paragraph{Supervised interventions.} \n\n% \\paragraph{Sparse autoencoders (SAEs).} \n\n% \\paragraph{Rank-1 representation finetuning (ReFT).} \n\n% \\subsection{Datasets}\n\n% \\paragraph{Random seed sentences.} For \\textit{text}, we use the \\texttt{ag-news} dataset. For \\textit{code}, we use the \\texttt{Rosetta} code dataset. For \\textit{math}, we use the \\texttt{Omni-Math} datasets. For each dataset, we subsample 1,000 examples to create our final sentence pool. \n\n\\subsection{Evaluation}\\label{sec:evaluation}\n\n\\paragraph{Datasets.} We synthetically generate training and validation datasets (see \\cref{sec:concept-dataset}) for 500 concepts, which we release as \\textsc{Concept500}. The concepts are sampled from the Neuronpedia SAE concept list for \\texttt{GemmaScope} as described in \\cref{sec:sae-concept-list}. For each concept, we include 144 examples for training and $\\approx$72 samples for evaluating concept detection.\\footnote{This varies based on valid hard negatives.} In this paper, we train and evaluate all methods, and report results on \\textsc{Concept500}. For SFT, we only train and evaluate on the first 20 concepts due to limited resources.\n\nFor evaluating steering, we use the instructions from the \\texttt{Alpaca-Eval} dataset~\\cite{alpaca_eval}. For each concept, we sample 10 instructions. We generate up to 128 tokens for each instruction over 14 steering factors. We split the instructions into two equal sets -- one for selecting the best factor and the other for evaluation.\n\nWe additionally release training and evaluation datasets for all 16K concepts in \\texttt{GemmaScope} as the \\textsc{Concept16K} dataset suite. We train and release SAE-scale dictionaries on this dataset only for the best-performing methods found on \\textsc{Concept500}. See \\cref{sec:data_stats} for dataset statistics and \\cref{sec:sdl-16k} for further experiments on \\textsc{Concept16K}.\n\n\\paragraph{Models.} Our evaluations rely on access to and control over the LLM's representations. To reduce training cost, we prefer to use models for which pretrained SAEs are available. We thus evaluate our methods on two open models, \\texttt{Gemma-2-2B-it} and \\texttt{Gemma-2-9B-it} (henceforth referred to without the \\texttt{-it} suffix), from the \\texttt{Gemma}-family, with corresponding SAEs released as \\texttt{GemmaScope}. We evaluate our methods with model representations from the residual streams of layers 10 and 20 for \\texttt{Gemma-2-2B} and layers 20 and 31 for \\texttt{Gemma-2-9B}. We use SAEs from \\texttt{GemmaScope} that are trained for these layers.\\footnote{For \\texttt{Gemma-2-2B}, we follow the common practice to use SAEs for the base LM, as SAEs are not available for the instruction-tuned model at the time of publication~\\cite{lieberum-etal-2024-gemma}.} To ensure a fair comparison, we perform separate hyperparameter-tuning for each method. Details can be found in \\cref{sec:hyperparameters}.\n\n \n% \\paragraph{Hyperparameters.} To ensure a fair comparison, we perform separate hyperparameter-tuning for each method that requires training. To select the best hyperparameters, we use a held-out dataset for development, \\textsc{Concept10}, which  covers 10 concepts each on our tasks. Details can be found in \\cref{sec:hyperparameters}. We minimise the loss with \\texttt{AdamW} with a linear scheduler for all methods that require training. \n\n% \\paragraph{Steering factors.} For any method that learns a rank-1 subspace, the subspace can be treated as a \\emph{concept} direction. A common practice is to steer a model in a direction of different magnitudes, varying the steering effect~\\cite{subramani-etal-2022-extracting}. In total, we evaluate with 14 steering factors (listed in \\cref{sec:hyperparameters}). For each steering factor, we calculate the final steering magnitude by multiplying the factor with maximal latent activation by projecting hidden representations to this rank-1 subspace for all training examples we created for sueprvised training. This is identical to getting the unnormalized token-level concept score, such as in \\cref{diffmean-eq1}.\n% For SAEs, we fetch the maximal latent activations from publicly released data.\\footnote{We use \\url{https://www.neuronpedia.org/} to fetch the maximum latent activations.} For SAEs, we also consider the reconstruction error term when adding the steering vector as in \\cref{eqn:sae-steer}. \n% To ensure statistical equivalence, we treat the sampling process for non-rank-1 dictionary learning methods (e.g., LoRA or LoReFT) as if they had steering factors, giving them an equal chance of being favoured regardless of whether a steering factor is needed. For example, we would sample $14\\times 5 = 70$ times to get the averaged scores. And for win-rate, we would sample $70$ times on a separate set for each concept to get the best generation. This is crucial as best-of-n is a common strategy to get the best generation~\\cite{stiennon2020learning}. This ensures that we treat e","chars":18000,"total_chars":134717,"offset":60000,"truncated":true,"next_offset":78000}