BrainSurgery: воспроизводимые и надежные декларативные манипуляции весами для редактирования и апсайклинга моделей

Аннотация

По мере масштабирования моделей глубокого обучения управление, проверка и модификация больших контрольных точек становятся всё более сложными задачами. Исследователям часто требуется изменять веса моделей для реструктуризации слоёв, приведения типов данных, низкорангового разложения и отладки архитектуры, однако такие рабочие процессы зачастую опираются на хрупкие ad-hoc скрипты на Python. В данной работе мы представляем BrainSurgery — инструмент для надёжной и воспроизводимой «тензорной хирургии» контрольных точек нейронных сетей, а также демонстрируем его работу на четырёх примерах и трёх тематических исследованиях — от апсайклинга моделей до извлечения LoRA. Абстрагируя форматы хранения и управление памятью, BrainSurgery выполняет сложные преобразования на основе декларативных YAML-планов. Он поддерживает структурные модификации, математические преобразования и изменение формы тензоров с помощью выразительных регулярных выражений и структурного таргетирования, а встроенные проверки валидируют размерности тензоров, типы данных и значения, предотвращая незаметные ошибки. Мы полагаем, что BrainSurgery обеспечит прочную основу для будущих исследований благодаря своим воспроизводимым и валидированным операциям.

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.