BrainSurgery : Manipulations Déclaratives des Poids Reproductibles et Fiables pour l’Édition et le Surcyclage de Modèles

Résumé

À mesure que les modèles d'apprentissage profond passent à l'échelle, la gestion, l'inspection et la modification de grands points de contrôle deviennent de plus en plus complexes. Les chercheurs ont souvent besoin d'altérer les poids des modèles pour la restructuration de couches, le transtypage de précision, la factorisation de bas rang et le débogage architectural, mais ces processus reposent fréquemment sur des scripts ad hoc fragiles en Python. Nous présentons ici BrainSurgery, un outil pour une « chirurgie tensorielle » robuste et reproductible sur les points de contrôle de réseaux de neurones, et fournissons une démonstration du système couvrant quatre exemples et trois études de cas, allant du recyclage de modèles à l'extraction LoRA. En faisant abstraction des formats de stockage et de la gestion de la mémoire, BrainSurgery exécute des transformations complexes via des plans déclaratifs en YAML. Il prend en charge les modifications structurelles, les transformations mathématiques et le remodelage des tenseurs grâce à un ciblage structurel et par expressions régulières expressif, tandis que des assertions intégrées valident les formes, types de données et valeurs des tenseurs pour éviter les erreurs silencieuses. Nous envisageons que BrainSurgery fournira une base solide pour la recherche future grâce à ses opérations reproductibles et validées.

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.