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

1次阅读

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

自监督学习近年来在计算机视觉领域取得了巨大成功，它通过从数据本身生成监督信号，避免了昂贵的人工标注成本。对比学习作为自监督学习的重要分支，通过拉近相似样本、推开不相似样本来学习特征表示。然而，传统对比学习方法（如 SimCLR、MoCo）严重依赖负样本，这带来了两个主要问题：

需要大量的负样本才能保证学习效果，这导致计算成本高昂
负样本选择不当会导致模型性能下降，即所谓的 ” 负样本陷阱 ”

BYOL（Bootstrap Your Own Latent）的创新之处在于，它完全不需要负样本，仅通过两个网络（online 和 target）的协同学习就能获得优秀的特征表示。

BYOL 的核心思想可以概括为：让 online 网络学会预测 target 网络对同一图像不同增强视图的特征表示。具体实现包含三大关键设计：

双分支网络架构
online 网络：包含编码器、预测头和 projection 头，通过梯度下降更新
target 网络：结构与 online 网络相同（不包括预测头），通过 EMA（指数移动平均）更新
对称预测任务
对同一图像生成两个随机增强视图 v 和 v ’
online 网络预测 target 网络对 v ’ 的特征表示
交换 v 和 v ’ 再做一次预测（对称设计）
EMA 更新机制
target 网络的参数是 online 网络的缓慢更新版本
更新公式：θ_target ← τθ_target + (1-τ)θ_online
τ 通常设置为 0.99-0.999，保证 target 网络稳定

这种设计巧妙地避免了模型坍塌（collapse）问题，即使没有负样本也能学习到有意义的特征表示。

下面是用 PyTorch 实现的简化版 BYOL 核心代码：

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

class MLPHead(nn.Module):
    """BYOL 的预测头和 projection 头"""
    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.996):
        super().__init__()
        self.tau = tau

        # online 网络
        self.online_encoder = backbone
        self.online_projector = MLPHead(backbone.output_dim, hidden_dim, out_dim)
        self.online_predictor = MLPHead(out_dim, hidden_dim, out_dim)

        # target 网络（初始时与 online 相同）self.target_encoder = copy.deepcopy(backbone)
        self.target_projector = copy.deepcopy(self.online_projector)

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

    @torch.no_grad()
    def update_target(self):
        """EMA 更新 target 网络"""
        for online, target in zip(chain(self.online_encoder.parameters(), self.online_projector.parameters()),
            chain(self.target_encoder.parameters(), self.target_projector.parameters())
        ):
            target.data = self.tau * target.data + (1 - self.tau) * online.data

    def forward(self, x1, x2):
        """
        输入：同一图像的两个增强视图 x1, x2
        返回：对称预测损失
        """
        # online 网络处理 x1
        h1 = self.online_encoder(x1)
        z1 = self.online_projector(h1)
        p1 = self.online_predictor(z1)

        # target 网络处理 x2
        with torch.no_grad():
            h2 = self.target_encoder(x2)
            z2 = self.target_projector(h2)
            z2.detach_()

        # 对称预测
        loss = - F.cosine_similarity(p1, z2, dim=-1).mean()

        # 对称处理（交换 x1 和 x2）h2 = self.online_encoder(x2)
        z2 = self.online_projector(h2)
        p2 = self.online_predictor(z2)

        with torch.no_grad():
            h1 = self.target_encoder(x1)
            z1 = self.target_projector(h1)
            z1.detach_()

        loss += - F.cosine_similarity(p2, z1, dim=-1).mean()

        return loss

在实现 BYOL 时，有几个关键超参数需要注意：