wsmdocs
Weight Space Mods

Edit model weights.
Measure the model.

WSM is a Python SDK for editing model weights. Change named tensors, compare evaluations before and after the edit, and save a standalone checkpoint.

Python ≥ 3.12 / PyTorch tensors / wsm 0.1.0 on PyPI

Why WSM?

WSM lets you write a weight edit as a Python program, inspect it before running, and test the resulting model. You can load the saved checkpoint in other tools without WSM.

A model exposes several places to intervene. A prompt changes its input. A retrieval system changes the information supplied with that input. A training procedure changes parameters using an objective and an update rule. With direct weight editing, you specify the parameter changes yourself.

That access matters when your question is about the transformation itself. What happens if you scale a particular tensor slice? Can a proposed update be restricted to a named layer? Did the exported checkpoint preserve the intended edit? These are experimental questions. They need a way to inspect weights, execute a precise operation, and compare the resulting model against an unchanged baseline.

You can already manipulate model parameters with PyTorch. WSM builds on that ability and gives the experiment a common execution path. A developer or coding agent supplies Python. The runner loads the model, measures the baseline, executes the program, measures the candidate, and writes a checkpoint with a record of what happened.

WSM runs and records these experiments. Finding edits that improve model behavior remains a research problem.

What is weight space programming?

Let θ denote a model's parameters and fθ(x) its computation on an input. A weight-editing program produces a candidate parameter set, θ′ = P(θ). The model you evaluate afterward is fθ′. The program specifies a parameter transformation. Evaluation determines what that transformation did to behavior.

We call this weight space programming. The term describes an approach, rather than a single algorithm or a feature unique to WSM. A program might scale values, add a structured update, combine compatible checkpoints, or implement a more involved editing method. In each case, code specifies the parameter changes.

Weight coordinates are not human-readable instructions. A tensor name identifies where a value lives in an implementation; it does not tell you which concept that value represents. A small numerical change can affect many outputs. A short program can be reproducible without being useful, interpretable, or safe to apply to a different model.

How this relates to existing work

Fine-tuning also changes weights, usually through repeated optimization steps. An explicit editing program and a training procedure are therefore related ways to produce new parameters. WSM's current model interface exposes detached tensor views and evaluation, rather than a differentiable training interface. It does not replace a training framework.

LoRA learns low-rank updates while freezing the pretrained weights. ROME investigates targeted changes to factual associations through rank-one parameter updates. Task arithmetic constructs weight differences between a pretrained model and its fine-tuned descendants, then studies operations on those differences.

These approaches motivate concrete questions about how parameter changes affect behavior. They have different assumptions, objectives, and evidence. WSM does not currently ship implementations of those methods, and access to tensors does not reproduce their results.

Prompting and retrieval remain useful when you need to supply instructions or current information. A persistent weight edit changes the model artifact itself. Choosing that route means accepting the responsibility to measure regressions and preserve the original model for comparison.

How do we do it?

Start with a local model snapshot and a hypothesis about a parameter transformation. The example below scales the first 16 elements of a normalization tensor by 0.95. It demonstrates the editing workflow without claiming to improve the model.

Install the published package with its model dependencies in a Python 3.12 or later environment.

python3 -m pip install 'wsm[editing]==0.1.0'

Load, evaluate, modify, evaluate

WeightModel.load loads a local snapshot without downloading model files or executing custom model repository Python. model.weights exposes named tensors. Slicing and in-place arithmetic operate on live model storage, so the next evaluation sees the modified values.

experiment.pyPython SDK
from wsm import WeightModel

model = WeightModel.load("/path/to/snapshot", device="cpu")
cases = [{"prompt": "Say hello.", "target": "Hello!"}]

before = model.evaluate(cases)
model.weights["model.norm.weight"][:16].mul_(0.95)
after = model.evaluate(cases)

print({"baseline": before, "candidate": after})
model.save("/path/to/new-checkpoint")

The tensor name is specific to compatible model architectures. Inspect your model's names and shapes before applying an edit. Tied parameter names can share storage, so changing one view can change another name's values too.

W′ = 0.95 × WIllustrative tensor slice
[wsm_]

Named tensors.
Explicit operations.

0.320 × 0.95 = 0.304
The highlighted value illustrates the arithmetic on one element. The code scales up to 16 elements of model.norm.weight. These values are not an experimental result.

Run the edit as a recorded experiment

For the CLI, put the mutation in a Python file defining edit(model). The runner owns the baseline and candidate evaluations around that function.

edit.pyCLI entry point
def edit(model):
    model.weights["model.norm.weight"][:16].mul_(0.95)

Create a cases file containing a JSON array of prompt and target objects. This one-case example checks the workflow; use cases appropriate to your actual hypothesis.

[{"prompt": "Say hello.", "target": "Hello!"}]

Save it as cases.json, replace the snapshot path, and choose an output directory that does not already exist.

mkdir -p runs
wsm execute edit.py --model /path/to/snapshot \
  --cases cases.json --output runs/edit --device cpu
runs/edit/run.json
Program source hash, execution state, baseline and candidate evaluations, and checkpoint location.
runs/edit/checkpoint/
A standalone model checkpoint with weights and tokenizer files.

Interpret the result before keeping it

Evaluation measures target-token negative log likelihood. WSM calculates a mean for each case, then an unweighted mean across cases. It does not generate a response or apply a chat template. A loss difference answers a question about these target tokens under this evaluation procedure.

The runner records that difference without automatically accepting or rejecting the candidate. Keep development cases separate from held-out checks. Test behavior beyond the examples used to design the edit, including unrelated tasks you want to preserve. The SDK's public cases are visible to the editing program; independent private grading requires a separate process.

For exact API behavior and execution details, read the Python editing guide.

Why this design?

Python expresses the transformation directly

An edit can use familiar control flow, tensor indexing, and arithmetic. We do not require a separate editing language. A coding agent can write the same edit(model) file a developer would write, without the SDK importing an agent framework. This keeps the program inspectable and the execution interface independent of who authored it.

Python files run with the invoking process's privileges. The standalone runner does not sandbox filesystem access, imports, or other Python behavior. The experimental trial system has a separate isolation mechanism; it is not a property of ordinary SDK execution.

Execution and acceptance answer different questions

The runner can determine whether a program executed and whether its checkpoint was published. It cannot infer whether your chosen evaluation covers the behavior you care about. A lower development loss alone does not establish that the model improved on other tasks.

For the same reason, a receipt is a record of an execution, not authenticated proof that the program was harmless or the evaluation scientifically sufficient. More capable search can overfit an exposed metric. Independent grading, regression tasks, and an explicit acceptance policy belong in the surrounding experiment.

The output must survive outside the experiment

An edit is useful to other tools only if they can load the result. Saving writes model weights and tokenizer files, checks finite tensors, reloads the staged checkpoint, and compares its tensors before publication. The destination is published atomically without overwriting an existing checkpoint, on filesystems that support exclusive atomic rename.

Those checks protect the artifact's integrity. They do not certify its model quality. A correctly saved checkpoint can still be a worse model. We keep that distinction visible in the architecture and verification documentation.

What's done?

The following implementation and experiment records describe what WSM can do and what remains unproven.

A published SDK and a shared execution path

Version 0.1.0 is available on PyPI, with a matching GitHub release. The supported public API is WeightModel and execute_program. The CLI and terminal editing action use the same program runner. You can load, edit, evaluate, and save without adopting an agent framework.

Tests for edits and checkpoint publication

The repository's model tests build tiny local models and exercise live tensor changes, evaluation, and checkpoint reloads. Failure cases cover non-finite weights, existing destinations, interruptions, and execution errors. The verification record distinguishes completed runs from tests that require optional dependencies. These checks establish implementation behavior, not general model improvement.

A GPU edit and an independent reload path

The documented Modal demonstration checks a deterministic edit on a pinned SmolLM2 snapshot, then checks the exported model in a separate CPU environment without the original snapshot or model cache. It examines intended tensor changes, unchanged tensors, and evaluation parity after reload. This tests the edit-and-export mechanism.

An experimental teacher and student integration

The retained trial verification record documents a coding-agent rollout, independent held-out grading on a small task, and a response-weighted student training update through Modal Spindle. The record reports lower held-out loss on that toy task and one confirmed student optimizer update. It does not demonstrate improved student editing ability.

These results apply to the recorded code and runtime. They do not establish reliable autonomous model improvement. Finding useful edits, generalizing beyond the evaluated tasks, and repeating gains across models remain open research questions.

Start with the quickstart to run an edit, or read the research limitations before designing an experiment around the results.

Documentation

Type to search the documentation.