LeJEPA：论文与代码综述

论文：《LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics》（arXiv:2511.08544，Randall Balestriero & Yann LeCun）
代码：本仓库将论文中的 Sketched Isotropic Gaussian Regularization (SIGReg) 实例化为 PyTorch 模块，可直接嵌入任意自监督预训练框架。

1. 论文要点速览

• SIGReg 理论基础：约束模型嵌入分布接近最优的各向同性高斯，从而最小化下游预测风险；论文给出显式上界，且只需单个正则权重 $\lambda$ 调节表征质量与收敛速度。
• 去除常见启发式：无 stop-gradient、无 teacher-student EMA、无动量编码器或预测头，纯粹依靠统计约束即可稳定训练 ViT、ConvNeXtV2 等大模型。
• 线性时间 + 内存：SIGReg 的随机切片 + 单变量统计测试仅需一次矩阵乘，与 batch size 与维度线性关系，适合 8×H200 等大规模训练。
• 统一多数据集表现：论文在 DTD、Aircraft、Cars、CIFAR、Flowers、Food、Pets 等 10 余个数据集上比较 I-JEPA，给出 1-shot / 10-shot / full-shot 结果，LeJEPA 在 ImageNet-1K 预训练 100 epoch 后已与 300 epoch I-JEPA 旗鼓相当。

README 中的 Key Features 概要也与上述对应，可用于 README ↔ 论文/代码的快速对照：

• Single trade-off hyperparameter → scripts/launch_*.md 中仅需设置 bstat_lambda。
• Linear time/memory complexity → lejepa/multivariate/slicing.py 的随机切片实现。
• Heuristics-free implementation → README 的默认配置关闭 predictor / teacher / scheduler trick。
• Distributed training-friendly (~50 LOC) → 通过共享 global_step 和 torch.distributed all_reduce 保持多卡确定性。

2. SIGReg：理论直观与数学推导

SIGReg (Sketched Isotropic Gaussian Regularization) 的核心直觉是：**与其通过对比损失（Contrastive Loss）去“拉近正例、推远负例”，不如直接让所有数据的特征分布塌缩成一个各向同性的高斯分布（Isotropic Gaussian）。**

2.1 为什么是高斯？

论文证明，如果特征空间 $Z$ 的分布 $p(z)$ 满足 $p(z) \propto e^{-\|z\|^2/2}$（标准正态），则下游线性分类任务的风险存在显式上界。即：特征分布越接近各向同性高斯，线性可分性越好。

2.2 降维打击：随机切片 (Slicing)

在高维空间直接计算分布距离（如 KL 散度或 Wasserstein 距离）非常困难。LeJEPA 利用 Radon 变换定理：如果一个高维分布在所有 1D 投影（切片）上都是标准正态分布，那么它本身就是标准多元正态分布。

2.3 距离度量：Epps-Pulley 检验

论文中默认使用 EppsPulley 统计量作为 $\text{Dist}$，它是基于特征函数（Characteristic Function）的距离。相比 KS 或 AD 检验，它更适合梯度下降优化。

$$ T_{EP} = \int \left| \frac{1}{N} \sum_{j=1}^N e^{i t y_j} - e^{-t^2/2} \right|^2 w(t) dt $$

代码实现（lejepa/univariate/epps_pulley.py）通过离散化积分点 $t$ 来高效计算该损失。

3. Demo 与性能速览

README 中的 GIF / 图像与性能表也搬运到此，方便直接浏览：

4. 预训练数据流与配置

根据 README.md 与 scripts/launch_*.md，LeJEPA 的输入管线遵循 DINO 风格的多视角增强，随后由 backbone + projector 输出嵌入，再通过 SIGReg 约束。

2.1 视角生成（多 Crop）

2.2 Backbone + Projector

• 支持 TIMM 中 60+ 种骨干（ViT-L/H、ConvNeXtV2-H、轻量 ViT 等），脚本 scripts/launch_inet10.py 通过 stable_pretraining 批量扫描。
• 投影头可配置为 MLP 或线性层，典型设置：embedding_dim = 512、projector_dim = 512。
• 多视角输出会被拼接/平均成 batch × tokens × dim 的嵌入张量，直接送入 SIGReg。

2.3 优化器与调度

• 优化器：AdamW；ViT 使用 weight_decay=5e-2，ResNet 使用 5e-4。
• 学习率：起始 5e-4，线性 warmup 后采用 cosine 衰减，终值为初始的 /1000。
• 精度：默认 bfloat16（bf16）混合精度以匹配 H100/H200 的原生指令。
• 视角策略：可在 Hydra 中切换是否使用 multi-crop、predictor、teacher-student 等选项；论文默认全部关闭，以凸显 SIGReg 的“无启发式”特性。

5. SIGReg 损失在代码中的实现

核心思想：将高维嵌入 $X \in \mathbb{R}^{N \times D}$ 投影到大量随机 1D 切片上，应用单变量统计检验，使每个切片都符合标准正态，从而逼近各向同性高斯。

5.1 随机切片 + 分布同步

from torch import nn
from torch.distributed._functional_collectives import all_reduce as functional_all_reduce

def all_reduce(x, op="AVG"):
    if dist.is_available() and dist.is_initialized():
        return functional_all_reduce(x, op.lower(), dist.group.WORLD)
    return x

class SlicingUnivariateTest(nn.Module):
    def __init__(self, univariate_test, num_slices, reduction="mean",
                 sampler="gaussian", clip_value=None):
        super().__init__()
        self.univariate_test = univariate_test
        self.num_slices = num_slices
        self.reduction = reduction
        self.clip_value = clip_value
        self.register_buffer("global_step", torch.zeros((), dtype=torch.long))
        self._generator = None
        self._generator_device = None

    def forward(self, x):
        with torch.no_grad():
            seed = all_reduce(self.global_step.clone(), op="MAX").item()
            g = self._get_generator(x.device, seed)
            proj = torch.randn((x.size(-1), self.num_slices),
                               device=x.device, generator=g)
            proj /= proj.norm(p=2, dim=0)
            self.global_step.add_(1)

        stats = self.univariate_test(x @ proj)
        if self.clip_value is not None:
            stats = stats.clamp_min(self.clip_value)
        return stats.mean() if self.reduction == "mean" else stats

要点：

• 每次前向都会同步 global_step 并复用同一随机数生成器，确保多 GPU 推理得到完全一致的投影矩阵 $A$。
• 投影后得到形状 (*, N, num_slices) 的 1D 样本，再交给下游单变量检验。

5.2 单变量 Epps–Pulley 统计

class EppsPulley(UnivariateTest):
    def __init__(self, t_max=3.0, n_points=17):
        super().__init__()
        t = torch.linspace(0, t_max, n_points)
        dt = t_max / (n_points - 1)
        weights = torch.full((n_points,), 2 * dt)
        weights[[0, -1]] = dt
        self.register_buffer("t", t)
        self.register_buffer("phi", (-0.5 * t**2).exp())
        self.register_buffer("weights", weights * self.phi)

    def forward(self, x):
        N = x.size(-2)
        x_t = x.unsqueeze(-1) * self.t
        cos_mean = torch.cos(x_t).mean(-3)
        sin_mean = torch.sin(x_t).mean(-3)
        cos_mean = self.dist_mean(cos_mean)
        sin_mean = self.dist_mean(sin_mean)
        err = (cos_mean - self.phi).square() + sin_mean.square()
        return (err @ self.weights) * N * self.world_size

该检验比较经验特征函数与标准正态特征函数的差异，返回正标量作为“偏离高斯”的能量，正好可直接作为损失。

5.3 其他统计（BHEP 示例）

class BHEP(MultivariatetTest):
    def __init__(self, beta=0.1):
        super().__init__()
        if beta <= 0:
            raise ValueError("beta must be positive")
        self.beta = beta

    def forward(self, x):
        x = self.prepare_data(x)
        N, D = x.shape
        beta2 = self.beta ** 2
        sq_norm = x.square().sum(dim=1)
        dist = -2 * x @ x.T + sq_norm[:, None] + sq_norm[None, :]
        lhs = torch.exp(-0.5 * beta2 * dist).sum() / (N ** 2)
        scaling = 2 / ((1 + beta2) ** (D / 2))
        rhs = scaling * torch.exp(sq_norm * (-beta2 / (2 + 2 * beta2))).sum() / N
        constant = 1 / ((1 + 2 * beta2) ** (D / 2))
        return lhs - rhs + constant

仓库还提供 Henze–Zirkler、Cramér–von Mises、Anderson–Darling、Shapiro–Wilk 等多种检验，可随意替换，以适配不同数据分布或灵敏度。

5.4 将 SIGReg 接入训练循环

import lejepa

backbone = build_model(...)
projector = build_projector(...)

univar = lejepa.univariate.EppsPulley(t_max=5.0, n_points=21)
sigreg = lejepa.multivariate.SlicingUnivariateTest(
    univariate_test=univar,
    num_slices=1024,
    reduction="mean",
)

def training_step(images):
    views = augment(images)         # 2 global + 6 local
    embeddings = projector(backbone(views))  # [B, tokens, dim]
    loss = sigreg(embeddings)
    loss.backward()
    optimizer.step()

（上例与 README 保持一致，可直接复制到 Lightning / PyTorch 训练脚本中。）

6. Ablation Study (消融实验)

基于 scripts/launch_*.md 和论文数据，以下是关键超参数对模型性能的影响总结。

6.1 投影头与切片数量 (launch_proj_ablation.md)

6.2 视图策略与 Batch Size (launch_views_ablation.md)

注意：LeJEPA 发现 SIGReg 使得大 Batch Size (如 4096+) 训练极其稳定，不像对比学习方法那样需要精细调整 Learning Rate Scaling。

6.3 正则化强度 $\lambda$

bstat_lambda 是 LeJEPA 唯一 需要调整的权衡参数（Trade-off）。

• $\lambda$ 太小（e.g., 0.001）：模型退化为单纯的重建/预测任务（如果开启了 predictor），或者特征坍缩。
• $\lambda$ 太大（e.g., 10.0）：特征被强行拉成高斯噪声，丢失图像语义信息。
• 推荐范围： 0.01 ~ 0.1（ImageNet 任务通常用 0.05）。

7. 优势对比：LeJEPA vs DINO / I-JEPA / MAE

总结： LeJEPA 用统计学约束替代了工程上的 Trick（如 EMA、Stop-gradient），在保持高性能的同时，大幅降低了调参门槛和显存需求。

8. 推理与评估流程

• 表征提取：冻结预训练模型，取最后两层的 CLS token（无 CLS 时平均 patch token）并拼接。
• 归一化：对拼接特征做 LayerNorm/BatchNorm，可提升线性探针稳定性。
• 线性探针：使用 AdamW，学习率与预训练同样采用 warmup + cosine；weight decay 仅 1e-6。
• 可视化：eval/ 文件夹提供原图与 PCA 回放、GIF 动画，帮助验证表示是否捕获全局语义。

9. 应用指南：迁移与调优

9.1 迁移到下游任务 (检测/分割)

LeJEPA 预训练的 Backbone 可以像 MAE 或 DINO 那样直接加载到 Detectron2 或 MMSegmentation 中。

1. 提取权重： 使用 torch.load 读取 checkpoint，提取 backbone. 开头的权重。
2. 移除 Projector： 丢弃 MLP Projector 和 SIGReg 模块，它们只在预训练时用于约束分布。
3. LayerNorm 缩放： 由于 LeJEPA 强约束特征为 $N(0,1)$，下游任务初始化时建议不对 Backbone 输出再做大幅度的 Scale/Shift 操作，或者使用学习率 warmup。

9.2 适配自定义数据集

如果你要在非 ImageNet 数据（如医疗影像、遥感图像）上训练 LeJEPA：

• 小数据集 (N < 50k)： 增加 bstat_lambda (0.1 ~ 0.5)。数据少时更容易过拟合，需要更强的正则化防止塌缩。
• 大数据集 (N > 1M)： 减小 bstat_lambda (0.01 ~ 0.05)。数据本身的多样性已经提供了很好的约束。
• 非自然图像： 调整 Augmentation。医疗影像可能不适合 RandomColorJitter，应改为亮度/对比度微调；遥感图像可增加 RandomRotation。

10. 分布式与确定性细节

• 所有统计量都使用 torch.distributed 的 all_reduce 聚合，确保 batch 拆分后仍等价于单机结果。
• SlicingUnivariateTest 用共享 global_step 控制随机投影，避免不同 GPU 采到不同切片导致噪声。
• 代码完全在 torch.no_grad() 下生成切片和积分点，不会额外占用显存；只需反向传播单个标量 loss。

11. 安装与快速上手

无需查阅 README 即可在此获取安装与最小示例：

• 依赖：Python ≥ 3.8、PyTorch ≥ 1.10、NumPy，可选 stable_pretraining 以使用仓库脚本。

pip install lejepa stable-pretraining hydra-core

最小 SIGReg 使用示例：

import lejepa

univariate_test = lejepa.univariate.EppsPulley(num_points=17)
loss_fn = lejepa.multivariate.SlicingUnivariateTest(
    univariate_test=univariate_test,
    num_slices=1024,
    reduction="mean"
)

embeddings = backbone_outputs  # 假设 shape = [batch, dim]
loss = loss_fn(embeddings)
loss.backward()

12. 复现步骤建议

1. 环境：pip install lejepa stable-pretraining hydra-core，确保 PyTorch≥1.10、Python≥3.8。
2. 数据：准备 ImageNet-1K 或更小的数据集（如 INet-10），按 README 的 multi-crop 规则构建 DataLoader。
3. 配置：参考 scripts/launch_inet10.py 中的 Hydra 参数（bstat_name、bstat_lambda、num_slices 等），按需要枚举骨干。
4. 训练：关闭 predictor/teacher 相关 flag，启用 bf16、AdamW、cosine 调度；记录 loss_sigreg 即可监控训练。
5. 评估：按 README 的线性探针流程在 DTD、Aircraft、CIFAR、Flowers、Food、Pets 等数据集上对比 I-JEPA / DINO。

具体的 Hydra 启动示例（与 README 中保持一致）如下：

HYDRA_FULL_ERROR=1 python scripts/je.py \
  --config-dir scripts/configs \
  --config-name base \
  +accelerator=single_gpu_sc \
  ++dataset_name="inet10" \
  ++bstat_name="epps_pulley" \
  ++bstat_num_slices=1000 \
  ++bstat_lambda=0.05 \
  ++embedding_dim=512 \
  ++projector_dim=512 \
  ++lr=3e-3 \
  ++weight_decay=3e-2 \
  ++n_views=8 \
  ++teacher_student=false \
  ++predictor=false

该命令直接复用了 README 的建议参数，可根据需要修改 backbone/节点数/视角策略。

13. 结语与展望

LeJEPA 证明了“好的表征 = 各向同性高斯”这一简单假设的强大威力。它不仅在 ImageNet 上取得了 SOTA 级别的性能，更重要的是，它为自监督学习提供了一个可解释、可证明、无 Trick 的新范式。

未来方向包括：

• 多模态扩展： 音频、视频、文本的 Embedding 同样可以被 SIGReg 约束。
• 生成模型结合： 将 LeJEPA 学习到的结构化 Latent Space 用于 Diffusion Model 的条件输入（如 JiT 项目）。