Let your coding agent
edit model weights.

Why WSM

WSM gives AI agents a way to experiment with a model's weights and use the results to choose their next edit. Its experimental agent loop runs Python edits, returns evaluation results, and keeps checkpoints and history for the next turn. You can also use the Python SDK directly.

Use WSM to test edits such as scaling a tensor slice or changing one layer, then compare the changed model against an unchanged baseline.

A developer or coding agent supplies a Python editing file. The standalone runner loads the model, measures the baseline, executes the program, measures the candidate, and writes a checkpoint with a record of what happened.

Finding edits that improve model behavior remains a research problem.

Current direction

WSM's goal is to make weight-space editing easier for others to use. The first priority is to make the edit-and-evaluate workflow reproducible, so you can rerun experiments and inspect what changed.

The SDK supports evaluation before and after an edit and checkpoint export. The experimental benchmark workflow adds frozen inputs, retained grades, and offline reports.

Next is a workflow for comparing saved checkpoints and recovering from unwanted edits. This is planned work.

What is weight space programming?

An AI model stores what it learned as numbers called weights. Weight space programming means writing code that changes those numbers, then testing what changed in the model's behavior.

With WSM, you can use Python to select a group of weights and modify them. For example, you can multiply a small group by 0.95, compare evaluations before and after the edit, and save the changed model.

Changing weights does not automatically make a model better. Individual weights do not neatly represent concepts such as a fact or a skill, so even a small edit can affect many outputs. Testing helps you find out whether the change helped or caused regressions.

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. Use a training framework for differentiable training.

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 WSM works

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.

uv venv --python 3.12 .venv
. .venv/bin/activate
uv 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.

Use a local snapshot with its tokenizer files and an unused output path whose parent directory exists.

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 depends on the model architecture. 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.

CLI coming soon

Use the Python SDK to edit and evaluate weights. The command-line interface is coming soon.

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. WSM requires no 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.

A receipt records execution. Assess program safety and evaluation coverage separately. 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. Evaluate the saved model on your tasks to check its quality. Read more about these checks in the architecture documentation.

What's done?

WSM has a published SDK, checkpoint publication tests, and recorded editing and training experiments.

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. 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 repository verification record distinguishes completed runs from tests that require optional dependencies. These checks cover editing and checkpoint publication.

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 editing-agent training workflow

WSM's experimental training integration uses the Tinker API to train the agent that writes weight edits. The target model is separate. WSM changes its weights by running the agent's Python code. Tinker is optional for the Python SDK.

The recorded demonstration used a Codex teacher to write edits, then trained a separate open-weight student on the teacher's recorded actions through Modal Spindle using the Tinker SDK. It confirmed one training update. Better student editing ability remains unproven. The hosted Tinker training path is implemented and awaits verification of a training update. Read the experimental trial guidefor setup and the distinction between the teacher, student, and target.

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.