Add VJP-delta activation steering

docs/.nav.yml

@@ -50,6 +50,7 @@ nav:
         - DirectionalAblation: "examples/notebooks/algorithms/directional_ablation.ipynb"
         - ITI: "examples/notebooks/algorithms/iti.ipynb"
         - PASTA: "examples/notebooks/algorithms/pasta.ipynb"
+        - VJPDelta: "examples/notebooks/algorithms/vjp_delta.ipynb"
       - Output control:
         - BestOfN: "examples/notebooks/algorithms/best_of_n.ipynb"
         - BudgetForcing: "examples/notebooks/algorithms/budget_forcing.ipynb"
@@ -114,6 +115,7 @@ nav:
             - Directional Ablation: reference/algorithms/state_control/directional_ablation.md
             - ITI: reference/algorithms/state_control/iti.md
             - PASTA: reference/algorithms/state_control/pasta.md
+            - VJPDelta: reference/algorithms/state_control/vjp_delta.md
         - Output control:
             - Base classes: reference/algorithms/output_control/base_output_control.md
             - Common library: reference/algorithms/output_control/common.md

docs/concepts/controls.md

@@ -152,6 +152,9 @@ patching. The toolkit implements:
 - `PASTA` ([API reference](../reference/algorithms/state_control/pasta.md), [notebook](../examples/notebooks/algorithms/pasta.ipynb))
     - *Description*: post-hoc attention steering[@zhang2024tell], rescaling attention to targeted prompt substrings at selected layers and heads. The `head_config` argument takes a dict or list of layers and heads, or a `HeadProfile` recipe that runs the paper's head-profiling stage as a steer-time fit on the loaded model (scoring each candidate head by its paired lift over an unsteered baseline) and freezes the resolved head map.
     - *Backends*: HF with `attn_implementation` `"eager"` or `"sdpa"` (attention-map writes have no engine form).
+- `VJPDelta` ([API reference](../reference/algorithms/state_control/vjp_delta.md), [notebook](../examples/notebooks/algorithms/vjp_delta.ipynb))
+    - *Description*: fits a target-state contrast and uses vector-Jacobian products to derive one normalized additive direction per earlier residual layer. `VJPDeltaFit` uses raw prompts, excludes `skip_first` positions and each row's final real token from the VJP spans, and averages each class independently before subtraction.
+    - *Backends*: HF for fitting. The frozen form uses the existing `ActivationAdapter` additive intervention.
 
 Reusable building blocks shared across the residual-stream methods (estimators, gating, selectors, transforms,
 steering vectors, hook utilities) are located in

docs/reference/algorithms/state_control/vjp_delta.md

@@ -0,0 +1,20 @@
+# VJPDelta
+
+::: steerability.algorithms.state_control.vjp_delta
+    handler: python
+    options:
+        show_if_no_docstring: true
+        show_source: true
+        show_root_heading: true
+        docstring_style: google
+        show_root_full_path: true
+        show_object_full_path: false
+        separate_signature: false
+        inherited_members: true
+        show_submodules: true
+        show_symbol_type_heading: true
+        show_symbol_type_toc: true
+        filters:
+          - "!.*Args$"
+          - "!^registry"
+          - "!^STEERING_METHOD"

examples/index.md

@@ -59,6 +59,8 @@ Algorithm notebooks demonstrate how each method (i.e., control) operates. The me
 
     :octicons-arrow-right-24: [PASTA](./notebooks/algorithms/pasta.ipynb)
 
+    :octicons-arrow-right-24: [VJPDelta](./notebooks/algorithms/vjp_delta.ipynb)
+
 -   __Output control__
 
     ---

examples/notebooks/algorithms/vjp_delta.ipynb

@@ -0,0 +1,170 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "24c9ee9c",
+   "metadata": {},
+   "source": [
+    "# VJP-delta\n",
+    "\n",
+    "`VJPDelta` fits an additive direction at each chosen source layer. It reads a contrast at a later target layer, then uses a vector-Jacobian product (VJP) to map that contrast back to each source layer. The fitted vectors use the usual state-control additive intervention during generation.\n",
+    "\n",
+    "VJP-delta follows [Clark, Michael J. (2026), _vjp-steering: contrastive steering vectors from vector-Jacobian products_](https://github.com/wassname/vjp-steering), adapting the [Jacobian lens](https://transformer-circuits.pub/2026/workspace/)."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2fe4a263",
+   "metadata": {},
+   "source": [
+    "## Setup\n",
+    "\n",
+    "This CPU demonstration loads a small Hugging Face Llama checkpoint through the public pipeline API."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "1c306fb6",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "[transformers] Model config: pad_token_id must be `None` or an integer within the vocabulary (between 0 and 31999), got -1. This may result in unexpected behavior.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "[transformers] The following generation flags are not valid and may be ignored: ['pad_token_id']. Set `TRANSFORMERS_VERBOSITY=info` for more details.\n"
+     ]
+    }
+   ],
+   "source": [
+    "from pathlib import Path\n",
+    "import tempfile\n",
+    "\n",
+    "from steerability.algorithms.core.steering_pipeline import SteeringPipeline\n",
+    "from steerability.algorithms.state_control.vjp_delta import VJPDelta\n",
+    "from steerability.spipe import SPipe\n",
+    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+    "\n",
+    "model_name = \"hf-internal-testing/tiny-random-LlamaForCausalLM\"\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
+    "tokenizer.pad_token = tokenizer.eos_token\n",
+    "model = AutoModelForCausalLM.from_pretrained(model_name)\n",
+    "fit_data = {\n",
+    "    \"positives\": [\"the cat sat\", \"the dog ran\"],\n",
+    "    \"negatives\": [\"dog ran fast\"],\n",
+    "}\n",
+    "control = VJPDelta(\n",
+    "    data=fit_data,\n",
+    "    target_layer=1,\n",
+    "    source_layer_ids=[0],\n",
+    "    skip_first=0,\n",
+    "    strength=0.5,\n",
+    ")\n",
+    "pipeline = SteeringPipeline(\n",
+    "    model=model,\n",
+    "    tokenizer=tokenizer,\n",
+    "    controls=[control],\n",
+    "    model_name_or_path=model_name,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5c87a2e7",
+   "metadata": {},
+   "source": [
+    "## Extract and steer\n",
+    "\n",
+    "The raw prompts have unequal positive and negative pool sizes. `steer()` extracts the vectors, then binds the standard additive intervention. The stored direction rows are unit norm."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "ca2ba708",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'layers': [0], 'reply': \"'agreedָָ'\"}\n"
+     ]
+    }
+   ],
+   "source": [
+    "pipeline.steer()\n",
+    "vector = control.export_state()[\"intervention_0/transform\"]\n",
+    "assert set(vector.directions) == {0}\n",
+    "assert all(abs(direction.norm().item() - 1.0) < 1e-5 for direction in vector.directions.values())\n",
+    "\n",
+    "reply = pipeline.generate(text=\"the cat\", max_new_tokens=3, do_sample=False)\n",
+    "print({\"layers\": sorted(vector.directions), \"reply\": repr(reply)})"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5b74b575",
+   "metadata": {},
+   "source": [
+    "## Freeze and reload\n",
+    "\n",
+    "The frozen form contains the fitted vectors as an `ActivationAdapter`. Reloading it resolves the stored additive artifact and does not run another VJP fit."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "651ceb9c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'bundle': '/tmp/tmpx73lemsj/vjp_delta_demo', 'reply_matches': True}\n"
+     ]
+    }
+   ],
+   "source": [
+    "bundle = Path(tempfile.mkdtemp()) / \"vjp_delta_demo\"\n",
+    "saved = pipeline.to_spipe().save(bundle)\n",
+    "reloaded = SPipe.load(saved).pipeline()\n",
+    "assert type(reloaded.state_controls[0]).__name__ == \"ActivationAdapter\"\n",
+    "reloaded.model, reloaded.tokenizer = model, tokenizer\n",
+    "reloaded.steer()\n",
+    "reloaded_reply = reloaded.generate(text=\"the cat\", max_new_tokens=3, do_sample=False)\n",
+    "assert reloaded_reply == reply\n",
+    "print({\"bundle\": str(saved), \"reply_matches\": reloaded_reply == reply})"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.13.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

steerability/algorithms/state_control/vjp_delta/__init__.py

@@ -0,0 +1,12 @@
+from .args import VJPDeltaArgs
+from .control import VJPDelta
+from .fit import VJPDeltaFit
+
+STEERING_METHOD = {
+    "category": "state_control",
+    "name": "vjp_delta",
+    "control": VJPDelta,
+    "args": VJPDeltaArgs,
+}
+
+__all__ = ["VJPDelta", "VJPDeltaArgs", "VJPDeltaFit"]

steerability/algorithms/state_control/vjp_delta/args.py

@@ -0,0 +1,67 @@
+"""Arguments for VJP-delta steering."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Sequence
+
+from steerability.algorithms.core.base_args import BaseArgs
+from steerability.algorithms.core.internals.data import LabeledExamples, as_labeled_examples
+from steerability.algorithms.state_control.common.sources import ArtifactSource
+from steerability.algorithms.state_control.common.steering_vector import SteeringVector
+from steerability.algorithms.state_control.common.token_scope import ScopeKind
+
+
+@dataclass
+class VJPDeltaArgs(BaseArgs):
+    """Arguments for `VJPDelta`.
+
+    Args:
+        steering_vector: A precomputed `SteeringVector` skips VJP fitting. An `ArtifactSource` resolves its own artifact.
+        data: Independent positive and negative raw prompt pools for `VJPDeltaFit`.
+        target_layer: Target layer for the contrast. None selects `num_layers - 3`.
+        source_layer_ids: Source layers for VJPs. None selects every layer before the target.
+        skip_first: Prefix positions excluded during gradient extraction.
+        max_length: Maximum tokenized fit-prompt length.
+        batch_size: Fit prompts per differentiable forward.
+        strength: Multiplier used when applying the normalized directions.
+        token_scope: Positions that receive the additive intervention during generation.
+        last_k: Required with `token_scope="last_k"`.
+        from_position: Required with `token_scope="from_position"`.
+    """
+
+    steering_vector: SteeringVector | ArtifactSource | None = None
+    data: LabeledExamples | dict | None = None
+    target_layer: int | None = None
+    source_layer_ids: Sequence[int] | None = None
+    skip_first: int = 16
+    max_length: int = 384
+    batch_size: int = 8
+    strength: float = 1.0
+    token_scope: ScopeKind = "after_prompt"
+    last_k: int | None = None
+    from_position: int | None = None
+
+    def __post_init__(self) -> None:
+        if (self.steering_vector is None) == (self.data is None):
+            raise ValueError("Provide exactly one of steering_vector or data.")
+        if isinstance(self.steering_vector, SteeringVector):
+            self.steering_vector.validate()
+        if self.data is not None and not isinstance(self.data, LabeledExamples):
+            self.data = as_labeled_examples(self.data)
+        if self.target_layer is not None and self.target_layer < 0:
+            raise ValueError("target_layer must be >= 0.")
+        if self.source_layer_ids is not None:
+            source_ids = tuple(int(layer_id) for layer_id in self.source_layer_ids)
+            if not source_ids or min(source_ids) < 0 or len(set(source_ids)) != len(source_ids):
+                raise ValueError("source_layer_ids must be a non-empty sequence of unique integers >= 0.")
+            self.source_layer_ids = source_ids
+        if self.skip_first < 0:
+            raise ValueError("skip_first must be >= 0.")
+        if self.max_length < 2:
+            raise ValueError("max_length must be >= 2.")
+        if self.batch_size < 1:
+            raise ValueError("batch_size must be >= 1.")
+        if self.token_scope == "last_k" and (self.last_k is None or self.last_k < 1):
+            raise ValueError("last_k must be >= 1 when token_scope is 'last_k'.")
+        if self.token_scope == "from_position" and (self.from_position is None or self.from_position < 0):
+            raise ValueError("from_position must be >= 0 when token_scope is 'from_position'.")

steerability/algorithms/state_control/vjp_delta/control.py

@@ -0,0 +1,62 @@
+"""VJP-delta control."""
+from __future__ import annotations
+
+from steerability.algorithms.state_control.base import InterventionControl
+from steerability.algorithms.state_control.common.sources import _Precomputed
+from steerability.algorithms.state_control.common.specs import CoveredLayers, Intervention, TokenScope
+from steerability.algorithms.state_control.common.steering_vector import SteeringVector
+from steerability.algorithms.state_control.common.transforms import AdditiveTransform
+
+from .args import VJPDeltaArgs
+from .fit import VJPDeltaFit
+
+
+class VJPDelta(InterventionControl):
+    """VJP-delta activation steering.
+
+    The control fits one normalized additive direction per source layer during `steer()`. Its
+    target contrast is the positive-minus-negative final unpadded target state. The fit applies
+    that contrast as a cotangent at valid target tokens and averages valid source-token gradients
+    per prompt before separately averaging the positive and negative classes. At generation it
+    uses the standard additive intervention and token scopes.
+
+    A precomputed `SteeringVector` skips VJP fitting. A supplied `ArtifactSource` resolves its own
+    artifact. The frozen form is `ActivationAdapter`, so a reloaded `.spipe` resolves the stored
+    vectors without a VJP fit.
+
+    Reference:
+
+    - Clark, Michael J. (2026). "vjp-steering: contrastive steering vectors from
+      vector-Jacobian products."
+      [https://github.com/wassname/vjp-steering](https://github.com/wassname/vjp-steering)
+      Adapts the [Jacobian lens](https://transformer-circuits.pub/2026/workspace/).
+    """
+
+    Args = VJPDeltaArgs
+    supports_batching = True
+
+    def _configure(self) -> None:
+        if self.steering_vector is None:
+            source = VJPDeltaFit(
+                data=self.data,
+                target_layer=self.target_layer,
+                source_layer_ids=self.source_layer_ids,
+                skip_first=self.skip_first,
+                max_length=self.max_length,
+                batch_size=self.batch_size,
+            )
+        elif isinstance(self.steering_vector, SteeringVector):
+            source = _Precomputed(self.steering_vector.clone())
+        else:
+            source = self.steering_vector
+        self._template = (
+            Intervention(
+                layers=CoveredLayers(),
+                transform=AdditiveTransform(source, strength=self.strength),
+                scope=TokenScope(self.token_scope, last_k=self.last_k, from_position=self.from_position),
+            ),
+        )
+
+    def cleanup(self) -> None:
+        """Drop fitted intervention tensors and their bound artifacts."""
+        self.interventions = ()