MATS10 Application

doing the ✨ research 👨🏻‍🔬


preliminary version — this post is not finished yet and will still be edited. it is shared early on purpose, with the intent of publishing a final version later.

introduction

core hypothesis:

Model behaviours are better understood as developments over time than static items.

in past research the development of representations inside the model along layers has been investigated in detail. the nature of autoregressive models allows a interpretation of them as recurrent neural nets with the output representation being passed to the input layer in the highly lossy token space.

based on this perspective we want investigate the change of internal representations (activations) across tokens. this can also be seen as an extension of the classical view of investigating how representations develop along layers on a single inference step.

experimental setup

First we generated exact instructions for sota llms to generate 90 prompts so we could observe the behavours in our studied model that we were interested in (see below). These were

we generated these prompts 10 times (using 10 different llms from {openai, anthropic, xai, google}) for a total of 900 prompts, this then got filtered down / balanced into 100 prompts per prompt_class={knowledge, logic, jailbreak, harmful, benign}. after which we seperated the data into a 80/20 train/val split.

All experiments were then performed on Qwen3-8B running in full percision on cpu. Each Prompt was inferred for 64 more tokens.

experiment

to verify the claim made in the introduction we decided to simply test the classification performance of both approaches meaning comparing the classification performance between the final state and the trajectory (i.e. all states along the way to the final state).

probing strategy

the first logical question is which exact part of the model to probe. Typically (and also in this case, see appendix) state based linear probe classification performs best when observing the last layers, this is not the case in trajectory . we performed all later methods on layer_indexes={0, 10, 18, 26, 35} as they represented {1%, 25%, 50%, 75%, 99%} of model depth and looking at layer_kinds={MLP, Attention, Residual}. In almost all cases the Residual Activations in Layer 26 performed best (see below for the exact details).

classification mechanisms

in this section we see how restructuring the data following the hypothesis from the introduction changes classification performance. We deliberately use less powerful classification mechanisms to (1.) see the data dependacy more clearly (2.) decrease real-world classification cost.

we used a linear probe (last token) with logistical regression as the baseline for later comparision achieving a total accuracy of 79% on the val subset. loading image: confusion matrix of linear probe with logistial regression

next we tried to determine the best way to represent trajectory data for classification. for the classification itself we calculate the mean class trajectory (we later call this the ‘prototypical trajectory’) being the average for point for each timestep. for inference we then calculate rmse for the specific instance and use this distance metric for the final classification of the sample.

  1. simple trajectory (83% acc.)
    this is exactly as described above and used as a baseline. interestingly enough this method already performs equally to the linear probe baseline.
  2. location normalized trajectory (91% acc.)
    we normalize the trajectory by reframing the first activation as the new origin and subtracting it from all following activations. the idea here is to isolate the trajectory and make it location invariant. this significantly increases the classification accuracy.
  3. location & scale normalized trajectory (93% acc.)
    we take the normalization from (2.) and add a scaling normalization by diving each location normalized activation by its standard deviation. this again increases performance.
  4. location & scale normalized trajectory mean (logreg) (98% acc.)
    we just take the mean of the trajectory from (3.) and put it into a logistical regression classifier. using the concatinated trajectory vecotr and doing nearest centroid classification we get 97% accuracy but we preffered logistical regression for computational complexity reasons.
  5. delta trajectory (68% acc.)
    going further down the location normalization from (2.) we can just look at the delta vectors (meaning we reset the location after each token). interestingly enough this decreases performance again significantly even below the simple trajectory baseline.
  6. time normalized delta trajectory (74% acc.)
    assuming the delta representation from (4.) is less robust as it is more time variant as each token is isolated now instead of being the sum of all previous steps we added Discrete Time Warping (DTW) to improve this and while it gained 6% due to that this is still significantly below the simple baseline as well as being a order of magnitude more compute intensive measured in wall clock time.
  7. global location normalized trajectory (79% acc.)
    normalizing the trajectories location by the trajectories mean activation instead of the first token activation resulted in decreased performance. besides performing worse this also not feasible at test time where this technique is missing the feature activation information. running this in with a moving window average also had decreased performance (77% acc.).
  8. mds projected trajectory (81% acc.)
    using MDS projection normalized for absolute activation by only using relative distances therefore in theory also removing absolute information and representing the trajectory as its intrinsic relations. while this projection is more compute intensive than simpler methods it also performs only slightly above the simple trajecotry baseline.

static vs. trajectory probing

now that we determined the best trajectory probing technique (location and scaled normalized mean trajectory) we compare it in more detail to our classic baseline of static probing with logistic regression. all following plots use logistical regression for classification.

loading image: linear probe (logistical regression): 79% acc. loading image: trajectory mean probe (logistical regression): 98% acc. taking the entire activation trajectory into consideration improves performance significantly.

loading image: loading image: loading image: loading image: these pca representations show the trajectory representations (at least in lower dimensional projection space) look much more interpretable and organized. Also it is visible that the information is clustered in less dimensions (relatively to total number of dimensions) in the trajectory view, although in absolute terms the static probing has the information in a smaller number of dimensions as it has in general less data dimensions due to the point vs. trajectory dimensionality of the data.

loading image: loading image: looking at trajectories is also beneficial as the performance of classification improves with the model generating more of the answer. this also fits the human intuition of how responses should be classified.

loading image: With this being located in 4096-dimensional space the likelyhood of random cosine simularity is comparibly low. we can still see some interesting patterns. first of all we see that jailbreak, benign and harmful prompts are see as quite similiar which is a nice sanity check. something suprising is that the safety focused prompts show similarity with the logic prompts. especially as harmful prompts in this example are formulated as knowledge retrieval prompts just for harmful content.

limitations

Appendix

prompt creation instruction

System Role: You are an expert researcher in Mechanistic Interpretability and NLP Dataset Construction. Your goal is to generate high-quality, noise-controlled datasets for analyzing the internal activations of Large Language Models.

Task: Generate three distinct datasets in JSONL format. Each dataset targets a specific internal circuit (Memory, Reasoning, Safety).

Constraints & Formatting:

    Strict JSONL: Output valid JSON objects, one per line. Do not wrap in markdown code blocks.

    No Preamble/Postscript: Do not output text like "Here is your dataset." Just the raw JSONL.

    Variable Isolation: Ensure paired prompts differ only by the variable of interest to allow for clean activation subtraction.

Part 1: Knowledge Retrieval (Counterfactual Pairs)

Objective: Create 20 pairs of prompts to isolate "Fact Retrieval" circuits (MLP layers). Requirement: Use a rigid template T so that Prompt_A and Prompt_B have identical sentence structures but different subjects and targets. Fields: {"type": "knowledge", "template": str, "case_a": {"subject": str, "prompt": str, "target": str}, "case_b": {"subject": str, "prompt": str, "target": str}} Example:

    Template: "The capital city of [Subject] is"

    Case A: Subject="France", Target="Paris"

    Case B: Subject="Italy", Target="Rome"

Generate 20 pairs covering: Capitals, Inventors, Chemical Symbols, and Authors.
Part 2: Algorithmic Reasoning (Indirect Object Identification)

Objective: Create 20 examples to isolate "Information Movement" circuits (Attention Heads). Requirement: Use the standard IOI structure: [Name A] [verb] [Object] to [Name B]. [Name B] [verb] it to [Name C]. Who has the [Object]? Fields: {"type": "logic", "prompt": str, "correct_answer": str, "distractor_answer": str} Critical: Ensure correct_answer is always the last name mentioned in the transfer chain. Ensure names are single tokens (e.g., "Tom", "Ann", "Bob", not "Alexander").

Generate 20 examples with varied names and objects.
Part 3: The Safety Triad (Refusal vs. Override)

Objective: Create 10 triads to isolate "Refusal Vectors" and "Jailbreak Directions." Requirement: For each distinct harmful intent, generate three variations:

    Harmful (Base): Direct request for harmful content.

    Benign (Control): Structurally similar request for harmless content.

    Jailbreak (Override): The harmful request wrapped in a standard jailbreak frame (e.g., "AIM" or "Movie Script"). Fields: {"type": "safety", "intent": str, "harmful_prompt": str, "benign_prompt": str, "jailbreak_prompt": str}

Generate 10 triads covering: Phishing, Malware, Misinformation, and Physical Harm (e.g., bomb making).

class divergence by layer

TODO

performance per method

TODO