BYOL对比学习论文精读：从理论到实践的新手指南

1次阅读

共计 3928 个字符，预计需要花费 10 分钟才能阅读完成。

自监督学习近年来成为机器学习领域的热门方向，它允许模型从未标注的数据中自动学习有用的特征表示。传统的对比学习方法（如 SimCLR）通过将同一图像的不同增强视图（正样本）拉近，同时将不同图像的视图（负样本）推远来学习特征。这种方法虽然有效，但也存在明显的局限性：

需要大量的负样本才能获得良好的性能
负样本的存储和计算成本高昂
负样本的质量直接影响模型效果

BYOL（Bootstrap Your Own Latent）的创新之处在于它完全摆脱了对负样本的依赖，仅通过正样本间的对比就能学习到高质量的特征表示。这对于计算资源有限的新手来说是个好消息。

BYOL 的巧妙设计主要体现在三个关键组件上：

BYOL 使用两个结构相同但参数不同的网络：

Online 网络：包含编码器 fθ、预测头 gθ 和投影头 qθ
Target 网络：包含编码器 fξ 和投影头 qξ

这种双网络设计是 BYOL 能够避免崩溃（即所有输出都收敛到相同值）的关键。

Online 网络中的预测头 gθ 是一个小型 MLP，其作用是将 online 网络的特征空间映射到 target 网络的特征空间。这个看似简单的组件实际上是防止模型崩溃的关键：

强制 online 网络需要 ” 预测 ”target 网络的特征
打破了对称性，防止两个网络互相模仿

Target 网络的参数 ξ 不是通过梯度下降直接更新，而是通过 online 网络参数 θ 的指数移动平均 (EMA) 更新：

ξ ← τξ + (1-τ)θ

其中 τ 是 EMA 系数（通常接近 1，如 0.99）。这种缓慢更新确保了 target 网络比 online 网络更稳定，提供了更可靠的学习目标。

下面我们实现一个简化版的 BYOL 模型，只保留核心组件以便理解：

import torch
import torch.nn as nn
import torch.nn.functional as F

class MLPHead(nn.Module):
    def __init__(self, in_dim, hidden_dim=4096, out_dim=256):
        super().__init__()
        self.layer1 = nn.Sequential(nn.Linear(in_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(inplace=True)
        )
        self.layer2 = nn.Linear(hidden_dim, out_dim)

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        return x

class BYOL(nn.Module):
    def __init__(self, backbone, hidden_dim=4096, out_dim=256, tau=0.99):
        super().__init__()
        self.backbone = backbone
        self.online_projector = MLPHead(backbone.output_dim, hidden_dim, out_dim)
        self.online_predictor = MLPHead(out_dim, hidden_dim, out_dim)

        # 初始化 target 网络
        self.target_backbone = backbone
        self.target_projector = MLPHead(backbone.output_dim, hidden_dim, out_dim)

        # 冻结 target 网络参数
        for param in self.target_backbone.parameters():
            param.requires_grad = False
        for param in self.target_projector.parameters():
            param.requires_grad = False

        self.tau = tau

    @torch.no_grad()
    def update_target_network(self):
        # EMA 更新 target 网络
        for online_param, target_param in zip(self.backbone.parameters(), self.target_backbone.parameters()):
            target_param.data = self.tau * target_param.data + (1 - self.tau) * online_param.data
        for online_param, target_param in zip(self.online_projector.parameters(), self.target_projector.parameters()):
            target_param.data = self.tau * target_param.data + (1 - self.tau) * online_param.data

    def forward(self, view1, view2):
        # online 网络处理第一个视图
        online_z1 = self.backbone(view1)
        online_proj1 = self.online_projector(online_z1)
        online_pred1 = self.online_predictor(online_proj1)

        # online 网络处理第二个视图
        online_z2 = self.backbone(view2)
        online_proj2 = self.online_projector(online_z2)
        online_pred2 = self.online_predictor(online_proj2)

        # target 网络处理
        with torch.no_grad():
            target_z1 = self.target_backbone(view2)
            target_proj1 = self.target_projector(target_z1)
            target_proj1 = F.normalize(target_proj1, dim=1)

            target_z2 = self.target_backbone(view1)
            target_proj2 = self.target_projector(target_z2)
            target_proj2 = F.normalize(target_proj2, dim=1)

        # 计算对称损失
        loss1 = 2 - 2 * (online_pred1 * target_proj1).sum(dim=1).mean()
        loss2 = 2 - 2 * (online_pred2 * target_proj2).sum(dim=1).mean()

        return (loss1 + loss2) / 2

BYOL 的性能很大程度上依赖于高质量的数据增强。我们通常使用以下增强组合：

from torchvision import transforms

train_transform = transforms.Compose([transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.RandomApply([transforms.ColorJitter(0.4, 0.4, 0.2, 0.1)], p=0.8),
    transforms.RandomGrayscale(p=0.2),
    transforms.RandomApply([transforms.GaussianBlur(kernel_size=5)], p=0.5),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])