BrainSurgery: Manipulaciones de Pesos Declarativas, Reproducibles y Fiables para la Edición y Upcycling de Modelos

Resumen

A medida que los modelos de aprendizaje profundo escalan, gestionar, inspeccionar y modificar puntos de control (checkpoints) extensos se ha vuelto cada vez más desafiante. Los investigadores frecuentemente necesitan alterar los pesos del modelo para reestructuración de capas, conversión de precisión, factorización de bajo rango y depuración arquitectónica; sin embargo, estos flujos de trabajo suelen depender de frágiles scripts ad-hoc en Python. Aquí presentamos BrainSurgery, una herramienta para una "cirugía de tensores" robusta y reproducible sobre puntos de control de redes neuronales, junto con una demostración del sistema que abarca cuatro ejemplos y tres casos de estudio, desde el reciclaje de modelos hasta la extracción de LoRA. Al abstraer los formatos de almacenamiento y la gestión de memoria, BrainSurgery ejecuta transformaciones complejas mediante planes declarativos en YAML. Admite modificaciones estructurales, transformaciones matemáticas y reformateo de tensores mediante expresiones regulares y direccionamiento estructural, mientras que aserciones integradas validan formas, tipos de datos y valores de los tensores para evitar errores silenciosos. Prevemos que BrainSurgery proporcionará una base sólida para futuras investigaciones gracias a sus operaciones reproducibles y validadas.

English

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 BrainSurgery, 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, BrainSurgery 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 BrainSurgery will provide a strong foundation for future research through its reproducible and validated operations.