How JEPA Avoids Noisy Features: The Implicit Bias of Deep Linear Self Distillation Networks

📄 arXiv: 2407.03475v1 📥 PDF

作者: Etai Littwin, Omid Saremi, Madhu Advani, Vimal Thilak, Preetum Nakkiran, Chen Huang, Joshua Susskind

分类: cs.LG

发布日期: 2024-07-03

备注: Technical report

💡 一句话要点

通过线性自蒸馏网络隐式偏置，JEPA避免噪声特征干扰

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 自监督学习 联合嵌入预测架构 自蒸馏 隐式偏置 深度线性模型

📋 核心要点

  1. 自监督学习面临如何有效提取数据中抽象、鲁棒特征的挑战，避免模型过度关注噪声或不相关细节。
  2. 本文分析了JEPA架构的自蒸馏方法，发现其在线性模型中存在学习高影响力特征的隐式偏置。
  3. 研究结果表明，在潜在空间进行预测的JEPA方法，能够有效避免噪声特征，从而提升表征学习的性能。

📝 摘要（中文）

本文研究了自监督学习中表征学习的两种范式：联合嵌入预测架构(JEPA)和掩码自编码器(MAE)。JEPA通过预测语义相似输入的表征来学习，例如自蒸馏方法，在线编码器被训练来预测目标编码器的输出。与MAE在数据空间重建缺失输入不同，JEPA被认为更关注抽象特征而非细粒度的像素信息。本文通过分析深度线性模型的训练动态，揭示了JEPA避免噪声特征的机制：在线性模型中，JEPA偏向于学习高影响力的特征，即具有高回归系数的特征。该研究结果表明，在潜在空间中进行预测具有独特的隐式偏置，这可能解释了JEPA在实践中的成功。

🔬 方法详解

问题定义：自监督学习旨在学习有用的数据表征，但如何避免模型学习到噪声或不相关的低级特征是一个关键问题。现有的MAE方法在像素空间进行重建，容易受到像素级噪声的影响。JEPA架构，特别是自蒸馏方法，在潜在空间进行预测，经验上表现出对抽象特征的偏好，但其内在机制尚不明确。

核心思路：本文的核心思路是通过分析深度线性模型的训练动态，揭示JEPA架构避免噪声特征的机制。作者假设，在线性模型中观察到的现象可以推广到更复杂的非线性模型，从而解释JEPA在实践中的成功。通过简化模型，可以更容易地分析其隐式偏置。

技术框架：本文主要关注深度线性自蒸馏网络。该网络包含一个在线编码器和一个目标编码器。在线编码器被训练来预测目标编码器的输出。目标编码器通常是在线编码器的指数移动平均(EMA)。研究人员分析了该网络在训练过程中的权重变化，特别是不同特征的回归系数。

关键创新：本文最重要的技术创新点是揭示了JEPA架构在线性模型中存在学习高影响力特征的隐式偏置。这意味着JEPA倾向于学习那些对预测结果影响更大的特征，而忽略那些影响较小的特征，从而避免了噪声特征的干扰。这种隐式偏置是JEPA在潜在空间进行预测的固有属性。

关键设计：本文的关键设计在于使用深度线性模型进行分析，这使得研究人员能够更容易地推导出训练动态的解析解。此外，作者还关注特征的回归系数，将其作为衡量特征影响力的指标。通过分析回归系数的变化，作者能够揭示JEPA架构的隐式偏置。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

本文通过对深度线性自蒸馏网络的分析，揭示了JEPA架构存在学习高影响力特征的隐式偏置。该研究表明，在潜在空间进行预测可以有效避免噪声特征的干扰。虽然是在线性模型上进行的分析，但其结论可能推广到更复杂的非线性模型，为理解JEPA在实践中的成功提供了新的视角。

🎯 应用场景

该研究成果有助于更好地理解自监督学习算法的内在机制，并为设计更有效的自监督学习方法提供理论指导。潜在应用领域包括计算机视觉、自然语言处理等，可以提升图像分类、目标检测、语义分割等任务的性能。此外，该研究还可以应用于机器人学习，帮助机器人学习更鲁棒的环境表征。

📄 摘要（原文）

Two competing paradigms exist for self-supervised learning of data representations. Joint Embedding Predictive Architecture (JEPA) is a class of architectures in which semantically similar inputs are encoded into representations that are predictive of each other. A recent successful approach that falls under the JEPA framework is self-distillation, where an online encoder is trained to predict the output of the target encoder, sometimes using a lightweight predictor network. This is contrasted with the Masked AutoEncoder (MAE) paradigm, where an encoder and decoder are trained to reconstruct missing parts of the input in the data space rather, than its latent representation. A common motivation for using the JEPA approach over MAE is that the JEPA objective prioritizes abstract features over fine-grained pixel information (which can be unpredictable and uninformative). In this work, we seek to understand the mechanism behind this empirical observation by analyzing the training dynamics of deep linear models. We uncover a surprising mechanism: in a simplified linear setting where both approaches learn similar representations, JEPAs are biased to learn high-influence features, i.e., features characterized by having high regression coefficients. Our results point to a distinct implicit bias of predicting in latent space that may shed light on its success in practice.