Barlow Twins 损失函数解析：从原理到自监督学习实践

1次阅读

没有评论

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

在自监督学习中，模型需要从未标注的数据中学习有意义的特征表示。传统对比学习方法（如 SimCLR、MoCo）通过拉近正样本对、推开负样本对来实现这一目标。然而，这种方法存在一个显著问题：特征冗余。

特征冗余指的是学习到的特征表示中，不同维度之间存在高度相关性，导致信息重复。
这不仅降低了特征的表达能力，还浪费了模型容量。
在极端情况下，模型可能仅依赖少数几个特征维度而忽略其他维度。

Barlow Twins 提出了一种新颖的思路：直接最小化特征之间的互相关性。其损失函数由两部分组成：

对角一致性（Diagonal Consistency）：鼓励特征在相同维度上保持一致
非对角冗余（Off-Diagonal Redundancy）：惩罚不同维度之间的相关性

数学上，损失函数定义为：

$$
\mathcal{L}{BT} = \sum_i (1 – C^2
$$})^2 + \lambda \sum_{i \neq j} C_{ij

其中 $C$ 是标准化后的互相关矩阵，$\lambda$ 是权衡系数。

以下是完整的 Barlow Twins 实现，包含数据增强和损失计算：

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

class BarlowTwinsLoss(nn.Module):
    """
    Barlow Twins 损失函数的 PyTorch 实现
    参数:
        lambda_param (float): 权衡系数，控制冗余惩罚的强度
        feature_dim (int): 特征向量的维度
    """
    def __init__(self, lambda_param=5e-3, feature_dim=128):
        super().__init__()
        self.lambda_param = lambda_param
        self.bn = nn.BatchNorm1d(feature_dim, affine=False)

    def forward(self, z1, z2):
        """
        计算 Barlow Twins 损失
        参数:
            z1, z2 (Tensor): 两个增强视图的特征表示，形状为 (batch_size, feature_dim)
        返回:
            loss (Tensor): 计算得到的损失值
        """
        # 批归一化（零中心化）z1_norm = self.bn(z1)
        z2_norm = self.bn(z2)

        # 计算互相关矩阵
        batch_size = z1.size(0)
        c = torch.mm(z1_norm.T, z2_norm) / batch_size  # (feature_dim, feature_dim)

        # 计算损失
        on_diag = torch.diagonal(c).add_(-1).pow_(2).sum()
        off_diag = off_diagonal(c).pow_(2).sum()
        loss = on_diag + self.lambda_param * off_diag

        return loss

def off_diagonal(x):
    """返回矩阵中所有非对角元素的视图"""
    n, m = x.shape
    assert n == m
    return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()