Git Re-Basin：在置换对称性下合并模型

本文代码：Git Re-Basin：在置换对称性下合并模型。

摘要：

深度学习的成功得益于我们能够相对轻松地解决某些大规模的非凸优化问题。尽管非凸优化问题是NP难的，但简单的算法——通常是随机梯度下降的变体——在实践中却表现出惊人的有效性，能够很好地拟合大型神经网络。我们提出，当考虑隐藏单元的所有可能置换对称性后，神经网络的损失景观中实际上只存在（几乎）一个盆地。为此，我们提出了三种算法，用于对齐一个模型与参考模型之间的隐藏单元排列顺序。通过这种变换，我们可以得到一组功能等价的权重参数，它们位于靠近参考模型的一个近似凸的盆地内。实验结果表明，这一“单盆地”现象在多种模型架构和数据集上均成立，其中包括首次（据我们所知）在CIFAR-10和CIFAR-100数据集上独立训练的ResNet模型之间实现无屏障的线性模式连通性。此外，我们还发现模型宽度和训练时间与模式连通性之间存在一些有趣的关系，并在不同模型和数据集上进行了验证。最后，我们讨论了“单盆地”理论的局限性，包括对线性模式连通性假设的一个反例。

Git Re-Basin 快速上手指南

Git Re-Basin 是一个用于解决神经网络排列对称性问题的工具，能够将独立训练的模型权重进行排列对齐，使其落入同一个损失盆地（Loss Basin），从而实现模型的线性模式连接（Linear Mode Connectivity）和高效合并。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux, macOS 或 Windows (WSL 推荐)
• Python 版本：Python 3.8 或更高版本
• 前置依赖：
  • PyTorch (建议版本 >= 1.9)
  • NumPy
  • Matplotlib (用于可视化)
  • Git

建议创建一个独立的虚拟环境以避免依赖冲突：

python -m venv git-re-basin-env
source git-re-basin-env/bin/activate  # Windows 用户使用: git-re-basin-env\Scripts\activate

安装步骤

您可以直接从 GitHub 源码安装该工具。为了获得更快的下载速度，国内用户可配置 pip 使用清华或阿里云镜像源。

使用默认源安装：

pip install git+https://github.com/giladcohen/git-re-basin.git

推荐使用国内镜像源加速安装：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple git+https://github.com/giladcohen/git-re-basin.git

或者克隆仓库后进行本地安装（适合需要修改源码的开发者）：

git clone https://github.com/giladcohen/git-re-basin.git
cd git-re-basin
pip install -e .
# 国内加速版
# pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -e .

基本使用

本工具的核心功能是通过排列算法将一个模型的权重对齐到参考模型。以下是最简单的 Python 使用示例，演示如何加载两个独立训练的模型并进行对齐。

假设您已经有两个训练好的 PyTorch 模型实例 model_a (参考模型) 和 model_b (待对齐模型)。

import torch
from git_re_basin import match_permutations

# 1. 加载或初始化两个结构相同但权重不同的模型
# model_ref: 参考模型 (Reference Model)
# model_target: 目标模型 (Target Model)，需要被重排列
model_ref = ... 
model_target = ...

# 确保模型处于评估模式
model_ref.eval()
model_target.eval()

# 2. 执行排列匹配算法
# 该函数会返回重新排列后的目标模型权重，使其与参考模型在功能上等价且位于同一盆地
matched_state_dict = match_permutations(
    model_ref.state_dict(),
    model_target.state_dict(),
    model_architecture="resnet" # 根据实际架构选择，如 'mlp', 'resnet', 'vit' 等
)

# 3. 加载对齐后的权重到目标模型
model_target.load_state_dict(matched_state_dict)

# 4. 验证线性模式连接 (可选)
# 现在您可以对两个模型的权重进行线性插值，损失曲面将呈现近似凸性
alpha = 0.5
interpolated_state_dict = {
    key: (1 - alpha) * model_ref.state_dict()[key] + alpha * matched_state_dict[key]
    for key in model_ref.state_dict().keys()
}

# 将插值后的权重加载到新模型中进行测试
model_interpolated = ... # 初始化同架构模型
model_interpolated.load_state_dict(interpolated_state_dict)

注意：match_permutations 函数的具体参数（如 model_architecture）需根据您的网络结构进行调整。对于自定义架构，可能需要参考源码中的 permutation_matching.py 模块实现特定的层映射逻辑。