B RAIN S URGERY: Reproducible and Reliable Declarative Weight Manipulations for Model Editing and Upcycling Gianluca Barmina* , Annemette Broch Pirchert* , Andrea Blasi Núñez Lukas Galke Poech, Peter Schneider-Kamp University of Southern Denmark {gbarmina,ampirchert,petersk,galke}@imada.sdu.dk
arXiv:2606.09707v1 [cs.LG] 8 Jun 2026
Abstract
range of research areas. We briefly describe below the importance and relevance of post-hoc manipulation techniques in four distinct research areas, before we introduce and present our framework that provides the technical tools for facilitating all these use cases through performing principled and validated re-arrangements and edits of the model parameters (metaphorically, a “brain surgery”).
As deep learning models scale, managing, inspecting, and modifying large checkpoints has become increasingly challenging. Researchers often need to alter model weights for layer restructuring, precision casting, low-rank factorization, and architectural debugging, yet these workflows often rely on fragile ad-hoc Python scripts. Here, we introduce B RAIN S URGERY, a tool for robust and reproducible “tensor surgery” on neural network checkpoints, and provide a system demonstration covering four examples and three case studies from model upcycling to LoRA extraction. By abstracting storage formats and memory management, B RAIN S URGERY executes complex transformations through declarative YAML plans. It supports structural modifications, mathematical transformations, and tensor reshaping through expressive regex and structural targeting, while built-in assertions validate tensor shapes, data types, and values to prevent silent errors. We envision that B RAIN S URGERY will provide a strong foundation for future research through its reproducible and validated operations.
Model merging and task arithmetic A growing body of work demonstrates that meaningful knowledge can be transferred, combined, or suppressed by performing arithmetic directly in the weight space of pretrained models. Ilharco et al. (2023) introduced the concept of task vectors, directions in weight space obtained by subtracting pretrained weights from fine-tuned weights and showed that these vectors can be added or negated to compose or remove capabilities without any additional training. Building on this idea, Yadav et al. (2023) showed that naive weight averaging often fails due to sign conflicts and redundant parameters, and proposed a more principled merging strategy that resolves such interference. More broadly, model merging has emerged as an efficient paradigm for constructing multi-task learners that require neither joint training data nor separate parameter sets for each task (Yang et al., 2026). All of these methods ultimately reduce to sequences of tensor-level operations: addition, subtraction, scaling, and assignment, applied to specific layers of a neural network.
§ github.com/schneiderkamplab/brainsurgery
1
Introduction
The rapid proliferation of large-scale neural network models has transformed virtually every subfield of machine learning, from natural language processing to computer vision and beyond. While significant research effort has been devoted to designing training procedures and novel architectures, comparatively little attention has been paid to the post-hoc manipulation of trained model weights, a class of operations that has quietly become indispensable in both research and deployment settings. The ability to inspect, transform, compose, and verify neural network tensors in a principled and reproducible way underpins a surprisingly broad
Parameter-efficient adaptation and low-rank decomposition Low-Rank Adaptation (LoRA) (Hu et al., 2021) has become the dominant approach for fine-tuning large models under memory constraints, by decomposing weight updates into pairs of low-rank matrices. A critical but often overlooked step in the LoRA lifecycle is the integration of the adapter matrices back into the base weights prior to deployment, as well as the inverse oper-
* Equal contribution.
1
Figure 1: Overview of the B RAIN S URGERY workflow. Checkpoint rewrites are expressed as explicit declarative plans, inspected interactively, and validated through executable checks such as assert and diff. The depicted plan fragment is illustrative and includes advanced operations such as phlora, reflecting that the same workflow supports both simple tensor edits and more complex expert-rewriting pipelines.
ation of decomposing a full-rank weight matrix into low-rank factors for analysis or re-composition. Performing these operations correctly, across potentially hundreds of layers and with proper bookkeeping of tensor names, would benefit from tooling that operates directly on checkpoint files rather than through a full model loading pipeline.
(De Lange et al., 2022). Methods such as Elastic Weight Consolidation (Kirkpatrick et al., 2017) address this by constraining the update magnitude of individual weights according to their estimated importance, effectively requiring fine-grained, pertensor scaling and masking operations at the checkpoint level. Reproducing, extending, or debugging such methods demands direct, inspectable access to individual weight tensors.
Pruning and sparsification Model compression through pruning remains a major research direction, spanning unstructured weight removal, structured channel or head pruning, and the theoretical study of sparse subnetworks (Cheng et al., 2024; He and Xiao, 2023). Empirical studies in this space routinely require researchers to zero out specific weight tensors, delete entire parameter groups, clamp weight magnitudes, or verify that targeted sparsity patterns have been correctly applied. These operations must be performed with surgical precision: modifying the wrong subset of tensors, or failing to verify the result, can silently degrade model performance in ways that are difficult to diagnose after the fact.
Basic re-arrangements Beyond the research applications above, a substantial fraction of practical deep learning work involves adapting existing checkpoints to slightly different architectures or deployment targets: renaming layers, reshaping or transposing weight matrices, changing numerical precision, sharding large checkpoints across devices, and verifying that the resulting files are structurally sound. These tasks are currently handled through ad-hoc, one-off scripts that are difficult to audit, share, or reproduce. We need brainsurgeries Despite the breadth and importance of these use cases, the community lacks a unified, general-purpose tool for tensor-level manipulation of neural network checkpoints. Existing solutions are either tied to specific frameworks, focus solely on interpretability and activation manipulation rather than weights, or offer only a lim-
Continual learning and catastrophic forgetting When a neural network is sequentially fine-tuned on new tasks, it tends to overwrite weights that were important for previously learned tasks, a phenomenon known as catastrophic forgetting 2
ited set of operations. By making these operations composable, verifiable, and reproducible, B RAIN S URGERY fills a gap in the neural network research toolchain and lowers the barrier to a wide class of weight-space experiments that currently require bespoke, fragile scripts. Our contributions can be summarized as follows: • We present B RAIN S URGERY, a toolkit for fast and flexible tensor surgery on model checkpoints. The tool supports a comprehensive range of operations, including arithmetic composition, structural transformations, low-rank factorization and reconstruction, and a suite of verification primitives, enabling fine-grained model customization. • We provide code-free interaction modes, including a Web UI and declarative YAML plans, that are format-agnostic, operating natively on both safetensors and PyTorch checkpoints without loading any model code or instantiating any framework objects. This enables quick and reproducible setups while avoiding potential code incompatibilities. • We validate the correctness of model modifications using the built-in assertion mechanism, compare results against standard code-based implementations of the same operations, and present a model upcycling use case.
2
sarily focused on interpretability, but limited in the number of supported operations and often lacking fine-grained control, which prevents complete and detailed customization of models. Furthermore, leveraging the full capabilities of existing methods typically requires writing and executing custom code, introducing additional overhead and potential incompatibilities. Unlike previous approaches, B RAIN S URGERY provides a robust, purpose-built framework with an extensive set of operations for fine-grained modification of neural architecture weights. Several prior works are also restricted to a subset of architectures, whereas B RAIN S URGERY is architecture-agnostic. Its primary objective is to enable the application of operations and the modification of models in a way that allows them to be reused as-is, without the need for custom code, but directly through the definition of YAML plans. This does not preclude the use of B RAIN S URGERY for studying the effects of such operations on models for interpretability purposes – quite on the contrary, it enables a wide range of novel introspective and interventional applications.
3
B RAIN S URGERY
3.1
Design Principles
The design of B RAIN S URGERY is centered on providing a robust, transparent, and scalable framework for the surgical manipulation of neural network weights. Its architecture is guided by the following principles: 1. Declarative specification (OLY Grammar): Rather than requiring imperative scripts, B RAIN S URGERY employs a domain-specific language called OLY (One-Line YAML) and a structured YAML-based configuration. This allows users to declare what transformations should occur (e.g., weight scaling, merging, or pruning) rather than how to implement them. Separating specification from execution ensures legible and reproducible transformations. 2. Scalability for large models: Recognizing the memory constraints of modern Large Language Models (LLMs), B RAIN S URGERY is designed for performance. It implements sharded reading and writing for safetensors and provides multiple storage providers (inmemory and arena). The arena provider allows for out-of-core processing, enabling the editing of models that exceed the available system RAM. 3. Structural and pattern-based addressing:
Related Work
Several works have investigated model internals such as activations and weights. Many focus on the interpretability of language models (Zhao et al., 2024), injecting new knowledge into models by modifying their weights (Meng et al., 2022a,b; Gupta et al., 2024), acting in-real-time on hidden states through get and set operations and performing activation patching (Fiotto-Kaufman et al., 2024; Dumas, 2025; Belrose et al., 2023; Nanda and Bloom, 2022), or extracting concepts through attribution-based and concept-based methods (Poché et al., 2025). Others are more general, enabling manipulation of model weights through merging weights across different models (Goddard et al., 2024) or through of optimization-based techniques (Lepori et al., 2023). All prior works fall into one or both of the following categories. The first concerns model modifications whose sole purpose is internal analysis, focusing on interpretability and often targeting activations rather than weights. The second concerns targeted internal model modifications, not neces3
Input/output and memory management B RAIN S URGERY supports both safetensors files and standard PyTorch checkpoints (.pt, .bin), allowing operations on different formats without requiring any conversion. Checkpoint files for large models, such as large language models, can be very large; B RAIN S URGERY handles this by applying sharding to the modified checkpoints, allowing them to be saved as shards with customizable sizes.
Precision in “surgery” requires the ability to target specific layers or groups of parameters. The tool supports advanced pattern matching, including regular expressions and structured path patterns. This allows users to apply operations across complex architectures (e.g., targeting all attention.wv weights across 80 layers) with a single command. 4. Interactive and multi-modal interaction: To bridge the gap between automated pipelines and exploratory research, B RAIN S URGERY offers multiple interfaces. The batch CLI facilitates integration into CI/CD and training loops, while the Interactive CLI and Web UI allow researchers to experiment with weight edits in real-time, visualizing the results of individual operations before committing them to a final checkpoint. 5. Auditability and reproducibility: A core principle of the framework is the ability to track and reproduce edits. The tool features a summarize function that emits the exact sequence of transformations actually executed. This creates a “surgical log” that can be stored alongside edited models, ensuring that any weight modification is fully transparent and reproducible by other researchers, even if it was performed interactively. 3.2
Transformations Transformations (or more consise, transforms) are operations that can be applied to weight tensors of neural networks. These include the following type of operations: structural management (copy, move, delete, split, concat tensors), shape and type (reshape, permute, cast to different type), mathematical (insert values, sum, substract, dot product, matrix multiplication, scale by a scalar, clamp to a range), generation and initialization (fill a tensor with different modes e.g. constant, random), special (phlora, which splits a 2D target tensor into low-rank factors based on a specific rank). Tensor targeting and slicing Most transforms in B RAIN S URGERY require to specify source and/or destination tensors. This can be done by regex string matching or by structured expression system, allowing more flexibility and easy tensor targeting. Tensor slicing features are also provided, allowing to apply transforms also to subsections of tensors.
Features
The main features of B RAIN S URGERY can be divided into five categories: execution and reproducibility, input/output and memory management, tensor targeting and slicing, transformations, and inspection and validation.
Inspection and validation There are operations allowing to inspect tensors, e.g. diff to compare tensors and dump (with different formats) to summarize them. An assertion mechanism is also included, allowing to perform safety checks during a B RAIN S URGERY pipeline. A demonstration of this mechanism for validating B RAIN S URGERY is detailed in Section 4.1.
Execution and reproducibility Two execution modes are available: interactive mode and batch mode. In interactive mode it is possible to execute transformations on-the-fly through a CLI equipped with history and autocompletion. In batch mode, instead, a sequence of previously configured transformations is executed directly through YAML files (see Section 3.3), without any need to write code or interact with the CLI. In both cases reproducibility can be guaranteed. In batch mode, YAML configurations define a plan that can be replicated. In interactive mode, it is possible to create reproducibility summaries of the operations applied, producing YAML configurations that can then be used in batch mode to apply the same operations, making it easy to save exploratory interactive sessions as a reproducible script or even resume them.
Extensibility The framework is designed to be extensible: new transforms can be introduced by implementing a small Python class and placing it in the designated transforms directory, requiring no modifications to the core codebase. This allows users and contributors to grow the library of available transforms to suit custom workflows and model architectures. Memory management B RAIN S URGERY supports multiple memory providers for handling model weights and intermediate tensors. Notably, 4
the memory-mapped arena provider extends beyond what libraries such as safetensors typically offer: rather than memory-mapping only the model weights, it memory-maps all intermediate tensors and model copies as well. This allows large models to be manipulated efficiently without exhausting system RAM. 3.3
verifying the post-conditions using built-in assertions. If any operation deviates from its expected behavior, the engine’s strict ‘assert’ barriers immediately halt execution. This suite validates the system’s correctness across several core domains: • Namespace and memory management: The system successfully isolates state by creating, renaming, and removing virtual model aliases. Assertions like exists and not: exists confirm that garbage collection and pointer assignments function safely without memory leaks. • Arithmetic and in-place transformations: We perform step-by-step arithmetic tests, such as cloning a tensor x, computing x + x, and verifying the result against a deterministically scaled 2x tensor. Using appropriate assertions, we mathematically prove that both out-of-place (e.g. add) and in-place (e.g. add_) operations yield identical, correct outputs. • Structural and type transformations: The plan splits tensors into chunks and concatenates them back together, verifying via pairwise equality that no data is lost during structural manipulation. Additional checks confirm that reshape, permute, and datatype cast operations result in the exact dimensionalities (via assert: shape) and types (via assert: dtype) expected. • Advanced factorizations: For complex routines like Post-Hoc Low-Rank Adaptation (PHLoRA), the plan splits a 2D weight matrix into constituent A and B low-rank factors (Vasani et al., 2025) • I/O and state fidelity: To test lossless persistence, single tensors are saved to safetensors artifacts and reloaded into new destinations. Furthermore, a pristine checkpoint is loaded into an isolated alias and compared against the mutated environment using regex-based batch assertions, ensuring exact 1:1 parity for unmodified layers. By chaining these minimal atomic operations with continuous runtime validation, this validation plan shows that B RAIN S URGERY executes complex, stateful tensor surgeries deterministically. The assertion framework effectively transforms the tool into its own verifiable testbed, guaranteeing the strict precision required for reproducible scientific neural network editing.
B RAIN S URGERY Plans
The simplest, fastest, and code-free way to perform brainsurgeries is through the definition of a B RAIN S URGERY plan in YAML format, consisting of the following fields: • input: path to the model checkpoint (e.g., a safetensors file). • transforms: a sequence of transforms to apply, specifying the target and/or destination tensors along with the required parameters, via regex or the structured expression system described in Section 3.2. • output (optional): path of the modified model, output format, and shard size. The advantages of defining plans via YAML files are many. No code is required, hence, no environment setup, no model loading, no potential conflicts to resolve. Plans are easier and faster to set up, leading also to better readability. Each plan is fully reproducible, meaning that, once defined, it can be easily re-applied to the same starting model, yielding the same modifications. 3.4
Web UI
In addition to its command-line interfaces, B RAIN S URGERY provides a browser-based Web UI for interactive checkpoint inspection and editing. It allows users to browse tensor structure, apply transforms incrementally, and review the effects of edits before exporting the resulting checkpoint. Appendix A shows a screenshot of the WebUI.
4
Validation/Evaluation
4.1
Validation via Assertion Mechanism
To verify the operational correctness of B RAIN S URGERY, we developed a validation B RAIN S URGERY plan (as defined in Section 3.3) entirely within the tool’s own declarative framework. This approach leverages B RAIN S URGERY’s native assertion mechanism to validate transformations sequentially at runtime. The validation plan operates by executing minimal, controlled tensor mutations and immediately
4.2
Validation via PyTorch Equivalence
We validated the B RAIN S URGERY workflow by implementing a raw PyTorch equivalent of the same validation plan used in Section 4.1 and then comparing both executions in lockstep after every trans5
mean perplexity ratio (post/original) of 1.0 and top1 agreement of 100%. These results show that, for the tested prompts, B RAIN S URGERY preserves the model’s predictive behavior while enabling structured checkpoint transformations.
form. Each transform in the B RAIN S URGERY plan was mirrored by a corresponding PyTorch operation, and we performed step-by-step state comparisons (tensor presence, shape, dtype, and values) to verify equivalence at each stage. This procedure showed that the B RAIN S URGERY plan and the raw PyTorch implementation produce equivalent results transform-by-transform. Beyond correctness, we observed a clear usability and development-effort advantage for plans. B RAIN S URGERY plans are declarative and require no custom coding, which reduces debugging overhead and lowers the expertise needed to build and maintain transformation pipelines. They are also significantly more compact: the plan is 100 lines, while the equivalent raw PyTorch implementation is 421 lines (both excluding comments and blank lines), making it more than 4 times shorter. In practice, this makes B RAIN S URGERY plans faster to author, easier to review, more re-usable, and less error-prone than writing the same pipeline directly in imperative PyTorch code. 4.3
5
Declarative Tensor Surgery
This section connects the B RAIN S URGERY design principles, feature categories, and validation methodology described in Sections 3 and 4 to concrete checkpoint rewrites. Each example compares an imperative baseline written with Python, regular expressions, and PyTorch against the corresponding declarative B RAIN S URGERY fragment, illustrating how explicit plans make tensor surgery more structured, auditable, reproducible, and verifiable. The examples instantiate the same categories discussed above: model-scale tensor targeting, structural and type transformations, advanced factorizations such as PHLoRA, and validation through executable assertions and reference diffs. Additional standalone examples of slice copying, executable assertions, dense-to-expert (mixture of experts, MoE) upcycling, and in-place low-rank expert rewriting are provided in Appendix B; the latter uses subtract_, phlora_, and add_.
Validation via Inference Preservation
We validated the correctness of B RAIN S URGERY by applying a sequence of transforms to a checkpoint and then reversing them, restoring the model to its original state – this is what we refer to as the post-surgery checkpoint. We then verified that the post-surgery checkpoint remains usable for language generation with both qualitative and quantitative tests.
Expert rewrites Dense-to-expert MoE upcycling, shown in Appendix B, exercises namespace and state-management behavior through alias-level copying and deletion, as well as structural transformation through sliced router initialization and shape assertions. Figure 2 shows the full PHLoRA workflow rather than only the inner tensor rewrite: the imperative baseline includes checkpoint loading, format handling, PHLoRA factorization, dtype conversion, deletion, local assertions, reference comparison, and sharded output, while B RAIN S URGERY records the same workflow as one declarative plan.
Qualitative prompt-based checks. We ran inference on a set of 50 prompts and manually verified that the post-surgery model loaded successfully and produced coherent continuations, indicating that the transform pipeline did not break end-toend generation behavior. Quantitative consistency checks. We also compared the original checkpoint and the postsurgery checkpoint on the same prompt set using lightweight regression metrics: last-token logit cosine similarity, prompt-level perplexity, top-1 nexttoken agreement. As noted earlier, for the postsurgery checkpoint we apply transforms forward and backward in order to first modify and then restore the original state of the checkpoint, therefore we expect to have perfect or near-perfect metrics. Across 50 prompts, we observed near-identical outputs with both mean cosine similarity of and
Bulk tensor targeting The example in Figure 3 shows model-scale checkpoint editing through regex-based tensor targeting. The imperative baseline must compile a pattern, iterate over checkpoint names, and mutate matching tensors manually. In the B RAIN S URGERY fragment, the same target family and operation are stated directly: scale_ applies to all matching attention projection weights. Even for this small rewrite, the declarative plan makes the intended edit easier to inspect. 6
Imperative Python/PyTorch baseline
B RAIN S URGERY plan
from pathlib import Path import json import torch from safetensors.torch import load_file, save_file
inputs: - model::models/input.safetensors - ref::models/reference.safetensors transforms: - copy: from: "(.*experts\.1\..*)\.weight", to: "\1.delta" - subtract_: from: "(.*experts)\.0\.(.*)", to: "\1.1.\2.delta" - phlora: target: "(.*experts\.1\..*)\.delta" target_a: "\1.phlora_a" target_b: "\1.phlora_b" rank: 64 - cast_: target: ".*experts\.1\.phlora_(a|b)" to: float16 - delete: target: ".*experts\.1\..*\.delta" - assert: dtype: { of: ".*experts\.1\..*.phlora_(a|b)", is: float16 } - assert: not: { exists: ".*experts\.1\..*\.weight" }
input_path = Path("models/input.safetensors") source = load_file(str(input_path)) if input_path.suffix == ". safetensors" else torch.load(input_path, weights_only=True) ref = load_file("models/reference.safetensors") out = dict(source) for layer in range(16): prefix = f"model.layers.{layer}.mlp.experts" for proj in ("gate_proj", "up_proj", "down_proj"): e0 = f"{prefix}.0.{proj}.weight" e1 = f"{prefix}.1.{proj}.weight" delta = source[e1] - source[e0] u, s, vh = torch.linalg.svd(delta, full_matrices=False) sqrt_s = s[:64].sqrt() a = sqrt_s[:, None] * vh[:64, :] b = u[:, :64] * sqrt_s out[f"{prefix}.1.{proj}.phlora_a.weight"] = a.to( dtype=torch.float16, device=source[e1].device ) out[f"{prefix}.1.{proj}.phlora_b.weight"] = b.to( dtype=torch.float16, device=source[e1].device ) del out[e1]
output: path: models/output format: safetensors shard: 1GB
assert out["model.layers.0.mlp.experts.1.gate_proj.phlora_a.weight"] .dtype == torch.float16 assert "model.layers.0.mlp.experts.1.gate_proj.weight" not in out out_dir = Path("models/output") max_bytes = 1 << 30 out_dir.mkdir(parents=True, exist_ok=True) shards, cur, cur_size = [], {}, 0 for name, tensor in sd.items(): size = tensor.numel() * tensor.element_size() if cur and cur_size + size > max_bytes: shards.append(cur) cur, cur_size = {}, 0 cur[name] = tensor cur_size += size if cur: shards.append(cur) weight_map = {} for idx, shard in enumerate(shards, start=1): shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors" save_file(shard, str(out_dir / shard_name)) for name in shard: weight_map[name] = shard_name (out_dir / "model.safetensors.index.json").write_text( json.dumps({"weight_map": weight_map}), encoding="utf-8" )
Figure 2: Full PHLoRA workflow with validation. When assertions, reference comparison, checkpoint I/O, and sharded output are included, the imperative baseline must configure loading, mutation, validation, and persistence explicitly, while B RAIN S URGERY keeps the workflow in one declarative plan.
Imperative Python/Re baseline
Tensor surgery validation The local assertions and reference comparison in Figure 2 instantiate the validation methodology described in Section 4. The same mechanism scales from local post-conditions, such as dtype and deletion checks, to end-to-end agreement with an independent reference via diff, which reports missing-on-left, missing-on-right, and differing tensors.
import re import torch sd = torch.load("models/input.pt") pattern = re.compile(r".*self_attn\..*_proj\.weight") for name, tensor in sd.items(): if pattern.fullmatch(name): sd[name] = tensor * 0.5 torch.save(sd, "models/output.pt")
B RAIN S URGERY transform inputs: [ models/input.pt ] scale_: target: ".*self_attn\..*_proj\.weight", by: 0.5 output: models/output.pt
6
Figure 3: Bulk tensor targeting. The imperative baseline loops over matching checkpoint names; the B RAIN S URGERY fragment expresses the same regex target family and scale operation as one declarative transform.
Discussion
Our examples and case studies support four main claims about B RAIN S URGERY. First, it is expressive: operations such as scale_, copy, fill, delete, subtract_, phlora_, and phlora directly encode checkpoint manipulations that would otherwise be buried inside handwritten state-dict 7
Limitations
code. Second, it is consistent: the same targeting and reference language supports bulk edits, sliced references, assertions, dense-to-expert upcycling, low-rank rewriting, and PHLoRA factorization. Third, it is auditable: plans make intended rewrites reviewable rather than distributing logic across loops, conditionals, and in-place mutation. Finally, it supports reproducibility and validation: local claims can be checked with assert, while end-to-end agreement with an independent PyTorch reference can be checked with diff.
B RAIN S URGERY improves the rigor and reproducibility of checkpoint surgery, but does not remove the need for model-specific expertise when designing transformations. Diff based validation establishes equivalence to a reference transformation, not downstream quality, training stability, or runtime compatibility with every external framework. Some rewrites may still require framework-specific metadata, configuration changes, loader support, or custom interpretation, especially for factorized formats such as PHLoRA. Finally, the current evaluation focuses on checkpoint surgery and structural rewriting; broader benchmarking is still needed across larger models, distributed settings, and more diverse transformation families.
Beyond the transformations shown above, B RAIN S URGERY is also extensible and memorysavvy. New transforms can be added without modifying the core engine, and the memory-mapped arena provider can map intermediate tensors and model copies in addition to stored weights. The broader methodological point is that model-weight transformations are treated as first-class research artifacts rather than opaque implementation details.
Acknowledgements The research was supported in part by the Danish Foundation Models project, funded by the Danish government. This research was further supported in part by the MIST project, funded by the Novo Nordisk Foundation under grant reference number NNF25OC0103204. Part of the computation for this project was performed on the UCloud interactive HPC system managed by the eScience Center at the University of Southern Denmark.
This is particularly relevant for current work on expert architectures and memory-efficient lowrank adaptation, where checkpoint rewrites such as MoE upcycling and PHLoRA-style factorization are themselves part of the research method.
7
Conclusion References
B RAIN S URGERY turns checkpoint surgery from ad-hoc scripting into a declarative, auditable, and verifiable workflow. Across bulk targeting, slicing, executable assertions, dense-to-expert MoE upcycling, low-rank expert rewriting, and PHLoRA factorization, the examples show that explicit tensorsurgery primitives can express realistic modeltransformation workflows as reusable plans.
Nora Belrose, Igor Ostrovsky, Lev McKinney, Zach Furman, Logan Smith, Danny Halawi, Stella Biderman, and Jacob Steinhardt. 2023. Eliciting latent predictions from transformers with the tuned lens. arXiv preprint arXiv:2303.08112. Hongrong Cheng, Miao Zhang, and Javen Qinfeng Shi. 2024. A survey on deep neural network pruning: Taxonomy, comparison, analysis, and recommendations. IEEE Transactions on Pattern Analysis and Machine Intelligence, 46(12):10558–10578.
Built-in reference diffing and lightweight prompt-level regression checks support this workflow structurally and behaviorally: the former verifies agreement with independent implementations, while the latter showed near-identical predictive behavior before and after reversible checkpoint surgery in the tested setting. The B RAIN S URGERY Web UI brings plan construction, execution, preview impact, checkpoint diffing, and execution summaries into one interface, reinforcing the same goal: checkpoint surgery should be explicit, inspectable, and reproducible rather than hidden inside one-off scripts.
Matthias De Lange, Rahaf Aljundi, Marc Masana, Sarah Parisot, Xu Jia, Ales Leonardis, Gregory Slabaugh, and Tinne Tuytelaars. 2022. A continual learning survey: Defying forgetting in classification tasks. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(7):3366–3385. Clément Dumas. 2025. nnterp: A standardized interface for mechanistic interpretability of transformers. arXiv preprint arXiv:2511.14465. Jaden Fiotto-Kaufman, Alexander R. Loftus, Eric Todd, Jannik Brinkmann, Koyena Pal, Dmitrii Troitskii, Michael Ripa, Adam Belfki, Can Rager, Caden Juang,
8
Aaron Mueller, Samuel Marks, Arnab Sen Sharma, Francesca Lucchetti, Nikhil Prakash, Carla Brodley, Arjun Guha, Jonathan Bell, Byron C. Wallace, and David Bau. 2024. Nnsight and ndif: Democratizing access to open-weight foundation model internals. arXiv preprint arXiv:2407.14561.
and Fanny Jourdan. 2025. Interpreto: An explainability library for transformers. arXiv preprint arXiv:2512.09730. Bhoomit Vasani, Jack FitzGerald, Anjie Fang, and Sushmit Vaish. 2025. Phlora: data-free post-hoc lowrank adapter extraction from full-rank checkpoint. Preprint, arXiv:2509.10971.
Charles Goddard, Shamane Siriwardhana, Malikeh Ehghaghi, Luke Meyers, Vlad Karpukhin, Brian Benedict, Mark McQuade, and Jacob Solawetz. 2024. Arcee’s mergekit: A toolkit for merging large language models. arXiv preprint arXiv:2403.13257.
Prateek Yadav, Derek Tam, Leshem Choshen, Colin A. Raffel, and Mohit Bansal. 2023. TIES-merging: Resolving interference when merging models. In Advances in Neural Information Processing Systems, volume 36.
Akshat Gupta, Dev Sajnani, and Gopala Anumanchipalli. 2024. A unified framework for model editing. In Findings of the Association for Computational Linguistics: EMNLP 2024, pages 15403– 15418, Miami, Florida, USA. Association for Computational Linguistics.
Enneng Yang, Li Shen, Guibing Guo, Xingwei Wang, Xiaochun Cao, Jie Zhang, and Dacheng Tao. 2026. Model merging in llms, mllms, and beyond: Methods, theories, applications, and opportunities. ACM Computing Surveys, 58(8):1–41.
Yang He and Lingao Xiao. 2023. Structured pruning for deep convolutional neural networks: A survey. IEEE transactions on pattern analysis and machine intelligence, 46(5):2900–2919.
Haiyan Zhao, Hanjie Chen, Fan Yang, Ninghao Liu, Huiqi Deng, Hengyi Cai, Shuaiqiang Wang, Dawei Yin, and Mengnan Du. 2024. Explainability for large language models: A survey. ACM Trans. Intell. Syst. Technol., 15(2).
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. 2021. Lora: Low-rank adaptation of large language models. Preprint, arXiv:2106.09685.
A
B RAIN S URGERY Web UI
The B RAIN S URGERY Web UI for interactive checkpoint inspection, transform execution, previewing edit effects, and checkpoint export. After installtion, the B RAIN S URGERY Web UI can be accessed via the command; brainsurgery webui.
Gabriel Ilharco, Marco Tulio Ribeiro, Mitchell Wortsman, Suchin Gururangan, Ludwig Schmidt, Hannaneh Hajishirzi, and Ali Farhadi. 2023. Editing models with task arithmetic. In The Eleventh International Conference on Learning Representations. James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A. Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska, Demis Hassabis, Claudia Clopath, Dharshan Kumaran, and Raia Hadsell. 2017. Overcoming catastrophic forgetting in neural networks. Proceedings of the National Academy of Sciences, 114(13):3521–3526.
B
Additional B RAIN S URGERY vs Imperative Baseline
This Appendix presents B RAIN S URGERY through a compact progression of 5 examples, 3 case studies. Supplementary (Section 5) Examples (Section B.1; Figure 7; Figure 8; Figure 9; Figure 10); Figure 11/Cases (Section B.2; Figure 12; Figure 13; Figure 14). We recognize, that there is not only one imperative way to express the corresponding rewrite in Python, regular expressions (re) and PyTorch (torch). The same checkpoint transformation can often be realized through different combinations of loops, indexing, mutation, helper logic, and intermediate state, even when the intended effect is the same. By contrast, once the relevant references are known, B RAIN S URGERY keeps the rewrite in a more stable declarative form that more directly captures the semantic intent of the operation, making it more explicit, expressive, auditable, consistent, and reproducible across implementations. Throughout, the emphasis is on what each rewrite does to the checkpoint, why that effect
Michael A. Lepori, Ellie Pavlick, and Thomas Serre. 2023. Neurosurgeon: A toolkit for subnetwork analysis. arXiv preprint arXiv:2309.00244. Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. 2022a. Locating and editing factual associations in gpt. In Advances in Neural Information Processing Systems 35 (NeurIPS 2022). Kevin Meng, Arnab Sen Sharma, Alex Andonian, Yonatan Belinkov, and David Bau. 2022b. Massediting memory in a transformer. arXiv preprint arXiv:2210.07229. Neel Nanda and Joseph Bloom. 2022. Transformerlens. https://github.com/TransformerLensOrg/ TransformerLens. Antonin Poché, Thomas Mullor, Gabriele Sarti, Frédéric Boisnard, Corentin Friedrich, Charlotte Claye, François Hoofd, Raphael Bernas, Céline Hudelot,
9
Figure 4: B RAIN S URGERY Web UI figure showing model dump.
Figure 5: B RAIN S URGERY Web UI figure showing model move.
Figure 6: B RAIN S URGERY Web UI figure showing zoom-in on diff between the original model and the rewritten model after applying scale_.
is useful, and how explicit plans turn checkpoint manipulation and its validation into reviewable research artifacts.
with the corresponding B RAIN S URGERY transform fragments. When a block is cropped from a longer script or plan, [...] marks omitted continuation. The appendix gives the isolated slice-copy and assertion
Case studies compare larger imperative rewrites 10
Imperative Python/Re baseline
examples separately; in the main text, those mechanisms are shown where they are used in realistic rewrites. B.1
import re import torch sd = torch.load("models/input.pt") pattern = re.compile(r".*self_attn\..*_proj\.weight") for name, tensor in sd.items(): if pattern.fullmatch(name): sd[name] = tensor * 0.5 torch.save(sd, "models/output.pt")
Examples
Example: Targeting with Slices The example in Figure 7 shows precise local tensor surgery. The comparison is about copy plus a slice reference in from.
B RAIN S URGERY transform inputs: [ models/input.pt ] transforms: - scale_: target: ".*self_attn\..*_proj\.weight", by: 0.5 output: models/output.pt
Imperative baseline w = sd["model.layers.0.self_attn.q_proj.weight"] sd["tmp"] = w[:128, :128].clone()
Figure 9: Bulk tensor targeting. The imperative baseline loops over matching checkpoint names; the B RAIN S URGERY fragment expresses the same regex target family and scale operation as one declarative transform.
B RAIN S URGERY transform fragment - copy: from: ".*\.0\..*\.self_attn.q_proj.*::[:128, :128]", to: "tmp"
Figure 7: Example ensor slicing. Both sides copy the same [:128, :128] block into the same destination tensor slot.
Prefix Rewrite The example in Figure 10 shows a pure structural rewrite: all tensors under one checkpoint prefix are moved under another prefix.
Example: Verification as Executable Invariants The example in Figure 8 example shows that B RAIN S URGERY is also a language for validation. The inline checks correspond directly to assert, exists, shape, and equal.
Imperative Python/Re baseline
Imperative PyTorch baseline
B RAIN S URGERY transform fragment
import torch
move: from: "text_model\.(.*)", to: "model.\1"
import re pattern = re.compile(r"text_model\.(.*)") for name in list(sd): match = pattern.fullmatch(name) if match: sd[f"model.{match.group(1)}"] = sd.pop(name)
assert "layers.0.gate.weight" in out assert out["layers.0.gate.weight"].shape == (2, 2048) assert torch.equal( src["layers.0.gate.weight"][:16, :16], out["layers.0.gate.weight"][:16, :16], ) assert "layers.0.gate.bias" not in out
Figure 10: Prefix rewrite. The imperative baseline loops over checkpoint names and manually rewrites matching keys; the B RAIN S URGERY fragment expresses the same regex capture and move as one declarative transform.
B RAIN S URGERY transform fragment - assert: exists: "out::layers.0.gate.weight" - assert: shape: { of: "out::layers.0.gate.weight", is: [2, 2048] } - assert: equal: left: "src::layers.0.gate.weight::[:16, :16]" right: "out::layers.0.gate.weight::[:16, :16]" - assert: not: { exists: "out::layers.0.gate.bias" }
Example: Tensor Surgery Validation Figure 11 shows the B RAIN S URGERY validation artifact and diff.
Figure 8: Example validation as executable invariants. Both sides check the same existence, shape, equality, and deletion post-conditions.
B.2
Case Studies
Case Study: Dense-to-Expert MoE Upcycling Figure 12 expands the dense-to-expert MoE example from an inner rewrite into a full checkpoint workflow. The imperative baseline loads two dense checkpoints, copies projections into expert slots, initializes the router from a sliced source tensor, deletes the original dense projections, checks local post-conditions, compares against a reference checkpoint, and saves sharded output. The B RAIN S URGERY plan records the same workflow declaratively.
Example: Bulk Tensor Targeting The example in Figure 9 (as in main text Figure 3) shows regex-based model-scale targeting. The imperative baseline must compile a pattern and loop over the state dict; the B RAIN S URGERY fragment states the same target family and operation in one line. The operation is explicit as scale_, rather than being hidden inside a handwritten loop over matching tensor names. 11
Imperative PyTorch baseline
B RAIN S URGERY validation artifact
import torch from safetensors.torch import load_file
inputs: - yaml::models/example_yaml_output - ref::models/example_reference_output
yaml_out = load_file("models/example_yaml_output") ref_out = load_file("models/example_reference_output")
transforms: - diff: { mode: aliases, left_alias: ref, right_alias: yaml }
missing_on_left = sorted(set(ref_out) - set(yaml_out)) missing_on_right = sorted(set(yaml_out) - set(ref_out)) differing = [] for name in sorted(set(yaml_out) & set(ref_out)): if yaml_out[name].shape != ref_out[name].shape: differing.append(name) elif not torch.equal(yaml_out[name], ref_out[name]): differing.append(name) print("Diff: yaml <-> ref") print("Missing on left:" + "\n- ".join(missing_on_left)) print("Missing on right:" + "\n- ".join(missing_on_right)) print("Differing:" + "\n- ".join(differing))
Figure 11: Validation with diff. Local invariants can be checked with assert, while end-to-end agreement with an independent reference can be checked by diffing the reference output alias against the output produced by the B RAIN S URGERY plan.
Case Study: Expert Rewrites/PHLoRA Factorization Figure 13 shows the full PHLoRA workflow rather than only the inner tensor rewrite (as in main text Figure 2): the imperative baseline includes checkpoint loading, format handling, PHLoRA factorization, dtype conversion, deletion, local assertions, and sharded output, while B RAIN S URGERY records the same workflow as one declarative plan. Case Study: Low-Rank Expert Rewrite Figure 14 gives the same full-workflow treatment for the in-place low-rank expert rewrite. Unlike PHLoRA factorization, which writes explicit factor tensors, this rewrite keeps the dense expert slot and replaces it with the anchor expert plus a ranklimited approximation of the expert delta.
12
Imperative Python/PyTorch baseline
B RAIN S URGERY plan
from pathlib import Path import json import torch from safetensors.torch import load_file, save_file
inputs: - m0::models/dense_a.safetensors - m1::models/dense_b.safetensors - ref::models/moe_reference.safetensors output: path: models/moe_output format: safetensors shard: 1GB transforms: - copy: { from: "m0::model.layers\.(\d+)\.mlp\.(.*_proj)\.weight", to: "m0::model.layers.\1.mlp.experts.0.\2.weight" } - copy: { from: "m1::model.layers\.(\d+)\.mlp\.(.*_proj)\.weight", to: "m0::model.layers.\1.mlp.experts.1.\2.weight" } - fill: from: "m0::model.layers\.(\d+)\.self_attn\.q_proj\.weight::[:2,:]" to: "m0::model.layers.\1.mlp.gate.weight" mode: constant value: 0 - delete: { target: "m0::model.layers\.(\d+)\.mlp\.(.*_proj)\.weight" } - assert: shape: { of: "m0::model.layers.0.mlp.gate.weight", is: [2, 2048] } - assert: not: { exists: "m0::model.layers.0.mlp.gate_proj.weight" } - assert: all: - equal: { left: "m0::(.+)", right: "ref::\1" } - equal: { left: "ref::(.+)", right: "m0::\1" }
def load_checkpoint(path): return load_file(str(path)) if path.suffix == ".safetensors" else torch.load(path, weights_only=True) def save_sharded_safetensors(sd, out_dir, max_bytes): out_dir.mkdir(parents=True, exist_ok=True) shards, cur, cur_size = [], {}, 0 for name, tensor in sd.items(): size = tensor.numel() * tensor.element_size() if cur and cur_size + size > max_bytes: shards.append(cur) cur, cur_size = {}, 0 cur[name] = tensor cur_size += size if cur: shards.append(cur) weight_map = {} for idx, shard in enumerate(shards, start=1): shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors" save_file(shard, str(out_dir / shard_name)) for name in shard: weight_map[name] = shard_name (out_dir / "model.safetensors.index.json").write_text( json.dumps({"weight_map": weight_map}), encoding="utf-8" ) def assert_same_state_dict(left, right): missing_l = sorted(set(right) - set(left)) missing_r = sorted(set(left) - set(right)) differing = [k for k in set(left) & set(right) if left[k].shape != right[k].shape or not torch.equal(left[k], right[k])] assert missing_l == missing_r == differing == [] dense_a = load_checkpoint(Path("models/dense_a.safetensors")) dense_b = load_checkpoint(Path("models/dense_b.safetensors")) ref = load_checkpoint(Path("models/moe_reference.safetensors")) out = dict(dense_a) for layer in range(16): for expert, dense_sd in ((0, dense_a), (1, dense_b)): for proj in ("gate_proj", "up_proj", "down_proj"): src = f"model.layers.{layer}.mlp.{proj}.weight" dst = f"model.layers.{layer}.mlp.experts.{expert}.{proj}. weight" out[dst] = dense_sd[src].clone() q = f"model.layers.{layer}.self_attn.q_proj.weight" out[f"model.layers.{layer}.mlp.gate.weight"] = torch.zeros_like( dense_a[q][:2, :] ) for proj in ("gate_proj", "up_proj", "down_proj"): del out[f"model.layers.{layer}.mlp.{proj}.weight"] assert out["model.layers.0.mlp.gate.weight"].shape[0] == 2 assert "model.layers.0.mlp.gate_proj.weight" not in out assert_same_state_dict(out, ref) save_sharded_safetensors(out, Path("models/moe_output"), 1 << 30)
Figure 12: Full dense-to-expert MoE workflow with validation. Including checkpoint I/O, reference comparison, and sharded output makes the imperative baseline responsible for loading, mutation, validation, and persistence, while B RAIN S URGERY keeps the same structural rewrite and checks in one plan.
13
Imperative Python/PyTorch baseline
B RAIN S URGERY plan
from pathlib import Path import json import torch from safetensors.torch import load_file, save_file
inputs: - model::models/input.safetensors - ref::models/reference.safetensors transforms: - copy: from: "(.*experts\.1\..*)\.weight", to: "\1.delta" - subtract_: from: "(.*experts)\.0\.(.*)", to: "\1.1.\2.delta" - phlora: target: "(.*experts\.1\..*)\.delta" target_a: "\1.phlora_a" target_b: "\1.phlora_b" rank: 64 - cast_: target: ".*experts\.1\.phlora_(a|b)" to: float16 - delete: target: ".*experts\.1\..*\.delta" - assert: dtype: { of: ".*experts\.1\..*.phlora_(a|b)", is: float16 } - assert: not: { exists: ".*experts\.1\..*\.weight" }
input_path = Path("models/input.safetensors") source = load_file(str(input_path)) if input_path.suffix == ". safetensors" else torch.load(input_path, weights_only=True) ref = load_file("models/reference.safetensors") out = dict(source) for layer in range(16): prefix = f"model.layers.{layer}.mlp.experts" for proj in ("gate_proj", "up_proj", "down_proj"): e0 = f"{prefix}.0.{proj}.weight" e1 = f"{prefix}.1.{proj}.weight" delta = source[e1] - source[e0] u, s, vh = torch.linalg.svd(delta, full_matrices=False) sqrt_s = s[:64].sqrt() a = sqrt_s[:, None] * vh[:64, :] b = u[:, :64] * sqrt_s out[f"{prefix}.1.{proj}.phlora_a.weight"] = a.to( dtype=torch.float16, device=source[e1].device ) out[f"{prefix}.1.{proj}.phlora_b.weight"] = b.to( dtype=torch.float16, device=source[e1].device ) del out[e1]
output: path: models/output format: safetensors shard: 1GB
assert out["model.layers.0.mlp.experts.1.gate_proj.phlora_a.weight"] .dtype == torch.float16 assert "model.layers.0.mlp.experts.1.gate_proj.weight" not in out out_dir = Path("models/output") max_bytes = 1 << 30 out_dir.mkdir(parents=True, exist_ok=True) shards, cur, cur_size = [], {}, 0 for name, tensor in sd.items(): size = tensor.numel() * tensor.element_size() if cur and cur_size + size > max_bytes: shards.append(cur) cur, cur_size = {}, 0 cur[name] = tensor cur_size += size if cur: shards.append(cur) weight_map = {} for idx, shard in enumerate(shards, start=1): shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors" save_file(shard, str(out_dir / shard_name)) for name in shard: weight_map[name] = shard_name (out_dir / "model.safetensors.index.json").write_text( json.dumps({"weight_map": weight_map}), encoding="utf-8" )
Figure 13: Full PHLoRA workflow with validation. When assertions, reference comparison, checkpoint I/O, and sharded output are included, the imperative baseline must configure loading, mutation, validation, and persistence explicitly, while B RAIN S URGERY keeps the workflow in one declarative plan.
14
Imperative Python/PyTorch baseline
B RAIN S URGERY plan
from pathlib import Path import json import torch from safetensors.torch import load_file, save_file
inputs: - model::models/input.safetensors - ref::models/low_rank_reference.safetensors transforms: - subtract_: from: "model.layers\.(\d+)\.mlp\.experts\.0\.(.*_proj)\.weight" to: "model.layers.\1.mlp.experts.1.\2.weight" - phlora_: target: "model.layers\.(\d+)\.mlp\.experts\.1\.(.*_proj)\.weight" rank: 64 - add_: from: "model.layers\.(\d+)\.mlp\.experts\.0\.(.*_proj)\.weight" to: "model.layers.\1.mlp.experts.1.\2.weight" - cast_: target: "model.layers\.(\d+)\.mlp\.experts\.1\.(.*_proj)\.weight" to: float16 - assert: dtype: of: "model.layers.0.mlp.experts.1.gate_proj.weight" is: float16 - assert: all: - equal: { left: "model::(.+)", right: "ref::\1" } - equal: { left: "ref::(.+)", right: "model::\1" }
def load_checkpoint(path): return load_file(str(path)) if path.suffix == ".safetensors" else torch.load(path, weights_only=True) def save_sharded_safetensors(sd, out_dir, max_bytes): out_dir.mkdir(parents=True, exist_ok=True) shards, cur, cur_size = [], {}, 0 for name, tensor in sd.items(): size = tensor.numel() * tensor.element_size() if cur and cur_size + size > max_bytes: shards.append(cur) cur, cur_size = {}, 0 cur[name] = tensor cur_size += size if cur: shards.append(cur) weight_map = {} for idx, shard in enumerate(shards, start=1): shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors" save_file(shard, str(out_dir / shard_name)) for name in shard: weight_map[name] = shard_name (out_dir / "model.safetensors.index.json").write_text( json.dumps({"weight_map": weight_map}), encoding="utf-8" )
output: path: models/low_rank_output format: safetensors shard: 1GB
def assert_same_state_dict(left, right): missing_l = sorted(set(right) - set(left)) missing_r = sorted(set(left) - set(right)) differing = [k for k in set(left) & set(right) if left[k].shape != right[k].shape or not torch.equal(left[k], right[k])] assert missing_l == missing_r == differing == [] source = load_checkpoint(Path("models/input.safetensors")) ref = load_checkpoint(Path("models/low_rank_reference.safetensors")) out = dict(source) for layer in range(16): for proj in ("gate_proj", "up_proj", "down_proj"): e0 = f"model.layers.{layer}.mlp.experts.0.{proj}.weight" e1 = f"model.layers.{layer}.mlp.experts.1.{proj}.weight" delta = source[e1] - source[e0] u, s, vh = torch.linalg.svd(delta, full_matrices=False) approx = (u[:, :64] * s[:64]) @ vh[:64, :] out[e1] = (source[e0] + approx).to( dtype=torch.float16, device=source[e1].device, ) assert out["model.layers.0.mlp.experts.1.gate_proj.weight"].dtype == torch.float16 assert_same_state_dict(out, ref) save_sharded_safetensors(out, Path("models/low_rank_output"), 1 << 30)
Figure 14: Full in-place low-rank expert rewrite with validation. The imperative baseline spells out checkpoint loading, SVD-based low-rank reconstruction, dtype conversion, reference comparison, and sharded output; the B RAIN S URGERY plan expresses the same workflow with subtract_, phlora_, add_, cast_, assert, and diff.
15