diff --git a/.gitignore b/.gitignore index 500a01e8..71257afc 100644 --- a/.gitignore +++ b/.gitignore @@ -12,6 +12,7 @@ .cache # environments +.local/ .env .venv/ .venv**/ diff --git a/docs/.nav.yml b/docs/.nav.yml index d4ef49e8..2de381a7 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -47,6 +47,9 @@ nav: - AngularSteering: "examples/notebooks/algorithms/angular_steering.ipynb" - CAA: "examples/notebooks/algorithms/caa.ipynb" - CAST: "examples/notebooks/algorithms/cast.ipynb" + - CorDA-derived PCA: "examples/notebooks/algorithms/corda_pca.ipynb" + - S-space: "examples/notebooks/algorithms/sspace.ipynb" + - Linear-AcT: "examples/notebooks/algorithms/linear_act.ipynb" - DirectionalAblation: "examples/notebooks/algorithms/directional_ablation.ipynb" - ITI: "examples/notebooks/algorithms/iti.ipynb" - PASTA: "examples/notebooks/algorithms/pasta.ipynb" @@ -109,6 +112,9 @@ nav: - ActAdd: reference/algorithms/state_control/act_add.md - ActivationAdapter: reference/algorithms/state_control/activation_adapter.md - Angular Steering: reference/algorithms/state_control/angular_steering.md + - CorDA-derived PCA: reference/algorithms/state_control/corda_pca.md + - S-space: reference/algorithms/state_control/sspace.md + - Linear-AcT: reference/algorithms/state_control/linear_act.md - CAA: reference/algorithms/state_control/caa.md - CAST: reference/algorithms/state_control/cast.md - Directional Ablation: reference/algorithms/state_control/directional_ablation.md diff --git a/docs/concepts/controls.md b/docs/concepts/controls.md index 128ed76e..2bf67768 100644 --- a/docs/concepts/controls.md +++ b/docs/concepts/controls.md @@ -137,6 +137,15 @@ patching. The toolkit implements: - `AngularSteering` ([API reference](../reference/algorithms/state_control/angular_steering.md), [notebook](../examples/notebooks/algorithms/angular_steering.ipynb)) - *Description*: angular steering[@vu2025angular], rotating the hidden state within a per-layer 2D plane (feature axis + companion axis) to a target angle while leaving the orthogonal complement untouched. It is norm-preserving by construction, and vector addition and directional ablation are special cases. - *Backends*: HF, vLLM (`intervention_point="layer_output"` only, since the default norm-input placement is HF-only). +- `CordaPCA` ([API reference](../reference/algorithms/state_control/corda_pca.md), [notebook](../examples/notebooks/algorithms/corda_pca.ipynb)) + - *Description*: an activation-steering adaptation of CorDA's context-oriented decomposition, fitted from paired Linear inputs. + - *Backends*: HF. +- `SSpace` ([API reference](../reference/algorithms/state_control/sspace.md), [notebook](../examples/notebooks/algorithms/sspace.ipynb)) + - *Description*: weight-SVD steering with contrast-ranked coordinates and per-token cosine gates. + - *Backends*: HF. +- `LinearAcT` ([API reference](../reference/algorithms/state_control/linear_act.md), [notebook](../examples/notebooks/algorithms/linear_act.ipynb)) + - *Description*: coordinate-wise affine transport fitted by least squares on sorted activations. + - *Backends*: HF. - `CAA` ([API reference](../reference/algorithms/state_control/caa.md), [notebook](../examples/notebooks/algorithms/caa.ipynb)) - *Description*: contrastive activation addition[@panickssery2023steering], adding a learned mean-difference direction to the residual stream at a single layer. - *Backends*: HF, vLLM (norm-preserving configurations included). diff --git a/docs/reference/algorithms/state_control/corda_pca.md b/docs/reference/algorithms/state_control/corda_pca.md new file mode 100644 index 00000000..b30b73ad --- /dev/null +++ b/docs/reference/algorithms/state_control/corda_pca.md @@ -0,0 +1,44 @@ +# CorDA-derived PCA + +`CordaPCA` is an activation-steering adaptation of the context-oriented weight +decomposition in [Yang et al., 2024](https://arxiv.org/abs/2406.05223). Its steering +variant follows [steering-lite at 0a064ba](https://github.com/wassname/steering-lite/blob/0a064ba0c23a4998637ff41c5ab0fb5ca50a4271/src/steering_lite/variants/corda_pca.py). + +The CorDA paper uses weights and a calibration dataset to initialize trainable +adapters. Here, the decomposition is used to construct a steering vector instead. +For each target Linear module, the fit pools positive and negative input activations +to form a damped, uncentered second-moment matrix. This matrix and the module's +weights define the CorDA basis. Paired positive-minus-negative differences are +expressed in that basis; their first centered principal component is oriented toward +the mean difference and mapped back to an output-space steering vector. + +Inference adds `strength * direction` to each token's module output. Model weights +stay fixed, and only the resulting vector is needed at inference; the decomposition +is not retained. Supplying `directions` reuses fitted output vectors without fitting +again. The PCA vector has unit norm before reconstruction; the output vector need +not have unit norm. + + + +Numerically zero centered differences raise instead of selecting an arbitrary PCA +direction. This includes constant pair differences. The separate `pca_pairwise` +estimator elsewhere in steerability is not a `CordaPCA` option. + +::: steerability.algorithms.state_control.corda_pca + 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" diff --git a/docs/reference/algorithms/state_control/linear_act.md b/docs/reference/algorithms/state_control/linear_act.md new file mode 100644 index 00000000..02c6bfd3 --- /dev/null +++ b/docs/reference/algorithms/state_control/linear_act.md @@ -0,0 +1,35 @@ +# Linear-AcT + +`LinearAcT` adapts coordinate-wise affine activation transport from +[Rodriguez et al., ICLR 2025](https://openreview.net/forum?id=l2zFn6TIQi), following +[steering-lite at 0a064ba](https://github.com/wassname/steering-lite/blob/0a064ba0c23a4998637ff41c5ab0fb5ca50a4271/src/steering_lite/variants/linear_act.py). +The fit treats negative activations as the source and positive activations as the +target. It sorts samples independently within each coordinate and fits a scalar +slope and bias by least squares. Equal sample counts are required, but pairing is +not used after sorting. This is not a standard-deviation-ratio map. + +Inference applies `h + strength * (slope * h + bias - h)` at each selected decoder +layer's output, for every token. Only the slope and bias vectors are needed; model +weights stay fixed. `affine` accepts an already fitted map. This port implements the +coordinate-wise map, without support masking or sequential layerwise fitting. + + + +::: steerability.algorithms.state_control.linear_act + 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" diff --git a/docs/reference/algorithms/state_control/sspace.md b/docs/reference/algorithms/state_control/sspace.md new file mode 100644 index 00000000..858c2afb --- /dev/null +++ b/docs/reference/algorithms/state_control/sspace.md @@ -0,0 +1,44 @@ +# S-space + +`SSpace` implements a weight-SVD activation-steering variant. The method draws +on [S-Space Steering for Eval-Awareness Control in Reasoning Models](https://apartresearch.com/project/sspace-steering-for-evalawareness-control-in-reasoning-models-7j1i) +by Michael J Clark; this control follows [steering-lite at 0a064ba](https://github.com/wassname/steering-lite/blob/0a064ba0c23a4998637ff41c5ab0fb5ca50a4271/src/steering_lite/variants/sspace.py). + +For each target Linear module, the fit decomposes its weight matrix as +`W = U diag(s) Váµ€`. It expresses positive and negative module outputs in coordinates +`z = (output - bias) U / sqrt(s)` and takes their mean difference. `rank` retains the +coordinates with the largest absolute contrast, not necessarily the largest singular +values. These are weight-scaled coordinates; their activation covariance is not +necessarily identity. + +At inference, the retained basis is used to read each token's coordinates and map +the edit back to the module output. `cosine` scales the edit by the absolute cosine +with the fitted direction; `off` applies a constant edit. `signed` keeps the cosine's +sign: with positive strength it reinforces either pole of the axis, rather than +always pushing toward the positive examples. Reversing a direction leaves `signed` +unchanged, but reverses the `cosine` and `off` edits. + +Unlike CorDA-PCA's fixed output vector, the gated variants need the retained basis +at inference. Model weights stay fixed. For fp16/bf16 outputs, application uses +float32 arithmetic with autocast disabled, then restores the output dtype. + + + +::: steerability.algorithms.state_control.sspace + 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" diff --git a/examples/index.md b/examples/index.md index c4269db8..9f5605c3 100644 --- a/examples/index.md +++ b/examples/index.md @@ -53,6 +53,12 @@ Algorithm notebooks demonstrate how each method (i.e., control) operates. The me :octicons-arrow-right-24: [CAST](./notebooks/algorithms/cast.ipynb) + :octicons-arrow-right-24: [CorDA-derived PCA](./notebooks/algorithms/corda_pca.ipynb) + + :octicons-arrow-right-24: [S-space](./notebooks/algorithms/sspace.ipynb) + + :octicons-arrow-right-24: [Linear-AcT](./notebooks/algorithms/linear_act.ipynb) + :octicons-arrow-right-24: [DirectionalAblation](./notebooks/algorithms/directional_ablation.ipynb) :octicons-arrow-right-24: [ITI](./notebooks/algorithms/iti.ipynb) diff --git a/examples/notebooks/algorithms/corda_pca.ipynb b/examples/notebooks/algorithms/corda_pca.ipynb new file mode 100644 index 00000000..a07b6a5c --- /dev/null +++ b/examples/notebooks/algorithms/corda_pca.ipynb @@ -0,0 +1,179 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "51e1f617", + "metadata": {}, + "source": [ + "# CordaPCA\n", + "\n", + "This is steering-lite's CorDA-derived PCA variant, rather than CorDA fine-tuning. The covariance and paired PCA use Linear inputs.\n", + "\n", + "We use a random CPU Llama and token ids to demonstrate calibration and execution without downloads. The output checks establish a numerical effect, not behavior control. Run with the repository environment (`uv sync --extra all`).\n", + "\n", + "Implementation source: [steering-lite 0a064ba](https://github.com/wassname/steering-lite/tree/0a064ba0c23a4998637ff41c5ab0fb5ca50a4271/src/steering_lite/variants).\n", + "\n", + "Authored by PI/OpenAI." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "bc3c0f8e", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-16T06:48:51.771707Z", + "iopub.status.busy": "2026-09-16T06:48:51.771595Z", + "iopub.status.idle": "2026-09-16T06:48:56.741308Z", + "shell.execute_reply": "2026-09-16T06:48:56.740639Z" + } + }, + "outputs": [], + "source": [ + "from steerability.algorithms.state_control.corda_pca.control import CordaPCA\n", + "import torch\n", + "from tokenizers import Tokenizer, models, pre_tokenizers\n", + "from transformers import LlamaConfig, LlamaForCausalLM, PreTrainedTokenizerFast\n", + "from steerability.algorithms.core.steering_pipeline import SteeringPipeline\n", + "\n", + "torch.manual_seed(12)\n", + "torch.set_num_threads(1)\n", + "model = LlamaForCausalLM(LlamaConfig(\n", + " vocab_size=16, hidden_size=16, intermediate_size=32,\n", + " num_hidden_layers=2, num_attention_heads=2, num_key_value_heads=2,\n", + ")).eval()\n", + "raw = Tokenizer(models.WordLevel({\"\": 0, \"\": 1, \"\": 2, \"yes\": 3, \"no\": 4}, unk_token=\"\"))\n", + "raw.pre_tokenizer = pre_tokenizers.Whitespace()\n", + "tokenizer = PreTrainedTokenizerFast(tokenizer_object=raw, pad_token=\"\", bos_token=\"\", eos_token=\"\")\n", + "positive_ids = torch.tensor([[1, 3, 5], [1, 6, 7], [1, 8, 9], [1, 10, 11]])\n", + "negative_ids = torch.tensor([[1, 4, 6], [1, 7, 8], [1, 9, 10], [1, 11, 12]])\n" + ] + }, + { + "cell_type": "markdown", + "id": "a5724798", + "metadata": {}, + "source": [ + "## Calibration\n", + "\n", + "Each row is one prompt. These equal-length examples have no padding, so position `-1` is the last prompt token. For padded data, gather the last position whose attention mask is one. Capture from the same frozen model that will be steered. CorDA requires paired rows; Linear-AcT requires equal counts. The fit runs during `steer()`, and the inference edit applies at every token, including the prompt." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "903e54fa", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-16T06:48:56.743315Z", + "iopub.status.busy": "2026-09-16T06:48:56.743016Z", + "iopub.status.idle": "2026-09-16T06:48:56.807541Z", + "shell.execute_reply": "2026-09-16T06:48:56.806841Z" + } + }, + "outputs": [], + "source": [ + "target = \"model.layers.0.mlp.down_proj\"\n", + "module = model.get_submodule(target)\n", + "\n", + "def capture_last(ids):\n", + " captured = []\n", + " def capture(module, inputs, output):\n", + " captured.append(inputs[0][:, -1, :].detach().cpu())\n", + " handle = module.register_forward_hook(capture)\n", + " try:\n", + " with torch.no_grad():\n", + " model(input_ids=ids, attention_mask=torch.ones_like(ids))\n", + " finally:\n", + " handle.remove()\n", + " return captured[0]\n", + "\n", + "positive, negative = capture_last(positive_ids), capture_last(negative_ids)" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "18d90609", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-16T06:48:56.809267Z", + "iopub.status.busy": "2026-09-16T06:48:56.809117Z", + "iopub.status.idle": "2026-09-16T06:48:56.811949Z", + "shell.execute_reply": "2026-09-16T06:48:56.811188Z" + } + }, + "outputs": [], + "source": [ + "control = CordaPCA(positive_inputs={target: positive}, negative_inputs={target: negative}, rank=4, strength=0.5)" + ] + }, + { + "cell_type": "markdown", + "id": "bdf615c6", + "metadata": {}, + "source": [ + "## Apply\n", + "\n", + "Scoring the same continuation shows a nonzero change. Scoring again without the control checks that hooks do not remain on the model. For a pretrained model, use contrastive text and evaluate held-out behavior separately." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "d7cd5d91", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-16T06:48:56.813058Z", + "iopub.status.busy": "2026-09-16T06:48:56.812896Z", + "iopub.status.idle": "2026-09-16T06:48:57.352304Z", + "shell.execute_reply": "2026-09-16T06:48:57.351797Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Maximum log-probability change: 0.0005376338958740234\n", + "Generated token ids: [[10, 10, 10]]\n" + ] + } + ], + "source": [ + "baseline = SteeringPipeline(model=model, tokenizer=tokenizer, controls=[])\n", + "pipeline = SteeringPipeline(model=model, tokenizer=tokenizer, controls=[control])\n", + "baseline.steer()\n", + "pipeline.steer()\n", + "query, reference = torch.tensor([[1, 3, 4]]), torch.tensor([[5, 6]])\n", + "base_scores = baseline.compute_logprobs(query, ref_output_ids=reference)\n", + "steered_scores = pipeline.compute_logprobs(query, ref_output_ids=reference)\n", + "assert not torch.allclose(base_scores, steered_scores)\n", + "torch.testing.assert_close(baseline.compute_logprobs(query, ref_output_ids=reference), base_scores)\n", + "print(\"Maximum log-probability change:\", (steered_scores - base_scores).abs().max().item())\n", + "print(\"Generated token ids:\", pipeline.generate(input_ids=query, max_new_tokens=3, do_sample=False).tolist())" + ] + } + ], + "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 +} diff --git a/examples/notebooks/algorithms/linear_act.ipynb b/examples/notebooks/algorithms/linear_act.ipynb new file mode 100644 index 00000000..ed90bd6f --- /dev/null +++ b/examples/notebooks/algorithms/linear_act.ipynb @@ -0,0 +1,173 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0a95d2b4", + "metadata": {}, + "source": [ + "# LinearAcT\n", + "\n", + "[Linear-AcT](https://openreview.net/forum?id=l2zFn6TIQi) sorts each coordinate and fits affine least squares from negative to positive samples. This core map omits support masking and sequential layerwise fitting.\n", + "\n", + "We use a random CPU Llama and token ids to demonstrate calibration and execution without downloads. The output checks establish a numerical effect, not behavior control. Run with the repository environment (`uv sync --extra all`).\n", + "\n", + "Implementation source: [steering-lite 0a064ba](https://github.com/wassname/steering-lite/tree/0a064ba0c23a4998637ff41c5ab0fb5ca50a4271/src/steering_lite/variants).\n", + "\n", + "Authored by PI/OpenAI." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "3373cc9a", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-16T06:49:07.724813Z", + "iopub.status.busy": "2026-09-16T06:49:07.724710Z", + "iopub.status.idle": "2026-09-16T06:49:12.763578Z", + "shell.execute_reply": "2026-09-16T06:49:12.763095Z" + } + }, + "outputs": [], + "source": [ + "from steerability.algorithms.state_control.linear_act.control import LinearAcT\n", + "from steerability.algorithms.core.internals.capture import capture_hidden\n", + "import torch\n", + "from tokenizers import Tokenizer, models, pre_tokenizers\n", + "from transformers import LlamaConfig, LlamaForCausalLM, PreTrainedTokenizerFast\n", + "from steerability.algorithms.core.steering_pipeline import SteeringPipeline\n", + "\n", + "torch.manual_seed(12)\n", + "torch.set_num_threads(1)\n", + "model = LlamaForCausalLM(LlamaConfig(\n", + " vocab_size=16, hidden_size=16, intermediate_size=32,\n", + " num_hidden_layers=2, num_attention_heads=2, num_key_value_heads=2,\n", + ")).eval()\n", + "raw = Tokenizer(models.WordLevel({\"\": 0, \"\": 1, \"\": 2, \"yes\": 3, \"no\": 4}, unk_token=\"\"))\n", + "raw.pre_tokenizer = pre_tokenizers.Whitespace()\n", + "tokenizer = PreTrainedTokenizerFast(tokenizer_object=raw, pad_token=\"\", bos_token=\"\", eos_token=\"\")\n", + "positive_ids = torch.tensor([[1, 3, 5], [1, 6, 7], [1, 8, 9], [1, 10, 11]])\n", + "negative_ids = torch.tensor([[1, 4, 6], [1, 7, 8], [1, 9, 10], [1, 11, 12]])\n" + ] + }, + { + "cell_type": "markdown", + "id": "4d92e619", + "metadata": {}, + "source": [ + "## Calibration\n", + "\n", + "Each row is one prompt. These equal-length examples have no padding, so position `-1` is the last prompt token. For padded data, gather the last position whose attention mask is one. Capture from the same frozen model that will be steered. CorDA requires paired rows; Linear-AcT requires equal counts. The fit runs during `steer()`, and the inference edit applies at every token, including the prompt." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "8c82014c", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-16T06:49:12.765194Z", + "iopub.status.busy": "2026-09-16T06:49:12.764941Z", + "iopub.status.idle": "2026-09-16T06:49:13.065996Z", + "shell.execute_reply": "2026-09-16T06:49:13.065515Z" + } + }, + "outputs": [], + "source": [ + "def capture_last(ids):\n", + " with torch.no_grad():\n", + " hidden, mask = capture_hidden(\n", + " {\"input_ids\": ids, \"attention_mask\": torch.ones_like(ids)}, model=model,\n", + " location=\"layer_output\",\n", + " )\n", + " return hidden[0][:, -1, :]\n", + "\n", + "positive, negative = capture_last(positive_ids), capture_last(negative_ids)" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "9e450cb6", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-16T06:49:13.067705Z", + "iopub.status.busy": "2026-09-16T06:49:13.067556Z", + "iopub.status.idle": "2026-09-16T06:49:13.069872Z", + "shell.execute_reply": "2026-09-16T06:49:13.069471Z" + } + }, + "outputs": [], + "source": [ + "control = LinearAcT(positive_activations={0: positive}, negative_activations={0: negative}, strength=0.5)" + ] + }, + { + "cell_type": "markdown", + "id": "93e86df0", + "metadata": {}, + "source": [ + "## Apply\n", + "\n", + "Scoring the same continuation shows a nonzero change. Scoring again without the control checks that hooks do not remain on the model. For a pretrained model, use contrastive text and evaluate held-out behavior separately." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "37916235", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-16T06:49:13.071110Z", + "iopub.status.busy": "2026-09-16T06:49:13.070994Z", + "iopub.status.idle": "2026-09-16T06:49:13.093813Z", + "shell.execute_reply": "2026-09-16T06:49:13.093371Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Maximum log-probability change: 0.05680513381958008\n", + "Generated token ids: [[0, 0, 0]]\n" + ] + } + ], + "source": [ + "baseline = SteeringPipeline(model=model, tokenizer=tokenizer, controls=[])\n", + "pipeline = SteeringPipeline(model=model, tokenizer=tokenizer, controls=[control])\n", + "baseline.steer()\n", + "pipeline.steer()\n", + "query, reference = torch.tensor([[1, 3, 4]]), torch.tensor([[5, 6]])\n", + "base_scores = baseline.compute_logprobs(query, ref_output_ids=reference)\n", + "steered_scores = pipeline.compute_logprobs(query, ref_output_ids=reference)\n", + "assert not torch.allclose(base_scores, steered_scores)\n", + "torch.testing.assert_close(baseline.compute_logprobs(query, ref_output_ids=reference), base_scores)\n", + "print(\"Maximum log-probability change:\", (steered_scores - base_scores).abs().max().item())\n", + "print(\"Generated token ids:\", pipeline.generate(input_ids=query, max_new_tokens=3, do_sample=False).tolist())" + ] + } + ], + "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 +} diff --git a/examples/notebooks/algorithms/sspace.ipynb b/examples/notebooks/algorithms/sspace.ipynb new file mode 100644 index 00000000..6d147aba --- /dev/null +++ b/examples/notebooks/algorithms/sspace.ipynb @@ -0,0 +1,179 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "50229c87", + "metadata": {}, + "source": [ + "# SSpace\n", + "\n", + "This is wassname's experimental S-space method. Rank keeps the largest absolute S-space mean contrasts. The default absolute cosine gate depends on each token. [Project](https://apartresearch.com/project/sspace-steering-for-evalawareness-control-in-reasoning-models-7j1i).\n", + "\n", + "We use a random CPU Llama and token ids to demonstrate calibration and execution without downloads. The output checks establish a numerical effect, not behavior control. Run with the repository environment (`uv sync --extra all`).\n", + "\n", + "Implementation source: [steering-lite 0a064ba](https://github.com/wassname/steering-lite/tree/0a064ba0c23a4998637ff41c5ab0fb5ca50a4271/src/steering_lite/variants).\n", + "\n", + "Authored by PI/OpenAI." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "4ecd79a7", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-16T06:48:59.587554Z", + "iopub.status.busy": "2026-09-16T06:48:59.587432Z", + "iopub.status.idle": "2026-09-16T06:49:04.226460Z", + "shell.execute_reply": "2026-09-16T06:49:04.225889Z" + } + }, + "outputs": [], + "source": [ + "from steerability.algorithms.state_control.sspace.control import SSpace\n", + "import torch\n", + "from tokenizers import Tokenizer, models, pre_tokenizers\n", + "from transformers import LlamaConfig, LlamaForCausalLM, PreTrainedTokenizerFast\n", + "from steerability.algorithms.core.steering_pipeline import SteeringPipeline\n", + "\n", + "torch.manual_seed(12)\n", + "torch.set_num_threads(1)\n", + "model = LlamaForCausalLM(LlamaConfig(\n", + " vocab_size=16, hidden_size=16, intermediate_size=32,\n", + " num_hidden_layers=2, num_attention_heads=2, num_key_value_heads=2,\n", + ")).eval()\n", + "raw = Tokenizer(models.WordLevel({\"\": 0, \"\": 1, \"\": 2, \"yes\": 3, \"no\": 4}, unk_token=\"\"))\n", + "raw.pre_tokenizer = pre_tokenizers.Whitespace()\n", + "tokenizer = PreTrainedTokenizerFast(tokenizer_object=raw, pad_token=\"\", bos_token=\"\", eos_token=\"\")\n", + "positive_ids = torch.tensor([[1, 3, 5], [1, 6, 7], [1, 8, 9], [1, 10, 11]])\n", + "negative_ids = torch.tensor([[1, 4, 6], [1, 7, 8], [1, 9, 10], [1, 11, 12]])\n" + ] + }, + { + "cell_type": "markdown", + "id": "36cc4f5d", + "metadata": {}, + "source": [ + "## Calibration\n", + "\n", + "Each row is one prompt. These equal-length examples have no padding, so position `-1` is the last prompt token. For padded data, gather the last position whose attention mask is one. Capture from the same frozen model that will be steered. CorDA requires paired rows; Linear-AcT requires equal counts. The fit runs during `steer()`, and the inference edit applies at every token, including the prompt." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "94b9ffaf", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-16T06:49:04.228384Z", + "iopub.status.busy": "2026-09-16T06:49:04.227980Z", + "iopub.status.idle": "2026-09-16T06:49:04.290906Z", + "shell.execute_reply": "2026-09-16T06:49:04.290289Z" + } + }, + "outputs": [], + "source": [ + "target = \"model.layers.0.mlp.down_proj\"\n", + "module = model.get_submodule(target)\n", + "\n", + "def capture_last(ids):\n", + " captured = []\n", + " def capture(module, inputs, output):\n", + " captured.append(output[:, -1, :].detach().cpu())\n", + " handle = module.register_forward_hook(capture)\n", + " try:\n", + " with torch.no_grad():\n", + " model(input_ids=ids, attention_mask=torch.ones_like(ids))\n", + " finally:\n", + " handle.remove()\n", + " return captured[0]\n", + "\n", + "positive, negative = capture_last(positive_ids), capture_last(negative_ids)" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "1b4c9dda", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-16T06:49:04.292224Z", + "iopub.status.busy": "2026-09-16T06:49:04.292064Z", + "iopub.status.idle": "2026-09-16T06:49:04.294663Z", + "shell.execute_reply": "2026-09-16T06:49:04.294208Z" + } + }, + "outputs": [], + "source": [ + "control = SSpace(positive_outputs={target: positive}, negative_outputs={target: negative}, rank=4, strength=0.5)" + ] + }, + { + "cell_type": "markdown", + "id": "f57a9207", + "metadata": {}, + "source": [ + "## Apply\n", + "\n", + "Scoring the same continuation shows a nonzero change. Scoring again without the control checks that hooks do not remain on the model. For a pretrained model, use contrastive text and evaluate held-out behavior separately." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "b01fc826", + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-16T06:49:04.295786Z", + "iopub.status.busy": "2026-09-16T06:49:04.295650Z", + "iopub.status.idle": "2026-09-16T06:49:04.766672Z", + "shell.execute_reply": "2026-09-16T06:49:04.765958Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Maximum log-probability change: 0.05019950866699219\n", + "Generated token ids: [[10, 8, 8]]\n" + ] + } + ], + "source": [ + "baseline = SteeringPipeline(model=model, tokenizer=tokenizer, controls=[])\n", + "pipeline = SteeringPipeline(model=model, tokenizer=tokenizer, controls=[control])\n", + "baseline.steer()\n", + "pipeline.steer()\n", + "query, reference = torch.tensor([[1, 3, 4]]), torch.tensor([[5, 6]])\n", + "base_scores = baseline.compute_logprobs(query, ref_output_ids=reference)\n", + "steered_scores = pipeline.compute_logprobs(query, ref_output_ids=reference)\n", + "assert not torch.allclose(base_scores, steered_scores)\n", + "torch.testing.assert_close(baseline.compute_logprobs(query, ref_output_ids=reference), base_scores)\n", + "print(\"Maximum log-probability change:\", (steered_scores - base_scores).abs().max().item())\n", + "print(\"Generated token ids:\", pipeline.generate(input_ids=query, max_new_tokens=3, do_sample=False).tolist())" + ] + } + ], + "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 +} diff --git a/steerability/algorithms/core/base_control.py b/steerability/algorithms/core/base_control.py index 3533329d..1aa0f2e7 100644 --- a/steerability/algorithms/core/base_control.py +++ b/steerability/algorithms/core/base_control.py @@ -127,6 +127,14 @@ def export_state(self) -> dict[str, Any]: """ return {} + def export_state_classes(self) -> dict[str, str]: + """Artifact class overrides for exported state, keyed by state name. + + Most values infer their class from their Python type. Controls exporting raw tensors + from a calibrated fit override this mapping so frozen bundles retain the fit digest. + """ + return {} + def frozen_form(self, state: dict[str, Any]) -> tuple[str, dict[str, Any]]: """The `(registry method key, constructor kwargs)` of this control's frozen form. diff --git a/steerability/algorithms/core/internals/fingerprint.py b/steerability/algorithms/core/internals/fingerprint.py index 425b61b7..ec33a6e1 100644 --- a/steerability/algorithms/core/internals/fingerprint.py +++ b/steerability/algorithms/core/internals/fingerprint.py @@ -2,6 +2,7 @@ from __future__ import annotations import hashlib +import warnings from typing import TYPE_CHECKING import torch @@ -52,6 +53,24 @@ def model_fingerprint(model: PreTrainedModel) -> str: return digest.hexdigest()[:16] +def verify_model_fingerprint( + recorded: str | None, live: str | None, *, artifact_class: str, policy: str, +) -> None: + """Apply the frozen-artifact weight mismatch policy. Authored by PI/Astra.""" + if policy == "off" or not recorded or not live or recorded == live: + return + message = ( + f"Precomputed {artifact_class} artifact was produced on a different " + f"model (fingerprint {recorded!r} vs {live!r})." + + ("" if artifact_class == "calibrated" + else " Direction artifacts may transfer across fine-tunes of one " + "architecture; verify the behavior.") + ) + if artifact_class == "calibrated" and policy == "strict": + raise ValueError(message) + warnings.warn(message, UserWarning, stacklevel=2) + + def session_artifact_identity(session: SteeringSession | None) -> tuple[str, dict]: """`(model_type, meta)` recorded for a session-fitted artifact, from the session layout. diff --git a/steerability/algorithms/core/steering_pipeline.py b/steerability/algorithms/core/steering_pipeline.py index a729d5c1..4d07a1a2 100644 --- a/steerability/algorithms/core/steering_pipeline.py +++ b/steerability/algorithms/core/steering_pipeline.py @@ -646,6 +646,17 @@ def _run_control_steer(self, control, access: ModelAccess, venue_session, steer_ if "session" not in kwargs: scoped = ScopedSession(venue_session, type(control).__name__, access) kwargs = {**kwargs, "session": scoped} + tensor_records = getattr(control, "_spipe_tensor_records", ()) + if tensor_records and control._spipe_verify != "off": + from steerability.algorithms.core.internals.fingerprint import verify_model_fingerprint + + verification_session = venue_session if kwargs["session"] is None else kwargs["session"] + live_fingerprint = verification_session.layout.model_fingerprint + for record in tensor_records: + verify_model_fingerprint( + record.provenance.get("model_fingerprint"), live_fingerprint, + artifact_class=record.artifact_class, policy=control._spipe_verify, + ) model = self.model if access >= ModelAccess.MODULE else None control_name = type(control).__name__ logger.info("Steering %s (access=%s).", control_name, access.name.lower()) diff --git a/steerability/algorithms/state_control/common/sources.py b/steerability/algorithms/state_control/common/sources.py index dc4ab9a0..84cd8de0 100644 --- a/steerability/algorithms/state_control/common/sources.py +++ b/steerability/algorithms/state_control/common/sources.py @@ -20,6 +20,7 @@ from steerability.algorithms.core.execution.access import ModelAccess from steerability.algorithms.core.internals.capture import HiddenStateLocation from steerability.algorithms.core.internals.data import ContrastivePairs, as_contrastive_pairs +from steerability.algorithms.core.internals.fingerprint import verify_model_fingerprint from steerability.algorithms.state_control.common.estimators import ( ContrastiveDirectionEstimator, MeanDifferenceEstimator, @@ -290,17 +291,10 @@ def _check(self, layout) -> None: hard=True, ) - recorded = self.provenance.get("model_fingerprint") - live = getattr(layout, "model_fingerprint", None) - if recorded and live and recorded != live: - self._report( - f"Precomputed {self.artifact_class} artifact was produced on a different " - f"model (fingerprint {recorded!r} vs {live!r})." - + ("" if self.artifact_class == "calibrated" - else " Direction artifacts may transfer across fine-tunes of one " - "architecture; verify the behavior."), - hard=(self.artifact_class == "calibrated"), - ) + verify_model_fingerprint( + self.provenance.get("model_fingerprint"), getattr(layout, "model_fingerprint", None), + artifact_class=self.artifact_class, policy=self.policy, + ) def resolve( self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase, *, session: "SteeringSession | None" = None diff --git a/steerability/algorithms/state_control/corda_pca/__init__.py b/steerability/algorithms/state_control/corda_pca/__init__.py new file mode 100644 index 00000000..1aad157d --- /dev/null +++ b/steerability/algorithms/state_control/corda_pca/__init__.py @@ -0,0 +1,7 @@ +"""CorDA-derived PCA registration.""" +from .args import CordaPCAArgs +from .control import CordaPCA + +STEERING_METHOD = { + "category": "state_control", "name": "corda_pca", "control": CordaPCA, "args": CordaPCAArgs, +} diff --git a/steerability/algorithms/state_control/corda_pca/args.py b/steerability/algorithms/state_control/corda_pca/args.py new file mode 100644 index 00000000..22e7051b --- /dev/null +++ b/steerability/algorithms/state_control/corda_pca/args.py @@ -0,0 +1,40 @@ +"""CorDA-derived PCA arguments.""" +from dataclasses import dataclass + +import torch + +from steerability.algorithms.core.base_args import BaseArgs + + +@dataclass +class CordaPCAArgs(BaseArgs): + """Fit from paired last-token Linear inputs, keyed by exact module path. + + Args: + positive_inputs: Paired `[N, in_features]` positive samples per module. + negative_inputs: Negative samples with matching shapes and module paths. + directions: Frozen output offsets, supplied instead of calibration inputs. + rank: Leading singular modes to retain; -1 keeps all modes. + damping: Multiplier of the calibration mean square for covariance damping. + strength: Multiplier of the output offset; zero disables the edit. + """ + + positive_inputs: dict[str, torch.Tensor] | None = None + negative_inputs: dict[str, torch.Tensor] | None = None + directions: dict[str, torch.Tensor] | None = None + rank: int = -1 + damping: float = 0.01 + strength: float = 1.0 + + def __post_init__(self) -> None: + if self.rank != -1 and self.rank < 1: + raise ValueError("rank must be -1 or positive") + if self.damping <= 0: + raise ValueError("damping must be positive") + if self.directions is None: + if not self.positive_inputs or not self.negative_inputs: + raise ValueError("provide paired positive_inputs and negative_inputs") + if self.positive_inputs.keys() != self.negative_inputs.keys(): + raise ValueError("positive and negative module paths must match") + elif not self.directions or self.positive_inputs is not None or self.negative_inputs is not None: + raise ValueError("provide directions or calibration inputs, exclusively") diff --git a/steerability/algorithms/state_control/corda_pca/control.py b/steerability/algorithms/state_control/corda_pca/control.py new file mode 100644 index 00000000..1465b4de --- /dev/null +++ b/steerability/algorithms/state_control/corda_pca/control.py @@ -0,0 +1,118 @@ +"""CorDA-derived PCA steering.""" +import torch + +from steerability.algorithms.core.execution.access import ModelAccess +from steerability.algorithms.state_control.base import HookControl + +from .args import CordaPCAArgs + + +def fit_corda_direction( + weight: torch.Tensor, positive: torch.Tensor, negative: torch.Tensor, + rank: int = -1, damping: float = 0.01, +) -> torch.Tensor: + """Return an output-space offset from paired `[N, in_features]` inputs.""" + weight = weight.detach().float().cpu() + positive = positive.detach().float().cpu() + negative = negative.detach().float().cpu() + if positive.ndim != 2 or positive.shape != negative.shape or positive.shape[0] < 2: + raise ValueError("CorDA PCA needs at least two paired [N, in_features] samples") + if positive.shape[1] != weight.shape[1]: + raise ValueError("calibration input width does not match Linear.in_features") + context = torch.cat((positive, negative)) + lam = context.square().mean() * damping + if not torch.isfinite(lam) or lam <= 0: + raise ValueError("CorDA damping lambda must be finite and positive") + n = context.shape[0] + weighted = (weight @ context.T) @ context / n + lam * weight + u, singular, vh = torch.linalg.svd(weighted, full_matrices=False) + r = min(weight.shape) if rank < 0 else min(rank, min(weight.shape)) + u, sqrt_s, v = u[:, :r], singular[:r].sqrt(), vh[:r].T + b = context / n**0.5 + gram = torch.eye(n) + (b @ b.T) / lam + v_corda = v / lam - b.T @ torch.linalg.solve(gram, b @ v) / lam.square() + differences = ((positive - negative) @ v_corda) * sqrt_s + centered = differences - differences.mean(0) + _, variation, pcs = torch.linalg.svd(centered, full_matrices=False) + tolerance = torch.finfo(differences.dtype).eps * max(differences.shape) * differences.norm() + if variation[0] <= tolerance: + raise ValueError("CorDA centered paired differences need numerically nonzero variance") + direction = pcs[0] + direction = direction * torch.sign(differences.mean(0) @ direction + 1e-8) + return ((direction * sqrt_s) @ u.T).contiguous() + + +class CordaPCA(HookControl): + """Add a fixed output offset fitted from paired Linear inputs. + + The weights and pooled positive/negative input second moments define a CorDA + basis. Centered PCA on paired differences in that basis produces a direction, + which is mapped back to a fixed output vector. Inference adds the vector to + every token; it needs no decomposition, and model weights stay fixed. + `directions` accepts already fitted output vectors instead of calibration inputs. + + Reference: + + - Yang et al., "CorDA: Context-Oriented Decomposition Adaptation of Large Language Models" (2024). + https://arxiv.org/abs/2406.05223 + This control adapts the decomposition, not the paper's fine-tuning procedure. + - Implementation: wassname, steering-lite `corda_pca.py` at `0a064ba`. + https://github.com/wassname/steering-lite/blob/0a064ba0c23a4998637ff41c5ab0fb5ca50a4271/src/steering_lite/variants/corda_pca.py + """ + + Args = CordaPCAArgs + supports_batching = True + + def steer_access(self) -> ModelAccess: + return ModelAccess.MODULE + + def steer(self, model: torch.nn.Module, tokenizer=None, **kwargs) -> None: + names = self.positive_inputs if self.directions is None else self.directions + self.fitted_directions = {} + for name in names: + module = model.get_submodule(name) + if not isinstance(module, torch.nn.Linear): + raise TypeError(f"{name}: expected torch.nn.Linear with [out, in] weight") + if self.directions is None: + direction = fit_corda_direction( + module.weight, self.positive_inputs[name], self.negative_inputs[name], + self.rank, self.damping, + ) + else: + direction = self.directions[name].detach().cpu().clone() + if direction.shape != (module.out_features,) or not torch.isfinite(direction).all(): + raise ValueError(f"{name}: direction must be finite [out_features]") + self.fitted_directions[name] = direction + + def get_hooks(self, input_ids: torch.Tensor, runtime_kwargs: dict | None = None, **kwargs) -> dict: + hooks = [] + for name, direction in self.fitted_directions.items(): + runtime = {} + + def hook(module, inputs, kwargs, output, direction=direction, runtime=runtime): + if self.strength == 0: + return output + key = (output.device, output.dtype, self.strength) + if key not in runtime: + runtime[key] = self.strength * direction.to(output) + return output + runtime[key] + + hooks.append({"module": name, "hook_func": hook}) + return {"pre": [], "forward": hooks, "backward": []} + + def steer_fits(self) -> tuple[tuple[str, str], ...]: + return () if self.directions is not None else (("corda_pca", "direction"),) + + def export_state(self) -> dict[str, torch.Tensor]: + return self.fitted_directions + + def export_state_classes(self) -> dict[str, str]: + return {name: "direction" for name in self.export_state()} + + def frozen_form(self, state: dict[str, torch.Tensor]) -> tuple[str, dict]: + return "state_control/corda_pca", {"directions": state, "strength": self.strength} + + def fit_identity(self) -> tuple | None: + if self.directions is not None: + return None + return (self.positive_inputs, self.negative_inputs, self.rank, self.damping) diff --git a/steerability/algorithms/state_control/linear_act/__init__.py b/steerability/algorithms/state_control/linear_act/__init__.py new file mode 100644 index 00000000..1ce3bf85 --- /dev/null +++ b/steerability/algorithms/state_control/linear_act/__init__.py @@ -0,0 +1,7 @@ +"""Linear-AcT registration.""" +from .args import LinearAcTArgs +from .control import LinearAcT + +STEERING_METHOD = { + "category": "state_control", "name": "linear_act", "control": LinearAcT, "args": LinearAcTArgs, +} diff --git a/steerability/algorithms/state_control/linear_act/args.py b/steerability/algorithms/state_control/linear_act/args.py new file mode 100644 index 00000000..c6b77562 --- /dev/null +++ b/steerability/algorithms/state_control/linear_act/args.py @@ -0,0 +1,37 @@ +"""Linear-AcT arguments.""" +from dataclasses import dataclass + +import torch + +from steerability.algorithms.core.base_args import BaseArgs + + +@dataclass +class LinearAcTArgs(BaseArgs): + """Last-token decoder outputs `[N, hidden_size]`, keyed by layer index. + + Args: + positive_activations: Target distribution samples per layer. + negative_activations: Source samples with matching shapes and layers. + affine: Frozen `[2, hidden_size]` slope/bias tensors instead of samples. + strength: Interpolation from identity (0) to transport (1); extrapolation allowed. + """ + + positive_activations: dict[int, torch.Tensor] | None = None + negative_activations: dict[int, torch.Tensor] | None = None + affine: dict[int, torch.Tensor] | None = None + strength: float = 1.0 + + def __post_init__(self) -> None: + if self.affine is None: + if not self.positive_activations or not self.negative_activations: + raise ValueError("provide positive_activations and negative_activations") + if self.positive_activations.keys() != self.negative_activations.keys(): + raise ValueError("positive and negative layers must match") + keys = self.positive_activations + else: + if not self.affine or self.positive_activations is not None or self.negative_activations is not None: + raise ValueError("provide affine or calibration activations, exclusively") + keys = self.affine + if any(not isinstance(k, int) or k < 0 for k in keys): + raise ValueError("layer indices must be nonnegative integers") diff --git a/steerability/algorithms/state_control/linear_act/control.py b/steerability/algorithms/state_control/linear_act/control.py new file mode 100644 index 00000000..9e7f792d --- /dev/null +++ b/steerability/algorithms/state_control/linear_act/control.py @@ -0,0 +1,123 @@ +"""Coordinate-wise affine activation transport.""" +import torch + +from steerability.algorithms.state_control.base import InterventionControl +from steerability.algorithms.state_control.common.specs import Intervention, TokenScope +from steerability.algorithms.state_control.common.transforms.base import BaseTransform +from steerability.algorithms.state_control.common.transforms.context import TransformContext + +from .args import LinearAcTArgs + + +def fit_linear_act(positive: torch.Tensor, negative: torch.Tensor) -> torch.Tensor: + """Return `[slope, bias]` from coordinate-wise sorted least squares.""" + if positive.ndim != 2 or positive.shape != negative.shape or positive.shape[0] < 2: + raise ValueError("Linear-AcT needs equal [N, hidden_size] shapes with N >= 2") + source = negative.detach().float().cpu().sort(dim=0).values + target = positive.detach().float().cpu().sort(dim=0).values + mean_source = source.mean(0) + mean_target = target.mean(0) + source_centered = source - mean_source + target_centered = target - mean_target + denominator = source_centered.square().sum(0) + if (denominator <= 0).any(): + raise ValueError("Linear-AcT source coordinates must have nonzero variance") + slope = (source_centered * target_centered).sum(0) / denominator + return torch.stack((slope, mean_target - slope * mean_source)) + + +class _LinearAcTTransform(BaseTransform): + def __init__( + self, affine: dict[int, torch.Tensor] | None, + positive: dict[int, torch.Tensor] | None = None, + negative: dict[int, torch.Tensor] | None = None, strength: float = 1.0, + ) -> None: + self.affine = affine + self.positive = positive + self.negative = negative + self.strength = strength + self._validated = False + self._runtime_affine = {} + + @property + def is_bound(self) -> bool: + return self._validated + + @property + def covered_layer_ids(self) -> set[int]: + return set(self.affine if self.affine is not None else self.positive) + + def bind(self, ctx: TransformContext) -> "_LinearAcTTransform": + if self.affine is not None: + affine = {lid: value.detach().cpu().clone() for lid, value in self.affine.items()} + else: + affine = {lid: fit_linear_act(pos, self.negative[lid]) for lid, pos in self.positive.items()} + for lid, value in affine.items(): + if value.shape != (2, ctx.hidden_size) or not torch.isfinite(value).all(): + raise ValueError(f"layer {lid}: affine must be finite [2, hidden_size]") + bound = _LinearAcTTransform(affine, strength=self.strength) + bound._runtime_affine = { + (lid, ctx.device, ctx.dtype): value.to(device=ctx.device, dtype=ctx.dtype) + for lid, value in affine.items() + } + bound._validated = True + return bound + + def apply( + self, hidden_states: torch.Tensor, *, layer_id: int, token_mask: torch.BoolTensor, **kwargs, + ) -> torch.Tensor: + self._require_bound() + if self.strength == 0: + return hidden_states + key = (layer_id, hidden_states.device, hidden_states.dtype) + if key not in self._runtime_affine: + self._runtime_affine[key] = self.affine[layer_id].to(hidden_states) + slope, bias = self._runtime_affine[key] + transported = (1 - self.strength) * hidden_states + self.strength * (hidden_states * slope + bias) + return torch.where(token_mask.unsqueeze(-1), transported, hidden_states) + + +class LinearAcT(InterventionControl): + """Transport decoder-output coordinates with sorted affine least squares. + + Positive activations are the target and negative activations are the source. + Sorting is independent per coordinate, so fitting requires equal sample counts + but does not use row pairing. Inference needs only slope and bias vectors; the + map is interpolated with identity by `strength` and applied to every token. + Model weights stay fixed. Constant source coordinates raise. + + Reference: + + - Rodriguez et al., "Controlling Language and Diffusion Models by Transporting Activations", ICLR 2025. + https://openreview.net/forum?id=l2zFn6TIQi + - Implementation: wassname, steering-lite `linear_act.py` at `0a064ba`. + https://github.com/wassname/steering-lite/blob/0a064ba0c23a4998637ff41c5ab0fb5ca50a4271/src/steering_lite/variants/linear_act.py + """ + + Args = LinearAcTArgs + hook_only_hint = "Linear-AcT affine transport has no wire form; use the huggingface backend" + + def _configure(self) -> None: + transform = _LinearAcTTransform(self.affine, self.positive_activations, self.negative_activations, self.strength) + self._template = (Intervention(layers=tuple(sorted(transform.covered_layer_ids)), transform=transform, + scope=TokenScope("all")),) + + def steer_fits(self) -> tuple[tuple[str, str], ...]: + return () if self.affine is not None else (("linear_act", "calibrated"),) + + def export_state(self) -> dict[str, torch.Tensor]: + if not self.interventions: + return {} + return {str(lid): value for lid, value in self.interventions[0].transform.affine.items()} + + def export_state_classes(self) -> dict[str, str]: + return {name: "calibrated" for name in self.export_state()} + + def frozen_form(self, state: dict[str, torch.Tensor]) -> tuple[str, dict]: + return "state_control/linear_act", {"affine": {int(lid): value for lid, value in state.items()}, + "strength": self.strength} + + def fit_identity(self) -> tuple | None: + if self.affine is not None: + return None + return self.positive_activations, self.negative_activations diff --git a/steerability/algorithms/state_control/sspace/__init__.py b/steerability/algorithms/state_control/sspace/__init__.py new file mode 100644 index 00000000..398e6e2b --- /dev/null +++ b/steerability/algorithms/state_control/sspace/__init__.py @@ -0,0 +1,7 @@ +"""S-space registration.""" +from .args import SSpaceArgs +from .control import SSpace + +STEERING_METHOD = { + "category": "state_control", "name": "sspace", "control": SSpace, "args": SSpaceArgs, +} diff --git a/steerability/algorithms/state_control/sspace/args.py b/steerability/algorithms/state_control/sspace/args.py new file mode 100644 index 00000000..b25bee19 --- /dev/null +++ b/steerability/algorithms/state_control/sspace/args.py @@ -0,0 +1,40 @@ +"""S-space arguments.""" +from dataclasses import dataclass + +import torch + +from steerability.algorithms.core.base_args import BaseArgs + + +@dataclass +class SSpaceArgs(BaseArgs): + """Last-token Linear outputs keyed by module path, or frozen `artifacts`. + + Args: + positive_outputs: `[N, out_features]` positive samples per module. + negative_outputs: Negative samples; sample counts can differ. + artifacts: Frozen bases and directions instead of calibration outputs. + rank: Largest absolute S-space mean contrasts to retain; -1 keeps all. + gate: Absolute cosine (`cosine`), signed cosine (`signed`), or constant (`off`). + strength: Multiplier of the output edit; zero disables it. + """ + + positive_outputs: dict[str, torch.Tensor] | None = None + negative_outputs: dict[str, torch.Tensor] | None = None + artifacts: dict[str, dict[str, torch.Tensor]] | None = None + rank: int = -1 + gate: str = "cosine" + strength: float = 1.0 + + def __post_init__(self) -> None: + if self.rank != -1 and self.rank < 1: + raise ValueError("rank must be -1 or positive") + if self.gate not in ("cosine", "signed", "off"): + raise ValueError("gate must be cosine, signed, or off") + if self.artifacts is None: + if not self.positive_outputs or not self.negative_outputs: + raise ValueError("provide positive_outputs and negative_outputs") + if self.positive_outputs.keys() != self.negative_outputs.keys(): + raise ValueError("positive and negative module paths must match") + elif not self.artifacts or self.positive_outputs is not None or self.negative_outputs is not None: + raise ValueError("provide artifacts or calibration outputs, exclusively") diff --git a/steerability/algorithms/state_control/sspace/control.py b/steerability/algorithms/state_control/sspace/control.py new file mode 100644 index 00000000..0c994ad5 --- /dev/null +++ b/steerability/algorithms/state_control/sspace/control.py @@ -0,0 +1,157 @@ +"""Weight-SVD S-space steering.""" +from collections.abc import Callable + +import torch + +from steerability.algorithms.core.execution.access import ModelAccess +from steerability.algorithms.state_control.base import HookControl + +from .args import SSpaceArgs + + +def fit_sspace( + weight: torch.Tensor, bias: torch.Tensor | None, positive: torch.Tensor, negative: torch.Tensor, + rank: int = -1, +) -> dict[str, torch.Tensor]: + """Fit an S-space artifact from `[N, out_features]` Linear outputs.""" + weight = weight.detach().float().cpu() + positive = positive.detach().float().cpu() + negative = negative.detach().float().cpu() + if (positive.ndim != 2 or negative.ndim != 2 or not len(positive) or not len(negative) + or positive.shape[1] != weight.shape[0] or negative.shape[1] != weight.shape[0]): + raise ValueError("calibration outputs must be nonempty [N, out_features]") + u, s, _ = torch.linalg.svd(weight, full_matrices=False) + sqrt_s = s.sqrt() + if (sqrt_s <= 0).any(): + raise ValueError("S-space requires nonzero singular values") + b = torch.zeros(weight.shape[0]) if bias is None else bias.detach().float().cpu() + pos_s, neg_s = ((positive - b) @ u) / sqrt_s, ((negative - b) @ u) / sqrt_s + contrast = pos_s.mean(0) - neg_s.mean(0) + r = min(weight.shape) if rank < 0 else min(rank, min(weight.shape)) + indices = contrast.abs().topk(r).indices.sort().values + direction = contrast[indices] + direction = direction / (direction.norm() + 1e-8) + return {"u": u[:, indices].contiguous(), "sqrt_s": sqrt_s[indices].contiguous(), + "bias": b.clone(), "directions": direction.unsqueeze(0).contiguous()} + + +def _prepare_sspace_transform( + artifact: dict[str, torch.Tensor], strength: float, gate: str, +) -> Callable[[torch.Tensor], torch.Tensor]: + """Prepare token-independent arithmetic on runtime tensors. Authored by PI/Astra.""" + u, sqrt_s, directions, bias = (artifact[k] for k in ("u", "sqrt_s", "directions", "bias")) + with torch.autocast(u.device.type, enabled=False): + if gate == "off": + offset = strength * (directions.sum(0) * sqrt_s) @ u.T + elif gate in ("cosine", "signed"): + amplitudes = directions.norm(dim=-1) + unit = directions / (amplitudes.unsqueeze(-1) + 1e-8) + else: + raise ValueError("gate must be cosine, signed, or off") + + def apply(output: torch.Tensor) -> torch.Tensor: + with torch.autocast(output.device.type, enabled=False): + values = output.to(u.dtype) + if gate == "off": + return (values + offset).to(output.dtype) + projected = ((values - bias) @ u) / sqrt_s + projected = projected / (projected.norm(dim=-1, keepdim=True) + 1e-8) + cosine = projected @ unit.T + engagement = cosine.abs() if gate == "cosine" else cosine + delta = (engagement * amplitudes) @ unit + return (values + strength * (delta * sqrt_s) @ u.T).to(output.dtype) + + return apply + + +def apply_sspace( + output: torch.Tensor, artifact: dict[str, torch.Tensor], strength: float = 1.0, gate: str = "cosine", +) -> torch.Tensor: + """Apply cosine-scaled direction rows, using float32 for fp16/bf16 inputs.""" + if strength == 0: + return output + dtype = torch.float32 if output.dtype in (torch.float16, torch.bfloat16) else output.dtype + runtime = {k: v.to(device=output.device, dtype=dtype) for k, v in artifact.items()} + return _prepare_sspace_transform(runtime, strength, gate)(output) + + +class SSpace(HookControl): + """Apply weight-SVD steering to named Linear outputs. + + The weight SVD defines scaled coordinates for positive and negative outputs. + The fit retains coordinates with the largest absolute mean contrast, then + normalizes the contrast direction. Inference uses the retained basis to read + each token and reconstruct its edit. `cosine` uses absolute cosine similarity; + `signed` keeps its sign and reinforces either pole at positive strength; + `off` applies a constant edit. Model weights stay fixed. + + Reference: + + - Michael J Clark, "S-Space Steering for Eval-Awareness Control in Reasoning Models", Apart Research project. + https://apartresearch.com/project/sspace-steering-for-evalawareness-control-in-reasoning-models-7j1i + - Implementation: wassname, steering-lite `sspace.py` at `0a064ba`. + https://github.com/wassname/steering-lite/blob/0a064ba0c23a4998637ff41c5ab0fb5ca50a4271/src/steering_lite/variants/sspace.py + """ + + Args = SSpaceArgs + supports_batching = True + + def steer_access(self) -> ModelAccess: + return ModelAccess.MODULE + + def steer(self, model: torch.nn.Module, tokenizer=None, **kwargs) -> None: + names = self.positive_outputs if self.artifacts is None else self.artifacts + self.fitted = {} + for name in names: + module = model.get_submodule(name) + if not isinstance(module, torch.nn.Linear): + raise TypeError(f"{name}: expected torch.nn.Linear with [out, in] weight") + if self.artifacts is None: + artifact = fit_sspace( + module.weight, module.bias, self.positive_outputs[name], self.negative_outputs[name], self.rank, + ) + else: + artifact = {key: value.detach().cpu().clone() for key, value in self.artifacts[name].items()} + u, sqrt_s, directions, bias = (artifact[k] for k in ("u", "sqrt_s", "directions", "bias")) + if (u.ndim != 2 or u.shape[0] != module.out_features or sqrt_s.shape != (u.shape[1],) + or directions.ndim != 2 or directions.shape[1] != u.shape[1] + or bias.shape != (module.out_features,) or (sqrt_s <= 0).any() + or not all(torch.isfinite(v).all() for v in artifact.values())): + raise ValueError(f"{name}: invalid S-space artifact dimensions or values") + self.fitted[name] = artifact + + def get_hooks(self, input_ids: torch.Tensor, runtime_kwargs: dict | None = None, **kwargs) -> dict: + hooks = [] + for name, artifact in self.fitted.items(): + runtime = {} + + def hook(module, inputs, kwargs, output, artifact=artifact, runtime=runtime): + if self.strength == 0: + return output + dtype = torch.float32 if output.dtype in (torch.float16, torch.bfloat16) else output.dtype + key = (output.device, dtype, self.strength, self.gate) + if key not in runtime: + tensors = {k: v.to(device=output.device, dtype=dtype) for k, v in artifact.items()} + runtime[key] = _prepare_sspace_transform(tensors, self.strength, self.gate) + return runtime[key](output) + + hooks.append({"module": name, "hook_func": hook}) + return {"pre": [], "forward": hooks, "backward": []} + + def steer_fits(self) -> tuple[tuple[str, str], ...]: + return () if self.artifacts is not None else (("sspace", "calibrated"),) + + def export_state(self) -> dict[str, torch.Tensor]: + return {f"{name}:{key}": value for name, artifact in self.fitted.items() for key, value in artifact.items()} + + def export_state_classes(self) -> dict[str, str]: + return {name: "calibrated" for name in self.export_state()} + + def frozen_form(self, state: dict[str, torch.Tensor]) -> tuple[str, dict]: + artifacts = {name: {key: state[f"{name}:{key}"] for key in artifact} for name, artifact in self.fitted.items()} + return "state_control/sspace", {"artifacts": artifacts, "strength": self.strength, "gate": self.gate} + + def fit_identity(self) -> tuple | None: + if self.artifacts is not None: + return None + return self.positive_outputs, self.negative_outputs, self.rank diff --git a/steerability/spipe/freeze.py b/steerability/spipe/freeze.py index 8284b98a..8a5e8c45 100644 --- a/steerability/spipe/freeze.py +++ b/steerability/spipe/freeze.py @@ -129,18 +129,23 @@ def _freeze_control(control, ctx: EncodeContext, entry_path: str) -> Any: fit_digest = digest_of(fit_identity) if fit_identity is not None else None fit_source = fits[0][0] if fits else (type(control).__name__ if fit_digest is not None else None) - # encode each exported state value with its metadata installed, giving artifact sidecars - # and manifest records the fit provenance; content-equal values re-encoded later inside - # the frozen args reuse these records (first write wins) + # Encode each exported state value into the shared content store. The stored bytes may + # deduplicate, but every resolved control entry keeps its own fit provenance record. artifacts: dict[str, dict] = {} + state_classes = control.export_state_classes() remaining_fits = list(fits) for name, value in state.items(): - artifact_class = _artifact_class_of(value) + explicit_class = state_classes.get(name) + artifact_class = explicit_class or _artifact_class_of(value) source = None digest = None - matched = next((fit for fit in remaining_fits if fit[1] == artifact_class), None) + if explicit_class is not None: + matched = next((fit for fit in fits if fit[1] == artifact_class), None) + else: + matched = next((fit for fit in remaining_fits if fit[1] == artifact_class), None) + if matched is not None: + remaining_fits.remove(matched) if matched is not None: - remaining_fits.remove(matched) source, digest = matched[0], fit_digest elif fit_digest is not None and not fits: source, digest = fit_source, fit_digest @@ -154,7 +159,9 @@ def _freeze_control(control, ctx: EncodeContext, entry_path: str) -> Any: ids: list[str] = [] _collect_artifact_ids(encoded_value, ids) if ids: - artifacts[name] = ctx.records[ids[0]].manifest_entry() + record = ctx.records[ids[0]].manifest_entry() + record.update(artifact_class=artifact_class, source=source, fit_digest=digest) + artifacts[name] = record forms = control.frozen_form(state) if isinstance(forms, tuple): diff --git a/steerability/spipe/spipe.py b/steerability/spipe/spipe.py index 2af7ed07..57e836c7 100644 --- a/steerability/spipe/spipe.py +++ b/steerability/spipe/spipe.py @@ -552,6 +552,12 @@ def pipeline( if verify != "strict" and item["method"] == "structural_control/load_lora": args["allow_base_mismatch"] = True control = self._instantiate(item["method"], args, ctx) + # raw tensors have no source wrapper to check provenance at bind. -- PI/Astra + control._spipe_tensor_records = tuple( + ArtifactRecord.from_mapping(record) for record in item.get("artifacts", {}).values() + if record["type"] == "Tensor" and record["artifact_class"] in ("direction", "calibrated") + ) + control._spipe_verify = verify control.enabled = entry["enabled"] controls.append(control) else: diff --git a/tests/controls/test_corda_pca.py b/tests/controls/test_corda_pca.py new file mode 100644 index 00000000..10d784c9 --- /dev/null +++ b/tests/controls/test_corda_pca.py @@ -0,0 +1,103 @@ +"""CorDA formula and public pipeline regressions.""" +import pytest +import torch + +from steerability.algorithms.core.steering_pipeline import SteeringPipeline +from steerability.algorithms.state_control.corda_pca.control import CordaPCA, fit_corda_direction +from tests.utils.tiny_models import tiny_llama, wordlevel_tokenizer + +HIDDEN = 16 +MODULE = "model.layers.0.self_attn.o_proj" + + +@pytest.mark.parametrize("shape,rank", [((5, 9), -1), ((9, 5), 3), ((5, 5), 1)]) +def test_dense_covariance_formula(shape, rank): + torch.manual_seed(4) + w = torch.randn(*shape) + pos, neg = torch.randn(12, shape[1]) + 0.6, torch.randn(12, shape[1]) + x = torch.cat((pos, neg)) + covariance = x.T @ x / len(x) + 0.07 * x.square().mean() * torch.eye(shape[1]) + u, s, vh = torch.linalg.svd(w @ covariance, full_matrices=False) + r = min(shape) if rank == -1 else rank + z = (pos - neg) @ torch.linalg.solve(covariance, vh[:r].T) * s[:r].sqrt() + _, _, pc = torch.linalg.svd(z - z.mean(0), full_matrices=False) + v = pc[0] * torch.sign(z.mean(0) @ pc[0] + 1e-8) + expected = (v * s[:r].sqrt()) @ u[:, :r].T + actual = fit_corda_direction(w, pos, neg, rank, 0.07) + torch.testing.assert_close(actual, expected, atol=2e-4, rtol=2e-4) + assert actual.norm() > 0 + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_cached_offset_tracks_strength_without_tokenwise_scaling(dtype, device): + """Check cached scaling through the production hook. Authored by PI/Astra.""" + model = torch.nn.Module() + model.linear = torch.nn.Linear(3, 3) + control = CordaPCA(directions={"linear": torch.tensor([0.1, -0.3, 0.7])}) + control.steer(model) + hook = control.get_hooks(torch.ones(1, 1, dtype=torch.long))["forward"][0]["hook_func"] + output = torch.tensor([[[0.2, -1., 2.]]], dtype=dtype, device=device) + for strength in (0.7, -0.5, 0., 0.7): + control.strength = strength + expected = output + strength * control.fitted_directions["linear"].to(output) + for _ in range(2): + actual = hook(model.linear, (), {}, output) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + if strength == 0: + assert actual is output + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as profile: + hook(model.linear, (), {}, output) + operations = {event.key for event in profile.key_averages()} + assert "aten::add" in operations + assert "aten::mul" not in operations + + +@pytest.mark.parametrize("negative", [ + torch.tensor([[-2., 0., 1.], [0., 1., -1.], [2., -1., 0.]]), + torch.eye(3), + 1024 + torch.eye(3), +]) +@pytest.mark.parametrize("contrast", [0.0, 4.0]) +def test_reject_undefined_centered_pca(contrast, negative): + positive = negative + torch.tensor([0., contrast, 0.]) + with pytest.raises(ValueError, match="centered.*variance"): + fit_corda_direction(torch.diag(torch.tensor([3., 2., 1.])), positive, negative) + + +@pytest.mark.parametrize("scale", [1.0, 1e-4]) +def test_resolvable_variation_is_not_rejected_by_absolute_threshold(scale): + torch.manual_seed(42) + negative = torch.randn(8, 3) + positive = negative + scale * torch.randn(8, 3) + direction = fit_corda_direction(torch.eye(3), positive, negative) + assert torch.isfinite(direction).all() + assert direction.norm() > 0 + + +def test_hook_preserves_flattened_linear_output_shape(): + model = torch.nn.Module() + model.linear = torch.nn.Linear(HIDDEN, HIDDEN, bias=False) + direction = torch.arange(HIDDEN, dtype=torch.float32) + control = CordaPCA(directions={"linear": direction}, strength=0.5) + control.steer(model) + hook = control.get_hooks(torch.ones(1, 1, dtype=torch.long))["forward"][0]["hook_func"] + output = torch.randn(3, HIDDEN) + actual = hook(model.linear, (), {}, output) + assert actual.shape == output.shape + torch.testing.assert_close(actual, output + 0.5 * direction) + + +def test_control_fits_and_generates_without_retaining_hooks(): + torch.manual_seed(0) + model = tiny_llama(hidden=HIDDEN, heads=4) + control = CordaPCA( + positive_inputs={MODULE: torch.randn(6, HIDDEN) + 0.3}, + negative_inputs={MODULE: torch.randn(6, HIDDEN)}, + ) + pipeline = SteeringPipeline(controls=[control], model=model, tokenizer=wordlevel_tokenizer()) + pipeline.steer() + input_ids = torch.tensor([[1, 2, 3]]) + baseline = model(input_ids).logits.detach().clone() + output = pipeline.generate(input_ids=input_ids, max_new_tokens=1, do_sample=False, eos_token_id=None) + assert output.ndim == 2 + torch.testing.assert_close(model(input_ids).logits, baseline) diff --git a/tests/controls/test_fitted_control_artifacts.py b/tests/controls/test_fitted_control_artifacts.py new file mode 100644 index 00000000..a846c173 --- /dev/null +++ b/tests/controls/test_fitted_control_artifacts.py @@ -0,0 +1,218 @@ +"""Fit provenance and runtime tensor reuse for the three ports. Authored by PI/Astra.""" +import copy +import json +import warnings + +import pytest +import torch + +from steerability.algorithms.core.internals.fingerprint import model_fingerprint +from steerability.algorithms.core.steering_pipeline import SteeringPipeline +from steerability.algorithms.state_control.corda_pca.control import CordaPCA +from steerability.algorithms.state_control.linear_act.control import LinearAcT +from steerability.algorithms.state_control.sspace.control import SSpace +from steerability.spipe import SPipe, SpipeStaleError +from tests.utils.tiny_models import tiny_llama, wordlevel_tokenizer + +HIDDEN = 16 +MODULES = tuple(f"model.layers.{i}.self_attn.o_proj" for i in range(2)) + + +@pytest.fixture(params=[CordaPCA, SSpace, LinearAcT], ids=lambda cls: cls.__name__) +def fitted_pipeline(request): + torch.manual_seed(17) + positive = {name: torch.randn(6, HIDDEN) + 0.3 for name in MODULES} + negative = {name: torch.randn(6, HIDDEN) for name in MODULES} + cls = request.param + if cls is CordaPCA: + control = cls(positive_inputs=positive, negative_inputs=negative, rank=2) + elif cls is SSpace: + control = cls(positive_outputs=positive, negative_outputs=negative, rank=2) + else: + control = cls( + positive_activations={i: positive[name] for i, name in enumerate(MODULES)}, + negative_activations={i: negative[name] for i, name in enumerate(MODULES)}, + ) + pipeline = SteeringPipeline( + controls=[control], model=tiny_llama(hidden=HIDDEN, heads=4), tokenizer=wordlevel_tokenizer(), + ) + pipeline.steer() + return pipeline + + +def test_every_fitted_artifact_has_provenance_and_reload_preserves_scores(fitted_pipeline, tmp_path): + pipeline = fitted_pipeline + control = pipeline.state_controls[0] + query, answer = torch.tensor([[1, 2, 3]]), torch.tensor([[4, 5]]) + expected = pipeline.compute_logprobs(query, ref_output_ids=answer) + saved = pipeline.to_spipe(model_ref="tiny-llama").save(tmp_path / "fitted") + loaded = SPipe.load(saved) + records = loaded.manifest["controls"][0]["resolved"]["artifacts"] + assert set(records) == set(control.export_state()) + assert len(records) >= 2 + assert all(record["fit_digest"] for record in records.values()) + assert len({record["fit_digest"] for record in records.values()}) == 1 + assert all(record["source"] == control.steer_fits()[0][0] for record in records.values()) + rebuilt = loaded.pipeline() + rebuilt.model, rebuilt.tokenizer = pipeline.model, pipeline.tokenizer + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + rebuilt.steer() + assert not caught + torch.testing.assert_close(rebuilt.compute_logprobs(query, ref_output_ids=answer), expected) + + second_save = rebuilt.to_spipe(model_ref="tiny-llama").save(tmp_path / "precomputed") + manifest_path = second_save / "spipe.json" + manifest = json.loads(manifest_path.read_text()) + entry = manifest["controls"][0] + assert rebuilt.state_controls[0].steer_fits() == () + assert rebuilt.state_controls[0].fit_identity() is None + assert all(record["fit_digest"] is None and record["source"] is None + for record in entry["resolved"]["artifacts"].values()) + entry["args"]["strength"] = 0.7 + if isinstance(control, (CordaPCA, SSpace)): + entry["args"]["rank"] = 3 + if isinstance(control, CordaPCA): + entry["args"]["damping"] = 0.5 + manifest_path.write_text(json.dumps(manifest)) + SPipe.load(second_save) + + +@pytest.mark.parametrize("policy", ["strict", "warn", "off"]) +def test_frozen_tensors_check_changed_weights_before_installing(fitted_pipeline, tmp_path, policy): + saved = fitted_pipeline.to_spipe(model_ref="tiny-llama").save(tmp_path / "fitted") + other_model = copy.deepcopy(fitted_pipeline.model) + with torch.no_grad(): + next(other_model.parameters()).add_(0.01) + assert model_fingerprint(other_model) != model_fingerprint(fitted_pipeline.model) + rebuilt = SPipe.load(saved).pipeline( + model=other_model, tokenizer=fitted_pipeline.tokenizer, verify=policy, + ) + control = rebuilt.state_controls[0] + if policy == "strict" and not isinstance(control, CordaPCA): + with pytest.raises(ValueError, match="Precomputed calibrated artifact.*different model"): + rebuilt.steer() + assert not rebuilt._is_steered + if isinstance(control, SSpace): + assert not hasattr(control, "fitted") + else: + assert not control.interventions + elif policy == "off": + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + rebuilt.steer() + assert not caught + else: + with pytest.warns(UserWarning, match="Precomputed .* artifact.*different model"): + rebuilt.steer() + + +@pytest.mark.parametrize("policy", ["strict", "warn", "off"]) +@pytest.mark.parametrize("provenance", ["missing", "null"]) +def test_frozen_tensors_allow_unrecorded_identity(fitted_pipeline, tmp_path, policy, provenance): + saved = fitted_pipeline.to_spipe(model_ref="tiny-llama").save(tmp_path / "fitted") + manifest_path = saved / "spipe.json" + manifest = json.loads(manifest_path.read_text()) + for record in manifest["controls"][0]["resolved"]["artifacts"].values(): + if provenance == "missing": + del record["provenance"] + else: + record["provenance"] = None + manifest_path.write_text(json.dumps(manifest)) + rebuilt = SPipe.load(saved).pipeline( + model=fitted_pipeline.model, tokenizer=fitted_pipeline.tokenizer, verify=policy, + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + rebuilt.steer() + assert not caught + assert all(not record.provenance for record in rebuilt.state_controls[0]._spipe_tensor_records) + + +@pytest.mark.parametrize("fitted_pipeline", [CordaPCA, SSpace], indirect=True) +@pytest.mark.parametrize("changed_weights", [False, True]) +def test_frozen_hook_control_checks_model_with_explicit_none_session(fitted_pipeline, tmp_path, changed_weights): + saved = fitted_pipeline.to_spipe(model_ref="tiny-llama").save(tmp_path / "fitted") + other_model = copy.deepcopy(fitted_pipeline.model) + if changed_weights: + with torch.no_grad(): + next(other_model.parameters()).add_(0.01) + rebuilt = SPipe.load(saved).pipeline(model=other_model, tokenizer=fitted_pipeline.tokenizer) + if changed_weights and isinstance(rebuilt.state_controls[0], SSpace): + with pytest.raises(ValueError, match="Precomputed calibrated artifact.*different model"): + rebuilt.steer(session=None) + elif changed_weights: + with pytest.warns(UserWarning, match="Precomputed direction artifact.*different model"): + rebuilt.steer(session=None) + else: + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + rebuilt.steer(session=None) + assert not caught + + +def test_recipe_refit_does_not_check_discarded_frozen_weights(fitted_pipeline, tmp_path): + saved = fitted_pipeline.to_spipe(model_ref="tiny-llama").save(tmp_path / "fitted") + other_model = copy.deepcopy(fitted_pipeline.model) + with torch.no_grad(): + next(other_model.parameters()).add_(0.01) + rebuilt = SPipe.load(saved).pipeline( + model=other_model, tokenizer=fitted_pipeline.tokenizer, prefer="recipe", + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + rebuilt.steer() + assert not caught + + +def test_changed_fit_recipe_is_stale(fitted_pipeline, tmp_path): + saved = fitted_pipeline.to_spipe(model_ref="tiny-llama").save(tmp_path / "edited") + manifest_path = saved / "spipe.json" + manifest = json.loads(manifest_path.read_text()) + args = manifest["controls"][0]["args"] + if isinstance(fitted_pipeline.state_controls[0], LinearAcT): + args["positive_activations"]["$map"][1][1] = args["negative_activations"]["$map"][1][1] + else: + args["rank"] = 3 + manifest_path.write_text(json.dumps(manifest)) + with pytest.raises(SpipeStaleError, match="fit digest"): + SPipe.load(saved) + + +def test_runtime_reuses_converted_artifacts_without_changing_cpu_export(fitted_pipeline, monkeypatch, device): + control = fitted_pipeline.state_controls[0] + exported = control.export_state() + snapshots = {name: tensor.clone() for name, tensor in exported.items()} + originals = {id(tensor) for tensor in exported.values()} + conversions = [] + tensor_to = torch.Tensor.to + + def track_to(tensor, *args, **kwargs): + result = tensor_to(tensor, *args, **kwargs) + if id(tensor) in originals and result is not tensor: + conversions.append(id(tensor)) + return result + + monkeypatch.setattr(torch.Tensor, "to", track_to) + inputs = torch.ones(1, 3, dtype=torch.long, device=device) + hooks = control.get_hooks(inputs, model=fitted_pipeline.model)["forward"] + for dtype in (torch.float64, torch.float32, torch.float64): + output = torch.randn(1, 3, HIDDEN, dtype=dtype, device=device) + for index in range(2): + if isinstance(control, LinearAcT): + transform = control.interventions[0].transform + apply = lambda: transform.apply(output, layer_id=index, token_mask=inputs.bool()) + else: + entry = hooks[index] + module = fitted_pipeline.model.get_submodule(entry["module"]) + apply = lambda: entry["hook_func"](module, (), {}, output) + first = apply() + count = len(conversions) + second = apply() + assert len(conversions) == count + assert first.dtype == dtype + torch.testing.assert_close(first, second) + assert len(conversions) == len(exported) * (1 if device.type == "cpu" else 2) + for name, tensor in control.export_state().items(): + assert tensor.device.type == "cpu" + torch.testing.assert_close(tensor, snapshots[name], rtol=0, atol=0) diff --git a/tests/controls/test_linear_act.py b/tests/controls/test_linear_act.py new file mode 100644 index 00000000..ed1e0077 --- /dev/null +++ b/tests/controls/test_linear_act.py @@ -0,0 +1,148 @@ +"""Linear-AcT formula, scope and pipeline tests.""" +import json +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from steerability.algorithms.core.steering_pipeline import SteeringPipeline +from steerability.algorithms.state_control.linear_act.control import LinearAcT, _LinearAcTTransform, fit_linear_act +from steerability.spipe import SPipe, SpipeStaleError +from tests.utils.tiny_models import tiny_llama, wordlevel_tokenizer + +HIDDEN = 16 + + +def test_sorted_least_squares_not_gaussian_moment_map(): + negative = torch.tensor([[4., -1.], [0., 3.], [2., 0.], [1., 2.]]) + positive = torch.tensor([[2., 7.], [20., 2.], [1., 1.], [5., 11.]]) + actual = fit_linear_act(positive, negative) + expected = np.stack([np.linalg.lstsq( + np.column_stack((np.sort(negative[:, j].numpy()), np.ones(4))), + np.sort(positive[:, j].numpy()), rcond=None, + )[0] for j in range(2)], axis=1) + torch.testing.assert_close(actual, torch.from_numpy(expected).float()) + assert not torch.allclose(actual[0], positive.std(0) / negative.std(0), atol=1e-3) + torch.testing.assert_close(fit_linear_act(positive.flip(0), negative.roll(1, 0)), actual) + + +@pytest.mark.parametrize("strength", [0., 0.4, 1., -0.5, 1.5]) +def test_affine_strength_and_mask(strength): + affine = torch.tensor([[2., 0.5], [3., -2.]]) + transform = _LinearAcTTransform({0: affine}, strength=strength).bind( + SimpleNamespace(hidden_size=2, device=torch.device("cpu"), dtype=torch.float32), + ) + y = torch.tensor([[[1., 2.], [3., 4.]]]) + mask = torch.tensor([[True, False]]) + out = transform.apply(y, layer_id=0, token_mask=mask) + expected = y.clone() + expected[:, 0] = (1 - strength) * y[:, 0] + strength * (y[:, 0] * affine[0] + affine[1]) + torch.testing.assert_close(out, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("affine", [ + torch.ones(2, 1), + torch.tensor([[float("nan")] * HIDDEN, [0.] * HIDDEN]), +]) +def test_supplied_affine_is_validated_before_generation(affine): + control = LinearAcT(affine={0: affine}) + pipeline = SteeringPipeline(controls=[control], model=tiny_llama(hidden=HIDDEN, heads=4)) + with pytest.raises(ValueError, match="affine must be finite"): + pipeline.steer() + + +def test_control_binds_and_generates_with_calibrated_affine(): + torch.manual_seed(0) + model = tiny_llama(hidden=HIDDEN, heads=4) + control = LinearAcT( + positive_activations={0: torch.randn(5, HIDDEN) + 0.3}, + negative_activations={0: torch.randn(5, HIDDEN)}, + ) + pipeline = SteeringPipeline(controls=[control], model=model, tokenizer=wordlevel_tokenizer()) + pipeline.steer() + output = pipeline.generate( + input_ids=torch.tensor([[1, 2, 3]]), max_new_tokens=1, do_sample=False, eos_token_id=None, + ) + assert output.ndim == 2 + assert control.interventions[0].transform.is_bound + + +def test_supplied_affine_is_cloned_at_steer_time(): + affine = {0: torch.tensor([[2.] * HIDDEN, [3.] * HIDDEN])} + control = LinearAcT(affine=affine) + pipeline = SteeringPipeline( + controls=[control], model=tiny_llama(hidden=HIDDEN, heads=4), tokenizer=wordlevel_tokenizer(), + ) + pipeline.steer() + affine[0].fill_(float("nan")) + assert torch.isfinite(control.interventions[0].transform.affine[0]).all() + + +def test_precomputed_affine_freezes_without_loading_or_steering_model(tmp_path): + """An unbound precomputed control is already a complete recipe. Authored by PI/Astra.""" + affine = {0: torch.stack((torch.full((HIDDEN,), 1.5), torch.full((HIDDEN,), 0.2)))} + control = LinearAcT(affine=affine, strength=0.4) + pipeline = SteeringPipeline(model_name_or_path="tiny-llama", controls=[control]) + saved = pipeline.to_spipe(freeze=True).save(tmp_path / "precomputed.spipe") + assert pipeline.model is None + assert not pipeline._is_steered + assert control.export_state() == {} + loaded = SPipe.load(saved) + assert loaded.manifest["controls"][0]["resolved"] is None + assert loaded.manifest["lock"]["model_fingerprint"] is None + model, tokenizer = tiny_llama(hidden=HIDDEN, heads=4), wordlevel_tokenizer() + rebuilt = loaded.pipeline(model=model, tokenizer=tokenizer) + rebuilt.steer() + torch.testing.assert_close(rebuilt.state_controls[0].export_state()["0"], affine[0]) + pipeline.model, pipeline.tokenizer = model, tokenizer + pipeline.steer() + query, answer = torch.tensor([[1, 2, 3]]), torch.tensor([[4, 5]]) + torch.testing.assert_close( + rebuilt.compute_logprobs(query, ref_output_ids=answer), + pipeline.compute_logprobs(query, ref_output_ids=answer), + ) + + +def test_frozen_calibrated_affine_records_fit_digest(): + torch.manual_seed(0) + control = LinearAcT( + positive_activations={0: torch.randn(5, HIDDEN) + 0.3}, + negative_activations={0: torch.randn(5, HIDDEN)}, + ) + pipeline = SteeringPipeline( + controls=[control], model=tiny_llama(hidden=HIDDEN, heads=4), tokenizer=wordlevel_tokenizer(), + ) + pipeline.steer() + entry = pipeline.to_spipe(model_ref="tiny-llama").manifest["controls"][0]["resolved"] + assert all(record["artifact_class"] == "calibrated" for record in entry["artifacts"].values()) + assert all(record["fit_digest"] for record in entry["artifacts"].values()) + + +def test_frozen_calibrated_affine_rejects_changed_calibration(tmp_path): + torch.manual_seed(0) + control = LinearAcT( + positive_activations={0: torch.randn(5, HIDDEN) + 0.3}, + negative_activations={0: torch.randn(5, HIDDEN)}, + ) + pipeline = SteeringPipeline( + controls=[control], model=tiny_llama(hidden=HIDDEN, heads=4), tokenizer=wordlevel_tokenizer(), + ) + pipeline.steer() + saved = pipeline.to_spipe(model_ref="tiny-llama").save(tmp_path / "linear-act") + manifest_path = saved / "spipe.json" + manifest = json.loads(manifest_path.read_text()) + args = manifest["controls"][0]["args"] + args["positive_activations"]["$map"][0][1] = args["negative_activations"]["$map"][0][1] + manifest_path.write_text(json.dumps(manifest)) + with pytest.raises(SpipeStaleError, match="fit digest"): + SPipe.load(saved) + + +def test_degenerate_and_mismatched_samples_raise(): + with pytest.raises(ValueError, match="variance"): + fit_linear_act(torch.randn(3, 2), torch.ones(3, 2)) + with pytest.raises(ValueError, match="equal"): + fit_linear_act(torch.randn(3, 2), torch.randn(2, 2)) + with pytest.raises(ValueError, match="N >= 2"): + fit_linear_act(torch.randn(1, 2), torch.randn(1, 2)) diff --git a/tests/controls/test_sspace.py b/tests/controls/test_sspace.py new file mode 100644 index 00000000..23b8c0e4 --- /dev/null +++ b/tests/controls/test_sspace.py @@ -0,0 +1,207 @@ +"""S-space regressions through Linear modules and pipelines.""" +import json + +import pytest +import torch + +from steerability.algorithms.core.steering_pipeline import SteeringPipeline +from steerability.algorithms.state_control.sspace.control import SSpace, apply_sspace, fit_sspace +from steerability.spipe import SPipe, SpipeStaleError +from tests.utils.tiny_models import tiny_llama, wordlevel_tokenizer + +HIDDEN = 16 +MODULE = "model.layers.0.self_attn.o_proj" + + +@pytest.mark.parametrize("gate", ["cosine", "signed", "off"]) +def test_contrast_rank_bias_and_independent_gate_formula(gate): + weight = torch.diag(torch.tensor([9., 4., 1.])) + bias = torch.tensor([10., -7., 3.]) + negative = bias + torch.tensor([[1., 2., 0.], [-1., -2., 0.]]) + positive = negative + torch.tensor([0.03, 0.2, 2.]) + artifact = fit_sspace(weight, bias, positive, negative, rank=1) + torch.testing.assert_close(artifact["u"].abs(), torch.tensor([[0.], [0.], [1.]])) + y = bias + torch.tensor([[[2., 1., 4.], [0., 0., -3.], [0., 0., 0.]]]) + before = y.clone() + out = apply_sspace(y, artifact, 0.7, gate) + z = (y[..., 2] - bias[2]).unsqueeze(-1) + d = artifact["directions"] + unit = d / (d.norm(dim=-1, keepdim=True) + 1e-8) + cosine = (z / (z.norm(dim=-1, keepdim=True) + 1e-8)) @ unit.T + g = torch.ones_like(cosine) if gate == "off" else cosine.abs() if gate == "cosine" else cosine + expected = y.clone() + expected[..., 2] += (0.7 * g * d.norm() * unit.squeeze()).squeeze(-1) + torch.testing.assert_close(out, expected) + torch.testing.assert_close(y, before, rtol=0, atol=0) + torch.testing.assert_close(apply_sspace(y, artifact, 0, gate), y, rtol=0, atol=0) + + +@pytest.mark.parametrize("shape", [(4, 7), (7, 4)]) +def test_rectangular_weight_output_projection(shape): + torch.manual_seed(31) + weight = torch.randn(*shape) + bias = torch.randn(shape[0]) + pos_x, neg_x = torch.randn(9, shape[1]) + 1, torch.randn(9, shape[1]) + pos, neg = pos_x @ weight.T + bias, neg_x @ weight.T + bias + artifact = fit_sspace(weight, bias, pos, neg, rank=2) + u, s, vh = torch.linalg.svd(weight, full_matrices=False) + contrast = (pos_x @ vh.T * s.sqrt()).mean(0) - (neg_x @ vh.T * s.sqrt()).mean(0) + indices = contrast.abs().topk(2).indices.sort().values + direction = contrast[indices] / (contrast[indices].norm() + 1e-8) + torch.testing.assert_close(artifact["directions"][0], direction) + x = torch.randn(2, 3, shape[1]) + y = x @ weight.T + bias + z = x @ vh[indices].T * s[indices].sqrt() + unit = direction / (direction.norm() + 1e-8) + gate = ((z / (z.norm(dim=-1, keepdim=True) + 1e-8)) @ unit).abs() + expected = y + ((gate.unsqueeze(-1) * direction.norm() * unit) * s[indices].sqrt()) @ u[:, indices].T + torch.testing.assert_close(apply_sspace(y, artifact), expected, atol=2e-5, rtol=2e-5) + + +def test_stacked_directions_keep_independent_gates(): + artifact = {"u": torch.eye(2), "sqrt_s": torch.ones(2), "bias": torch.zeros(2), + "directions": torch.tensor([[2., 0.], [0., 3.]])} + y = torch.tensor([[[1., 0.], [0., -1.]]]) + expected = torch.tensor([[[3., 0.], [0., 2.]]]) + torch.testing.assert_close(apply_sspace(y, artifact), expected) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("autocast_dtype", [None, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("gate", ["cosine", "signed", "off"]) +def test_dtype_and_autocast_preserve_finite_formula(dtype, autocast_dtype, gate): + artifact = {"u": torch.eye(2), "sqrt_s": torch.ones(2), "bias": torch.zeros(2), + "directions": torch.tensor([[2., 0.], [0., 3.], [0., 0.]])} + output = torch.tensor([[[0., 0.], [2., -1.]]], dtype=dtype) + before = output.clone() + if gate == "off": + expected = output.double() + torch.tensor([2., 3.]) + else: + cosine = output.double() / (output.double().norm(dim=-1, keepdim=True) + 1e-8) + multiplier = cosine.abs() if gate == "cosine" else cosine + expected = output.double() + multiplier * torch.tensor([2., 3.]) + with torch.autocast("cpu", dtype=autocast_dtype, enabled=autocast_dtype is not None): + actual = apply_sspace(output, artifact, gate=gate) + assert apply_sspace(output, artifact, strength=0, gate=gate) is output + zero_artifact = dict(artifact, directions=torch.zeros(1, 2)) + torch.testing.assert_close(apply_sspace(output, zero_artifact, gate=gate), output) + assert actual.dtype == dtype + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, expected.to(dtype)) + assert not torch.equal(actual, output) + torch.testing.assert_close(output, before, rtol=0, atol=0) + + +def test_fp16_retains_small_basis_scales(): + artifact = {"u": torch.eye(2), "sqrt_s": torch.tensor([1e-8, 1.]), "bias": torch.zeros(2), + "directions": torch.tensor([[1., 1.]])} + output = torch.tensor([[[0., 1.]]], dtype=torch.float16) + expected = torch.tensor([[[0., 1. + 2**-0.5]]], dtype=output.dtype) + torch.testing.assert_close(apply_sspace(output, artifact), expected) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_cached_hook_preserves_small_basis_scales_under_autocast(dtype, device): + artifact = {"u": torch.eye(2), "sqrt_s": torch.tensor([1e-8, 1.]), "bias": torch.zeros(2), + "directions": torch.tensor([[1., 1.]])} + model = torch.nn.Module() + model.linear = torch.nn.Linear(2, 2) + control = SSpace(artifacts={"linear": artifact}) + control.steer(model) + hook = control.get_hooks(torch.ones(1, 1, dtype=torch.long))["forward"][0]["hook_func"] + output = torch.tensor([[[0., 1.]]], dtype=dtype, device=device) + expected = torch.tensor([[[0., 1. + 2**-0.5]]], dtype=dtype, device=device) + with torch.autocast(device.type, dtype=dtype): + for _ in range(2): + actual = hook(model.linear, (), {}, output) + assert actual.dtype == dtype + torch.testing.assert_close(actual, expected) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_cached_hook_reuses_constants_and_tracks_application_settings(dtype, device): + """Only token-dependent norms remain after preparing a hook. Authored by PI/Astra.""" + artifact = {"u": torch.tensor([[0.6, -0.8], [0.8, 0.6]]), + "sqrt_s": torch.tensor([1e-8, 2.]), "bias": torch.tensor([0.2, -0.3]), + "directions": torch.tensor([[2., 0.], [0., 3.], [0., 0.]])} + model = torch.nn.Module() + model.linear = torch.nn.Linear(2, 2) + control = SSpace(artifacts={"linear": artifact}) + control.steer(model) + hook = control.get_hooks(torch.ones(1, 1, dtype=torch.long))["forward"][0]["hook_func"] + output = torch.tensor([[[0., 1.], [2., -1.]]], dtype=dtype, device=device) + for gate in ("off", "cosine", "signed", "off"): + control.gate = gate + for strength in (0.7, -0.5, 0., 0.7): + control.strength = strength + expected = apply_sspace(output, artifact, strength, gate) + with torch.autocast(device.type, dtype=torch.bfloat16): + for _ in range(2): + actual = hook(model.linear, (), {}, output) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + if strength == 0: + assert actual is output + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as profile: + hook(model.linear, (), {}, output) + operations = {event.key: event.count for event in profile.key_averages()} + assert operations.get("aten::linalg_vector_norm", 0) == (0 if gate == "off" else 1) + assert "aten::sum" not in operations + if gate == "off": + assert "aten::mm" not in operations + + +def test_control_generates_and_frozen_fit_rejects_changed_rank(tmp_path): + torch.manual_seed(0) + model = tiny_llama(hidden=HIDDEN, heads=4) + control = SSpace( + positive_outputs={MODULE: torch.randn(6, HIDDEN) + 0.3}, + negative_outputs={MODULE: torch.randn(6, HIDDEN)}, + rank=1, + ) + pipeline = SteeringPipeline(controls=[control], model=model, tokenizer=wordlevel_tokenizer()) + pipeline.steer() + output = pipeline.generate( + input_ids=torch.tensor([[1, 2, 3]]), max_new_tokens=1, do_sample=False, eos_token_id=None, + ) + assert output.ndim == 2 + + saved = pipeline.to_spipe(model_ref="tiny-llama").save(tmp_path / "sspace") + manifest_path = saved / "spipe.json" + manifest = json.loads(manifest_path.read_text()) + artifacts = manifest["controls"][0]["resolved"]["artifacts"] + assert all(record["artifact_class"] == "calibrated" for record in artifacts.values()) + assert all(record["fit_digest"] for record in artifacts.values()) + manifest["controls"][0]["args"]["rank"] = 2 + manifest_path.write_text(json.dumps(manifest)) + with pytest.raises(SpipeStaleError, match="fit digest"): + SPipe.load(saved) + + +def test_shared_tensor_artifacts_keep_each_fit_digest(tmp_path): + torch.manual_seed(0) + controls = [ + SSpace( + positive_outputs={MODULE: torch.randn(6, HIDDEN) + offset}, + negative_outputs={MODULE: torch.randn(6, HIDDEN)}, + rank=rank, + ) + for offset, rank in ((0.3, 1), (-0.4, 2)) + ] + pipeline = SteeringPipeline( + controls=controls, + model=tiny_llama(hidden=HIDDEN, heads=4), + tokenizer=wordlevel_tokenizer(), + ) + pipeline.steer() + saved = pipeline.to_spipe(model_ref="tiny-llama").save(tmp_path / "two-sspace") + entries = SPipe.load(saved).manifest["controls"] + biases = [entry["resolved"]["artifacts"][f"{MODULE}:bias"] for entry in entries] + assert biases[0]["id"] == biases[1]["id"] + assert biases[0]["fit_digest"] != biases[1]["fit_digest"] + + +def test_reject_zero_singular_values_and_wrong_width(): + with pytest.raises(ValueError, match="singular"): + fit_sspace(torch.zeros(3, 4), None, torch.ones(2, 3), torch.ones(2, 3)) + with pytest.raises(ValueError, match="out_features"): + fit_sspace(torch.ones(3, 4), None, torch.ones(2, 4), torch.ones(2, 4))