pith. sign in

arxiv: 2604.23130 · v2 · pith:H3ZI6DWYnew · submitted 2026-04-25 · 💻 cs.CL · cs.AI

From Concept-Aligned Tokens to Vulnerable Features: Mechanistic Localization of Jailbreaks

Nilanjana Das , Mathew Dawit , Aman Chadha , Manas Gaur This is my paper

Pith reviewed 2026-07-04 15:41 UTC · model glm-5.2

classification 💻 cs.CL cs.AI
keywords sparse autoencodersjailbreakmechanistic interpretabilityfeature steeringLLM safetyGemma-2-2Badversarial robustness
0
0 comments X

The pith

Single token localizes jailbreak vulnerability in LLM internals

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

This paper tries to establish that jailbreak susceptibility in Gemma-2-2B can be traced to sparse, token-localized SAE feature subgroups concentrated in mid-to-late layers (16–25), and that a single prompt token aligned to a harmful concept is sufficient to identify these subgroups without broader cluster-level aggregation.

Core claim

The paper introduces a three-stage pipeline: (1) extract harmful concept descriptions from adversarial responses, (2) align them with concept-relevant prompt tokens via subspace similarity at layer 20, and (3) identify SAE feature subgroups through three grouping strategies. The central discovery is that single-token-driven grouping achieves harmfulness comparable to full cluster-based grouping, demonstrating that individual harmful prompt tokens are sufficient to localize vulnerability-relevant features. Mid-to-late layers (16–25) are consistently most vulnerable to steering across all three strategies.

What carries the argument

SAE feature subgroups identified via subspace similarity between harmful concept descriptions and prompt token activations, grouped by three methods: cluster-based, hierarchical-linkage, and single-token-driven, then amplified (steered) to measure harmfulness increase.

If this is right

  • Targeted feature-level interventions at layers 16–25 could complement or replace prompt-level defenses against jailbreaks.
  • Single-token-driven feature identification reduces computational cost of locating vulnerability-relevant features compared to cluster-based approaches.
  • Harm categories involving violence and incitement show disproportionate steerability, suggesting category-specific internal representations that may require category-specific safety interventions.
  • If vulnerability features are token-localized, safety training could potentially be refined to suppress specific feature subgroups rather than broad refusal directions.

Load-bearing premise

The paper assumes that amplifying SAE feature activations during inference (steering) and observing increased harmfulness constitutes evidence that these features are causally responsible for jailbreak vulnerability, but it does not verify that the same features are engaged during actual adversarial attacks like GCG suffixes.

What would settle it

If steering the identified single-token-driven feature subgroups at layers 16–25 does not increase harmfulness above baseline, or if natural jailbreak attacks engage entirely different feature subgroups than those identified by this pipeline, the claim that these features are vulnerability-relevant would fail.

Figures

Figures reproduced from arXiv: 2604.23130 by Aman Chadha, Manas Gaur, Mathew Dawit, Nilanjana Das.

Figure 1
Figure 1. Figure 1: Approach 1: This figure represents steering with cluster-based feature selection. Results demonstrate view at source ↗
Figure 2
Figure 2. Figure 2: Approach 3: This figure represents single-token driven technique of feature selection. Similar to the other view at source ↗
Figure 3
Figure 3. Figure 3: Figure demonstrating that the later layers are relatively more steerable than the early-mid layers on the view at source ↗
Figure 4
Figure 4. Figure 4: Approach 2: This figure represents steering with the hierarchical linkage-based feature selection. It shows view at source ↗
read the original abstract

Jailbreak attacks expose a persistent failure mode in safety-aligned LLMs: models can be pushed into harmful behavior, but the internal representations enabling this shift remain poorly localized. Recent mechanistic safety studies often explain such behavior through broad representational objects, including global refusal directions, activation steering vectors, and refusal-related SAE features. We instead ask whether jailbreak vulnerability can be traced to finer-grained, prompt-conditioned SAE feature subgroups. We introduce a token-driven mechanistic pipeline that decomposes the residual stream of Gemma-2-2B into Sparse Autoencoder (SAE) features and identifies feature subgroups associated with unsafe behavior. Using single-category unsafe examples from BeaverTails to reduce cross-category interference, we extract harmful concepts from adversarial responses and align them with concept-relevant prompt tokens through subspace similarity. We then apply three feature-grouping strategies: cluster-based, hierarchical-linkage, and single-token-driven, to identify SAE feature subgroups across all 26 layers. Finally, we amplify the top features in each subgroup and evaluate the resulting generations with a standardized harmfulness judge. Single-token-driven grouping achieves harmfulness comparable to full cluster-based grouping, showing that individual harmful prompt tokens are sufficient to localize vulnerability-relevant SAE feature subgroups without relying on broader cluster-level aggregation. These subgroups appear across early and mid-to-late layers, with stronger concentration in mid-to-late layers, where targeted steering exposes specific model vulnerabilities. Overall, our results suggest that jailbreak susceptibility can be traced to sparse, token-localized SAE feature subgroups, complementing prior accounts based on broad adversarial, refusal, or steering directions.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

4 major / 7 minor

Summary. The paper proposes a three-stage pipeline for localizing jailbreak-relevant SAE feature subgroups in Gemma-2-2B. Starting from adversarial prompt-response pairs in BeaverTails, the authors extract harmful concepts from adversarial responses using Grok, align these concepts to prompt tokens via subspace similarity at layer 20, and then apply three feature-grouping strategies (cluster-based, hierarchical-linkage, single-token-driven) to identify SAE feature subgroups across all 26 layers. Steering experiments show that mid-to-late layers (16–25) are consistently more steerable, and that single-token-driven grouping achieves comparable harmfulness to full cluster-based grouping. The central claim is that jailbreak vulnerability can be traced to sparse, token-localized SAE feature subgroups.

Significance. The paper tackles a timely question at the intersection of mechanistic interpretability and LLM safety. The use of pre-trained SAEs from Gemma Scope and the design of three distinct feature-grouping strategies with cross-layer analysis is a reasonable methodological contribution. The finding that single-token-driven grouping suffices for localization, if validated with proper controls, would be practically useful. However, the significance is substantially limited by the absence of control conditions and statistical rigor, which prevent the results from supporting the causal claims made.

major comments (4)
  1. §2 (Problem Formulation) and §3 (Baseline and Evaluation Protocol): The paper's central claim is that the identified SAE features are 'vulnerability-relevant' and that single-token grouping 'suffices to localize' them. However, no control condition is reported in which arbitrary SAE features, non-harmful-concept features, or randomly selected features from the same layers (16–25) are steered. Without this control, the observed harmfulness increase could simply reflect the fact that later-layer residual-stream perturbations are generally more disruptive to safety alignment, regardless of which features are amplified. This is load-bearing: if random-feature steering in layers 16–25 produces comparable harmfulness increases, the entire concept-alignment and token-grouping pipeline adds no specificity beyond 'steer later layers.' A control experiment with random or non-harmful features is a遮
  2. §3 (Baseline and Evaluation Protocol): Results are presented as raw counts of prompts with increased harmfulness (e.g., 'peak category-layer value of 10' in §4), with no error bars, confidence intervals, or significance tests. With 20 prompts per category, a count of 10/20 could easily arise by chance. The claim that single-token-driven grouping achieves 'comparable' harmfulness to cluster-based grouping (abstract, §4) is not supported by any statistical comparison. Without at least a bootstrap confidence interval or a paired test, the quantitative claims about relative method performance and layer concentration are not justified.
  3. §2 (Problem Formulation): The pipeline extracts harmful concepts from adversarial responses (z), aligns them to prompt tokens, finds SAE features for those tokens, steers those features, and measures increased harmfulness. This is partially circular: features are selected because they align with harmful concepts extracted from harmful responses, then their harmfulness is confirmed by steering. The paper does not verify that these same features are activated during actual adversarial attacks (e.g., GCG suffixes). The pipeline starts from adversarial responses and works backward to features, but does not test whether natural jailbreak attacks engage these same feature subgroups. Adding a verification step showing that real jailbreak attacks activate the identified features would strengthen the causal claim.
  4. §3 (Experimental Setup) and §4 (Results): The LLM judge (Grok) is acknowledged as potentially biased with no ground truth, yet no calibration, inter-annotator agreement, or human validation is reported. The 1–5 harmfulness scale is used to make comparative claims across methods and layers, but the reliability of these scores is unknown. At minimum, a human spot-check on a subset of responses (e.g., 50–100 samples) to validate judge agreement would substantially increase confidence in the reported results.
minor comments (7)
  1. The abstract states features 'appear across early and mid-to-late layers, with stronger concentration in mid-to-late layers,' while §4 states 'layer ~13 onward' for cluster-based and 'layers 16–25' for single-token. The specific layer ranges should be stated consistently.
  2. §2, Approach 1: The notation for feature sets and clusters (F, C, a′_k, a′′_k) is introduced rapidly and could benefit from a small worked example or a clearer step-by-step summary.
  3. §2, Approach 2: The threshold of 50 members for linkage clusters is stated without justification. A brief rationale or sensitivity analysis would improve clarity.
  4. Figures 1–4 are referenced as heatmaps but the color scale and axis labels are not described in the text. A brief caption explanation of what the color intensity represents would aid interpretation.
  5. §4 (Hierarchical Linkage-Based Steering): The statement 'if a linkage cluster is assigned a number greater than 16383, it may be made up of more than 50 features' is unclear. The relationship between cluster numbering and cluster size should be more informative.
  6. The reference to 'Figure §3' and 'Figure §4' in the text mixes section and figure numbering. Consistent referencing is needed.
  7. Appendix A.1 (Related Work) is placed in the appendix. Given the paper's brevity, integrating a condensed version into the main text would help readers unfamiliar with the SAE interpretability literature.

Simulated Author's Rebuttal

4 responses · 1 unresolved

We thank the referee for a careful and substantive review. The comments on control conditions, statistical rigor, circularity, and judge calibration are well-taken. Below we address each point and describe revisions we will make. We agree that the control experiment (random/non-harmful feature steering) is the most critical addition and will incorporate it in the revision.

read point-by-point responses

  1. Referee: No control condition with arbitrary/random/non-harmful SAE features steered in layers 16-25. Without this, observed harmfulness could reflect general later-layer disruptiveness rather than specificity of the concept-alignment and token-grouping pipeline.

    Authors: This is the most important critique and we fully agree. The current manuscript does not include a random-feature or non-harmful-concept control, and without it the specificity claim is not supported. We will add a control experiment in which we steer randomly sampled SAE features from the same layers (16-25) and the same number of features per layer as in our pipeline, using the same steering coefficient. We will also add a non-harmful-concept control: extracting benign concepts from safe responses, aligning them to prompt tokens via the same subspace-similarity procedure, and steering the resulting features. If random-feature steering in layers 16-25 produces comparable harmfulness increases to our pipeline-selected features, we will revise the specificity claims accordingly. We expect the control to show lower harmfulness than concept-aligned features because not all later-layer perturbations should produce unsafe outputs, but we acknowledge this is currently an expectation rather than a demonstrated result. We will report the control alongside the main results in the revised manuscript. revision: yes

  2. Referee: Results presented as raw counts with no error bars, confidence intervals, or significance tests. Claims of 'comparable' performance between single-token-driven and cluster-based grouping are not statistically supported.

    Authors: We agree that the current presentation lacks statistical rigor. With 20 prompts per category, raw counts alone are insufficient to support comparative claims. We will add bootstrap confidence intervals (1000 resamples) for the per-layer, per-category harmfulness counts. For the claim that single-token-driven grouping achieves 'comparable' harmfulness to cluster-based grouping, we will report a paired comparison (McNemar's test on per-prompt binary outcomes, plus a bootstrap on the difference in means) across all layer-category pairs. If the confidence intervals for the difference overlap zero, we will state that the methods are statistically indistinguishable; if not, we will revise the 'comparable' language. We will also add error bars to all figures. revision: yes

  3. Referee: The pipeline is partially circular: features selected for alignment with harmful concepts extracted from harmful responses, then harmfulness confirmed by steering. No verification that real jailbreak attacks (e.g., GCG suffixes) activate the same feature subgroups.

    Authors: This is a fair and important point. We acknowledge that the current pipeline starts from adversarial responses and works backward to features, and that this design has a degree of circularity. The referee is correct that testing whether real jailbreak attacks engage the identified feature subgroups would substantially strengthen the causal claim. We will add a verification experiment: we will generate GCG adversarial suffixes for a subset of prompts, run the attacked prompts through Gemma-2-2B, extract SAE feature activations at each layer, and measure the overlap between the activated features and our identified feature subgroups. We will report the fraction of identified subgroup features that activate above a threshold during actual GCG attacks, compared to a baseline of random feature sets. If the overlap is low, we will explicitly weaken the causal claim and reframe the contribution as identifying steerable vulnerability-relevant features rather than features that are necessarily engaged by natural jailbreak attacks. We cannot fully resolve the circularity concern without this experiment, so we list it as a partial revision pending results. revision: partial

  4. Referee: LLM judge (Grok) acknowledged as potentially biased with no calibration, inter-annotator agreement, or human validation. Reliability of 1-5 harmfulness scores is unknown.

    Authors: We agree that judge reliability should be validated. We will add a human spot-check: two annotators will independently score a random subset of 100 responses (sampled across methods, layers, and harm categories) on the same 1-5 scale. We will report inter-annotator agreement (Cohen's kappa or Krippendorff's alpha) and the agreement between the human annotations and the Grok judge scores. If agreement is moderate or lower, we will acknowledge this as a limitation and discuss its implications for the comparative claims. We will also report the distribution of judge scores to check for systematic biases (e.g., clustering at certain scale points). We note that we cannot fully eliminate judge bias concerns, but the human validation will at least characterize the reliability of the scores used for comparison. revision: yes

standing simulated objections not resolved
  • The circularity concern is partially inherent to the pipeline design: we extract harmful concepts from adversarial responses, align them to features, and then show that steering those features increases harmfulness. Even with the GCG verification experiment, the pipeline still starts from adversarial responses. If the GCG overlap is low, we will need to substantially reframe the claims, but we cannot eliminate the structural circularity without a fundamentally different approach (e.g., starting from attacks rather than responses). We acknowledge this as a genuine limitation that the revision can mitigate but not fully resolve.

Circularity Check

0 steps flagged

No significant circularity: the pipeline is methodologically tautological in a weak sense but contains no self-definitional reduction or self-citation chain.

full rationale

The paper's pipeline extracts harmful concepts from adversarial responses, aligns them to prompt tokens via cosine similarity, finds SAE features for those tokens, steers those features, and measures increased harmfulness. A reader might worry this is circular: features selected for alignment with harmful concepts produce harmful outputs when amplified. However, this is not circularity in the technical sense required by the framework. The pipeline inputs are (1) adversarial prompt-response pairs from BeaverTails, (2) a subspace generator model G from Wu et al. 2025, and (3) pre-trained SAEs from Lieberum et al. 2024. The output is a claim about which layers and feature subgroups are most steerable. No step in the derivation chain reduces to its own inputs by construction: the cosine similarity between subspace vectors and residual-stream activations is a genuine measurement, not a definition; the SAE feature extraction uses externally trained SAEs; the steering and harmfulness judging are independent operations. The key citations (Wu et al. 2025 for the subspace generator, Lieberum et al. 2024 for SAEs, Rufail et al. 2025 for layer-20 choice) are to external work by different author sets, not self-citations. The absence of a random-feature control condition (the skeptic's main concern) is a correctness/validity issue about whether the identified features are specifically vulnerability-relevant, not a circularity issue about whether the derivation reduces to its inputs by construction. The paper does not claim to predict a quantity that was fitted or defined into existence. The finding that single-token-driven grouping achieves comparable harmfulness to cluster-based grouping is an empirical comparison between two methods, not a tautological consequence of definitions. Score 2 reflects the mild methodological concern that the pipeline selects features for harmful-concept alignment and then demonstrates they produce harm, but this falls short of the specific reductions required for a circularity finding.

Axiom & Free-Parameter Ledger

6 free parameters · 5 axioms · 0 invented entities

The paper introduces no new entities, particles, forces, or postulated objects. It uses existing tools (SAEs, subspace generators, clustering) on an existing model (Gemma-2-2B) with an existing dataset (BeaverTails). The 'feature subgroups' are groupings of existing SAE features, not new entities. The free parameters are design choices (layer number, top-k counts, thresholds, sample size) rather than fitted constants.

free parameters (6)
  • Layer 20 for concept-token alignment = 20
    The choice to compute cosine similarity at layer 20 is justified by citing Rufail et al. 2025, but it is a fixed design choice that affects which tokens are selected as 'pos' and thus which features are identified.
  • Top-3 features for steering (cluster-based and linkage) = 3
    The number of top features selected from each cluster for steering is stated as 3 with no justification or sensitivity analysis.
  • Top-2 features per token per layer (single-token approach) = 2
    The single-token approach uses top-2 features per layer per token; this choice is not justified against alternatives.
  • Linkage cluster size threshold = 50
    The 50-member threshold for hierarchical linkage (§4) is introduced as ensuring specificity but is a post-hoc parameter with no sensitivity analysis.
  • Steering coefficient/magnitude = not stated
    The amplification factor applied to features during steering is not specified in the main text, making it impossible to assess whether the intervention strength is calibrated or arbitrary.
  • 20 prompts per category sample size = 20
    The sample of 20 prompts per category (280 total) is stated as a deliberate choice for mechanistic understanding but limits statistical power; no power analysis is provided.
axioms (5)
  • domain assumption SAE features are monosemantic and interpretable, such that amplifying a feature subgroup corresponds to amplifying a coherent concept.
    The entire pipeline depends on SAE features being meaningful units. This is assumed from Lieberum et al. 2024 but not independently verified for the adversarial safety domain in this paper.
  • domain assumption Cosine similarity between subspace vectors and residual stream activations at layer 20 is a valid measure of concept-token alignment.
    §2 states this computation as the basis for identifying 'pos' tokens. The validity of this similarity measure for concept alignment is assumed, citing Rufail et al. 2025 for the layer choice.
  • domain assumption An LLM judge (Grok-4-1-fast-non-reasoning) provides a reliable harmfulness score on a 1-5 scale.
    The evaluation protocol uses Grok as judge with the template from Qi et al. 2024. The paper acknowledges potential bias but treats scores as reliable enough for comparison. No calibration or human agreement is reported.
  • ad hoc to paper Features identified from adversarial responses and amplified via steering are causally relevant to jailbreak vulnerability.
    The paper infers causal relevance from steering efficacy (amplification increases harmfulness), but does not test whether suppressing these features prevents jailbreaks or whether they are engaged during actual adversarial attacks.
  • domain assumption Single-category BeaverTails prompts provide clean, unconfounded concept signals.
    §3 states the single-category constraint ensures extracted concepts are attributable to one harm type. This assumes category labels are mutually exclusive in practice, which may not hold for overlapping harm types.

pith-pipeline@v1.1.0-glm · 10897 in / 3470 out tokens · 157878 ms · 2026-07-04T15:41:22.433513+00:00 · methodology

Review history (2 revisions) →

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

2 extracted references · 2 canonical work pages · 2 internal anchors

  1. [1]

    Gemma 2: Improving Open Language Models at a Practical Size

    Understanding and enhancing the transferabil- ity of jailbreaking attacks.The Thirteenth Interna- tional Conference on Learning Representations. Anay Mehrotra, Manolis Zampetakis, Paul Kassianik, Blaine Nelson, Hyrum Anderson, Yaron Singer, and Amin Karbasi. 2024. Tree of attacks: Jailbreaking black-box llms automatically.Advances in Neural Information Pr...

    work page internal anchor Pith review Pith/arXiv arXiv 2024

  2. [2]

    Jailbroken: How Does LLM Safety Training Fail?

    Jailbroken: How does llm safety training fail? arXiv preprint arXiv:2307.02483. Zhengxuan Wu, Aryaman Arora, Atticus Geiger, Zheng Wang, Jing Huang, Dan Jurafsky, Christopher D Man- ning, and Christopher Potts. 2025. Axbench: Steer- ing LLMs? even simple baselines outperform sparse autoencoders. InF orty-second International Confer- ence on Machine Learni...

    work page internal anchor Pith review Pith/arXiv arXiv 2025