Source-linked AI summary

The First Token Is a Clue: Verbalizing Multi-Token Concepts from the J-lens

Xijie Gong, Tonghan Wang

arXiv:2608.31084v1cs.CL

TL;DR

J-lens cannot directly represent multi-token concepts because it reads hidden states as vocabulary-token rankings. This paper recovers first-token clues and complete-concept vectors using the frozen model, achieving higher readout and causal-swap performance than Template Lens across three models.

  • Problem

    J-lens provides vectors and scores for individual vocabulary tokens but lacks a representation for complete multi-token concepts.

  • Method

    The method uses J-lens to propose first tokens, lets the frozen model complete candidate concepts, recovers vectors from subsequent hidden states, and scores candidates with the vocabulary.

  • Results

    Across 496 multi-hop clozes, average Rank@10 was 43.1% versus 27.6% for Template Lens, while causal-swap succ@10 was 61.4% versus 26.2%.

  • Takeaways & Limitations

    First-token clues guide multi-token concept recovery, while subsequent hidden states provide vectors for readout and causal intervention.

  • Takeaways & Limitations

    The method depends on J-lens ranking a useful first token highly enough and focuses on canonical short English phrases, leaving longer or arbitrary expressions unestablished.

Abstract

from arXiv · show

The Jacobian Lens (J-lens) is a recent tool for interpreting LLMs. It reads a hidden state as a ranked list of vocabulary tokens, leaving multi-token concepts without a representation of their own. The original J-lens work addresses this limitation with Template Lens, which precomputes vectors for a fixed phrase vocabulary, and Oracle Lens, which fine-tunes components to propose phrases and reconstruct phrase vectors. We ask whether multi-token concepts and their vectors can instead be recovered directly from J-lens and the frozen model. We find that the first token of a multi-token concept is about as readable as a single-token concept. Given the correct first token and source prompt, the frozen model recovers the second token in 88.3% of two-token cases. We show that a vector for the complete concept can be recovered from subsequent hidden states in a single forward pass. We therefore use J-lens to propose first tokens and let the frozen model complete candidate concepts. We then recover a vector for each candidate and score it alongside the complete vocabulary. Across 496 multi-hop clozes on Gemma-3-12B-IT, Llama-3.1-8B, and Qwen3-14B, our method achieves an average $\mathrm{Rank@}10$ of 43.1%, compared with 27.6% for Template Lens. Without the J-lens clue, performance drops to 21.6%, showing that the first-token clue substantially improves readout. Causal concept swaps using the recovered vectors achieve an average $\mathrm{succ}@10$ of 61.4%, compared with 26.2% for Template Lens under the same intervention. These results show that first-token clues can guide multi-token concept recovery, while subsequent hidden states provide vectors for readout and intervention.

1 INTRODUCTION

The paper develops an inference-time method for recovering multi-token concepts and vectors from J-lens clues and a frozen model. It finds that first-token clues support concept completion, subsequent hidden states recover concept vectors, and the method improves readout and causal swaps over Template Lens.

  • Diagnosis: 54.6% of multi-token concepts had their correct first token in the J-lens top-10, close to 56.9% for single-token concepts.For two-token concepts, the correct first token and source prompt recovered the second token in 88.3% of cases.
  • Diagnosis: 97.4% Top-1 accuracy was achieved when recovered vectors identified matching ground-truth vectors across three models.The first fragment’s token vector and the mean of both fragment vectors never exceeded 59.6%.
  • Method: The method proposes first tokens with J-lens, completes candidates using the frozen model, recovers candidate vectors, and scores them jointly with the vocabulary.PROPOSE uses scaffold prompts; SCORE uses subsequent hidden states in a fixed carrier prompt.
  • Validation: 43.1% average Rank@10 was achieved, compared with 27.6% for Template Lens across 496 multi-hop clozes on three models.Removing the J-lens clue reduced Rank@10 to 21.6%.
  • Validation: 61.4% average succ@10 was achieved in causal concept swaps, compared with 26.2% for Template Lens under the same intervention.The recovered vectors supported replacing one intermediate concept with another.

2 PRELIMINARIES

J-lens propagates hidden states through layer-specific Jacobians and decodes them with the model’s unembedding into ranked vocabulary tokens. Because multi-token concepts lack vocabulary entries, extending readout requires identifying complete concepts and recovering their vectors.

  • 2.1 THE J-LENS READOUT: J-lens reads an intermediate hidden state through its effect on later model representations and produces a ranked list of vocabulary tokens.The model’s unembedding matrix supplies one vector for each vocabulary token.
  • 2.1 THE J-LENS READOUT: The J-lens ranked list is obtained by sorting the lens output over the vocabulary.The top-10 tokens are described as those the model is poised to verbalize.
  • 2.2 MULTI-TOKEN CONCEPT READOUT: A multi-token concept such as black hole has no vocabulary row, so J-lens provides no single vector or score for the complete concept.A constituent token such as black does not distinguish black hole from black coffee.
  • 2.2 MULTI-TOKEN CONCEPT READOUT: Extending J-lens requires identifying the complete concept and recovering its concept vector.Constituent-token sums or means cannot encode word order, assigning the same vector to black hole and hole black.
  • 2.2 MULTI-TOKEN CONCEPT READOUT: Template Lens uses a fixed phrase vocabulary and precomputed vectors, whereas Oracle Lens learns components to generate phrases and reconstruct their vectors.Both extensions add machinery operating directly on complete phrases.

3 WHAT IS ALREADY AVAILABLE?

J-lens supplies a useful first-token clue for multi-token concepts, while subsequent hidden states support recovery of vectors for complete concepts.

  • Multi-token readout requires both identifying the complete concept and assigning it a vector to score.
  • First-token Top-10 tests whether a multi-token concept’s first token is unusually difficult to read relative to single-token concepts.
  • 79.5–94.3%: given the correct first token, the frozen model recovers the second token in two-token cases.
  • Subsequent hidden states are used to recover a vector for a complete concept after the model is instructed to keep that concept in mind.
  • 97.4%: recovered vectors identify matching ground-truth vectors with at least this Top-1 accuracy across all three models.First-fragment and mean vectors never exceed 59.6%.

4 PROPOSING AND SCORING MULTI-TOKEN CONCEPTS

The method proposes complete concepts from J-lens first-token candidates, recovers each candidate’s vector from carrier-prompt hidden states, transports it, and scores it with the vocabulary.

  • 4.1 PROPOSING CANDIDATE CONCEPTS: PROPOSE selects possible first tokens from J-lens ranks and lets the frozen model complete each candidate from the source prompt.
  • 4.1 PROPOSING CANDIDATE CONCEPTS: Rank ranges receive independent search budgets, preserving representation for lower-ranked first tokens.
  • 4.1 PROPOSING CANDIDATE CONCEPTS: The fixed Scaffold Prompt keeps the J-lens seed as the first token while the model generates remaining tokens normally.
  • 4.1 PROPOSING CANDIDATE CONCEPTS: Candidates are deduplicated, shorter strict prefixes are removed, and a fixed number of final candidates is retained.
  • 4.2 SCORING CANDIDATE CONCEPTS: SCORE recovers a vector for each candidate from hidden states after the complete candidate is processed in a fixed Carrier Prompt.
  • 4.2 SCORING CANDIDATE CONCEPTS: Position-specific means remove shared components, while pooling weights are fixed using an independent calibration corpus.
  • 4.2 SCORING CANDIDATE CONCEPTS: Recovered vectors are transported from the carrier layer to the source layer by matching their images under two J-lens Jacobians.
  • 4.2 SCORING CANDIDATE CONCEPTS: Candidate concepts are ranked jointly with the complete vocabulary using the same scoring procedure.

5 EXPERIMENTS

The experiments evaluate unnamed multi-token concept readout and causal control using recovered vectors, comparing the method with Template Lens and ablations.

  • 5.1 END-TO-END READOUT: The full pipeline reads unnamed multi-token concepts by proposing candidates, recovering their vectors, and ranking them with the complete vocabulary.Rank@k counts prompts whose target reaches the top-k at any evaluated layer; unproposed targets count as failures.
  • 5.1 END-TO-END READOUT: The J-lens clue is necessary: removing it consistently reduces readout performance while leaving the source prompt and SCORE stage unchanged.This comparison isolates the contribution of the first-token clue rather than changing the scoring procedure.
  • 5.1 END-TO-END READOUT: Our method achieves the highest Rank@10 across all three models, outperforming Template Lens on the 496-prompt evaluation.The evaluation targets multi-token concepts that never appear in the source text.
  • 5.2 CAUSAL CONCEPT SWAP: The causal swap removes the source concept vector and adds the partner vector over a fixed short layer band.The partner-vector magnitude is measured from the partner prompt, so no fitted strength coefficient is required.
  • 5.2 CAUSAL CONCEPT SWAP: The intervention tests whether recovered vectors control the corresponding computation, not merely whether they correlate with the source concept.All vector constructions use the same swap operator, so success-rate differences are attributed to the vectors rather than separately tuned interventions.
  • 5.2 CAUSAL CONCEPT SWAP: The full swap gives the highest succ@10 across all three models, while addition alone, deletion alone, and vector controls produce weaker effects.The comparison includes Template Lens vectors under the same swap operator.

6 RELATED WORK

The paper relates its approach to lens-based readout, LLM-based activation interpretation, sub-token representation studies, and concept-vector methods.

  • Lens-based readout: J-lens, Logit Lens, and Tuned Lens make intermediate representations readable through different mappings that ultimately decode model states.J-lens uses an averaged Jacobian to connect an intermediate state to later model representations.
  • LLMs as activation interpreters: The frozen LLM is used to interpret candidate concepts without receiving the source activation or verbalizing the source activation directly.This distinguishes the approach from SelfIE and Patchscopes.
  • Representations beyond tokenizer granularity: Prior work suggests that internal computational units can exceed tokenizer granularity, with complete entities remaining represented as constituent-token information fades.Other work recovers word-level identity from subword-fragment representations.
  • Concept vectors and sparse features: Unlike sparse-feature and contrastive activation-steering methods, this approach starts from J-lens-ranked first tokens and recovers vectors from subsequent hidden states.The recovered vectors support both readout and causal control through the same vector-based scoring framework.

7 DISCUSSION

The results support recovering and controlling multi-token concepts beyond direct J-lens vocabulary-token output, while identifying candidate proposal as the main limitation.

  • Discussion: The first token of a multi-token concept is about as readable as an ordinary single-token concept, and subsequent hidden states retain a complete-concept vector.The recovered vector closely matches independently defined ground-truth vectors.
  • Discussion: Recovered vectors identify concepts absent from source prompts and support causal replacement of one intermediate concept with another.The paper therefore reports multi-token verbalization without a fixed phrase vocabulary or separately trained proposal and reconstruction components.
  • Limitations: Candidate proposal depends on J-lens ranking a useful first token highly enough for PROPOSE to consider it.If the first token ranks poorly, the correct concept may never enter the candidate set even when SCORE could identify it.
  • Limitations: The experiments focus on canonical multi-token concepts expressible as short English phrases, leaving long descriptions and arbitrary natural-language expressions unestablished.The frozen model must also have enough information in the source prompt to complete the concept from the selected first token.

A DATASET AND EXPERIMENTAL SETUP

The dataset uses counterfactual cloze pairs with unnamed multi-token intermediates, while setup details specify construction constraints, released J-lens fits, and a separate vector-recovery diagnostic.

  • Dataset construction: Each cloze follows a common frame in which a relation and description identify an unnamed multi-token intermediate before a usually single-word answer.Counterfactual pairs share the relation but differ in the intermediate, so exchanging intermediates exchanges answers.
  • Dataset construction: The intermediate is not a vocabulary entry, so J-lens exposes constituent-token clues without directly representing the complete concept or its endpoint.A high rank for the first token therefore provides a clue rather than a complete readout.
  • Dataset construction: GPT-5.5 writes clue sentences under fixed constraints requiring canonical multiword expressions, unnamed answers, shared relations, and producible English answers.Pre-measurement filters require intermediates to span several tokens under all three tokenizers.
  • Experimental setup: The study uses released J-lens fits without refitting, covering 47 Gemma layers, 31 Llama layers, and 39 Qwen layers.End-to-end readouts use the final input position with an appended literal space and leading-space concept tokens.
  • Experimental setup: The vector-recovery diagnostic splits 500 whole-token words into forced two-token fragments and ranks reconstructions by cosine against whole-token J-lens ground-truth vectors.The diagnostic includes both recognizable cuts such as blackmail → black+mail and mid-word fragments such as activity → act+ivity.

B.1 CANDIDATE GENERATION

Candidate generation uses J-lens-ranked roots and frozen-model completion, while subsequent hidden states support concept-vector recovery, transport, and scoring. The intervention applies these recovered vectors through a short residual-stream edit.

  • B.1 CANDIDATE GENERATION: PROPOSE selects roots independently from J-lens rank bands 1–20, 21–50, and 51–100 before frozen-model completion.Up to twenty eligible roots are considered in each band, with active widths 12, 6, and 6.
  • B.1 CANDIDATE GENERATION: The carrier prompt presents each generated candidate as a remembered concept, with the source prompt and final answer withheld during extraction.The extraction window uses tokens in the fixed suffix after the concept, while the concept’s final token is held out as a control.
  • B.1 CANDIDATE GENERATION: Calibration quantities, transport parameters, carrier layers, and extraction windows are fixed using data disjoint from evaluation prompts and labels.The procedure fixes µt, A, λ, ℓc, and T before evaluation, with model-specific carrier layers selected by mean paired cosine.
  • B.1 CANDIDATE GENERATION: Equation 2 matches the centered source representation to the candidate representation using ridge regularization, while Equation 3 implements diagonal-covariance whitening.The ridge objective approaches an exact match as λ → 0 when an exact match exists; whitening yields the linear score that best separates concept-present and concept-absent cases under the stated covariance approximation.
  • B.1 CANDIDATE GENERATION: Causal trials filter concepts and answers before intervention, then edit a short residual band by deleting the component along the concept vector and adding a prompt-derived magnitude.The edit runs as the forward pass proceeds, with deletion over the final twelve prompt positions and addition near the prompt’s end.

D COMPUTATIONAL COST

The method adds candidate-generation decoding and a carrier forward for each previously unseen concept, unlike baselines that use fixed phrase vocabularies or fine-tuning. This incurs measurable per-prompt runtime on one A800.

  • D COMPUTATIONAL COST: 0.4–0.8 s per prompt is the reported runtime for the extended J-lens pipeline on one A800.The added cost comes from candidate-generation decoding and one short carrier forward for each previously unseen concept.
  • D COMPUTATIONAL COST: Template Lens uses a fixed phrase vocabulary and a few hundred prefix forwards, whereas Oracle Lens fine-tunes two copies of the subject model.These baseline designs differ from the added per-prompt computation of the proposed method.
Loading 2608.31084v1…