BrainSurgery: モデル編集とアップサイクリングのための再現可能で信頼性のある宣言的な重み操作

要旨

ディープラーニングモデルの規模が拡大するにつれ、大規模なチェックポイントの管理、検査、修正はますます困難になっています。研究者はしばしば、層の再構築、精度変換、低ランク分解、アーキテクチャのデバッグのためにモデルの重みを変更する必要がありますが、これらのワークフローは脆弱なアドホックなPythonスクリプトに依存することがよくあります。ここでは、ニューラルネットワークのチェックポイントに対するロバストで再現可能な「テンソル手術」を実現するツールBrainSurgeryを紹介し、モデルアップサイクルからLoRA抽出までの4つの例と3つのケーススタディを含むシステムデモを提供します。ストレージ形式とメモリ管理を抽象化することにより、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.