BrainSurgery: Manipulações Declarativas de Pesos Reprodutíveis e Confiáveis para Edição e Reciclagem de Modelos

Resumo

À medida que os modelos de deep learning escalam, gerenciar, inspecionar e modificar grandes checkpoints tem se tornado cada vez mais desafiador. Pesquisadores frequentemente precisam alterar pesos de modelos para reestruturação de camadas, conversão de precisão, fatoração de posto baixo e depuração arquitetural, embora esses fluxos de trabalho frequentemente dependam de frágeis scripts Python ad-hoc. Aqui, apresentamos o BrainSurgery, uma ferramenta para "cirurgia tensorial" robusta e reprodutível em checkpoints de redes neurais, e fornecemos uma demonstração do sistema abrangendo quatro exemplos e três estudos de caso, desde reciclagem de modelos até extração de LoRA. Ao abstrair formatos de armazenamento e gerenciamento de memória, o BrainSurgery executa transformações complexas por meio de planos YAML declarativos. Ele suporta modificações estruturais, transformações matemáticas e redimensionamento tensorial por meio de expressões regulares expressivas e segmentação estrutural, enquanto asserções embutidas validam formas de tensores, tipos de dados e valores para evitar erros silenciosos. Prevemos que o BrainSurgery fornecerá uma base sólida para pesquisas futuras por meio de suas operações reprodutíveis e 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.