BYOL对比学习原理解析与自监督学习实践指南

1次阅读

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

自监督学习近年来在计算机视觉领域取得了显著进展，它通过从数据本身生成监督信号来训练模型，避免了大量人工标注的成本。对比学习（Contrastive Learning）是自监督学习中的一种重要方法，它通过拉近正样本对、推开负样本对的方式来学习有效的特征表示。然而，传统的对比学习方法（如 SimCLR）严重依赖负样本，这在实际应用中带来了一些挑战：

需要大量负样本才能保证学习效果
负样本的存储和计算成本高昂
负样本的质量直接影响模型性能

BYOL（Bootstrap Your Own Latent）提出了一种全新的思路：完全不需要负样本，仅通过在线网络（online network）和目标网络（target network）的协同训练就能学习到高质量的特征表示。这一创新不仅降低了计算成本，还提高了特征学习的稳定性。

传统对比学习方法通常需要：

构建正样本对（同一图像的不同增强视图）
构建负样本池（来自不同图像的样本）
使用对比损失（如 InfoNCE）拉近正样本对、推开负样本

BYOL 则完全抛弃了负样本，其核心思想是：让一个网络（在线网络）学习预测另一个网络（目标网络）的特征表示。这种自我引导的学习方式避免了特征坍塌（collapse）问题。

BYOL 包含两个结构相同但参数不同的网络：

在线网络（θ）：包含编码器 f_θ、投影头 g_θ 和预测头 q_θ
目标网络（ξ）：包含编码器 f_ξ 和投影头 g_ξ

训练过程如下：

对输入图像 x 生成两个随机增强视图 v 和 v ’
在线网络处理 v，得到表示 y_θ = g_θ(f_θ(v)) 和预测 z_θ = q_θ(y_θ)
目标网络处理 v ’，得到表示 y ’_ξ = g_ξ(f_ξ(v’))
最小化 L2 归一化后的预测和目标的 MSE 损失
通过 EMA（指数移动平均）更新目标网络参数

数学表达式为：

\mathcal{L}_{\theta,\xi} = \|\bar{z}_\theta - \bar{y}'_\xi\|_2^2 = 2 - 2 \cdot \frac{\langle z_\theta, y'_\xi \rangle}{\|z_\theta\|_2 \cdot \|y'_\xi\|_2}

预测头 q_θ 是 BYOL 防止特征坍塌的关键组件。如果没有预测头，两个网络可能会收敛到相同的平凡解（如输出常数）。预测头强制在线网络必须进行非线性变换才能匹配目标网络的输出。

EMA（Exponential Moving Average）更新策略让目标网络的参数变化更加平滑：

\xi \leftarrow \tau \xi + (1-\tau)\theta

其中 τ 通常设置为 0.99-0.999，这种缓慢更新确保了目标网络提供相对稳定的学习目标。

以下是 BYOL 的 PyTorch 实现核心代码：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms

# 编码器网络（以 ResNet18 为例）class Encoder(nn.Module):
    def __init__(self, backbone):
        super().__init__()
        self.backbone = backbone
        self.projection = nn.Sequential(nn.Linear(512, 1024),
            nn.BatchNorm1d(1024),
            nn.ReLU(),
            nn.Linear(1024, 256)
        )

    def forward(self, x):
        h = self.backbone(x)
        return self.projection(h)

# 预测头网络
class Predictor(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(256, 1024),
            nn.BatchNorm1d(1024),
            nn.ReLU(),
            nn.Linear(1024, 256)
        )

    def forward(self, x):
        return self.net(x)

# BYOL 完整模型
class BYOL(nn.Module):
    def __init__(self, backbone):
        super().__init__()
        self.online_encoder = Encoder(backbone)
        self.online_predictor = Predictor()

        # 目标网络初始与在线网络相同但不更新梯度
        self.target_encoder = Encoder(backbone)
        for param in self.target_encoder.parameters():
            param.requires_grad = False

        # 数据增强
        self.augment = transforms.Compose([transforms.RandomResizedCrop(32),
            transforms.RandomHorizontalFlip(),
            transforms.ColorJitter(0.4, 0.4, 0.4, 0.1),
            transforms.RandomGrayscale(p=0.2),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], 
                                std=[0.2023, 0.1994, 0.2010])
        ])

    def forward(self, x):
        # 生成两个增强视图
        v1 = self.augment(x)
        v2 = self.augment(x)

        # 在线网络处理
        y1 = self.online_encoder(v1)
        z1 = self.online_predictor(y1)

        # 目标网络处理
        with torch.no_grad():
            y2 = self.target_encoder(v2)

        # L2 归一化
        z1_norm = F.normalize(z1, dim=1)
        y2_norm = F.normalize(y2, dim=1)

        # 计算 MSE 损失
        loss = 2 - 2 * (z1_norm * y2_norm).sum(dim=1).mean()
        return loss

    def update_target(self, tau=0.99):
        # EMA 更新目标网络
        with torch.no_grad():
            for online, target in zip(self.online_encoder.parameters(), 
                                     self.target_encoder.parameters()):
                target.data = tau * target.data + (1 - tau) * online.data

BYOL 对学习率非常敏感，推荐使用余弦退火调度器：

from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR

model = BYOL(backbone)
optimizer = AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
scheduler = CosineAnnealingLR(optimizer, T_max=100, eta_min=1e-5)

理想批量大小为 4096，但在有限显存下可降至 256-512
使用梯度累积模拟大批量：

batch_size = 256
grad_accum = 16  # 实际等效批量 =256*16=4096

for epoch in range(100):
    for i, (x, _) in enumerate(train_loader):
        loss = model(x)
        loss.backward()

        if (i + 1) % grad_accum == 0:
            optimizer.step()
            optimizer.zero_grad()
            model.update_target()

使用 UMAP 或 t -SNE 可视化学习到的特征：

import umap
import matplotlib.pyplot as plt

# 提取特征
features = []
labels = []
with torch.no_grad():
    for x, y in test_loader:
        h = model.target_encoder.backbone(x)
        features.append(h)
        labels.append(y)

features = torch.cat(features).numpy()
labels = torch.cat(labels).numpy()

# UMAP 降维
reducer = umap.UMAP()
embedding = reducer.fit_transform(features)

# 可视化
plt.scatter(embedding[:,0], embedding[:,1], c=labels, cmap='Spectral', s=1)
plt.colorbar()
plt.show()