Gaussian Embeddings: How JEPAs Secretly Learn Your Data Density

Paper: arXiv:2510.05949 Code: itsanderz/jepa-score-eval（非官方 independent evaluation；作者官方 repo 未找到） Code reference: main @ 40de6866 (2026-03-20) — 非官方实现；论文自身只给出 PyTorch Listing

1. Motivation (研究动机)

这篇论文想回答一个很尖锐的问题：JEPAs 一直被当成“非生成式”的 representation learner，但如果训练目标要求 embedding 服从高熵、近似 Gaussian 的分布，那么 encoder 是否已经被迫学会了输入数据的 density？作者的答案是肯定的：只要 JEPA 成功训练出 Gaussian Embeddings，encoder 的 Jacobian 就必须在输入空间中局部压缩/扩张体积，而这种体积变化正对应数据密度。

当前主流 Self-Supervised Learning 有两条常见解释路径。第一条是 decoder / reconstruction 路径：通过从表示 $f_{\theta }(X)$ 重建 $X$ 来最大化一个 entropy 下界；问题是高维重建会把像素级细节、不可预测噪声和任务无关因素也纳入目标，训练成本和 inductive bias 都重。第二条是 JEPA / contrastive / teacher-student 路径：直接让表示具有 diversity、避免 collapse，同时让两个 view 的表示可互相预测。它的优点恰恰是“不显式建模 $p_X$”，因此通常被视为和 diffusion / score-based / energy-based generative modeling 分离的路线。

本文指出，这种分离并不彻底。JEPA 的 anti-collapse 项不只是防止所有样本映射到同一点；当它把 $f_{\theta }(X)$ 推向 Gaussian 或等价的高熵参考分布时，它实际上规定了输出密度 $p_{f(X)}$。如果输入密度 $p_X$ 非均匀，而输出密度又要变成近似均匀的 hypersphere / Gaussian geometry，那么 $f_{\theta }$ 的 Jacobian 必须在高密度区域和低密度区域做不同的 volume scaling。换句话说，density 没有被 decoder 显式写出来，却被藏在 encoder 的局部微分几何里。

这个问题值得研究有三个原因。第一，它改变了 JEPA 的定位：JEPA 不只是 representation learner，也可以被看成一种非参数 density estimator。第二，它给预训练模型提供了一个无需额外 head、无需重训的 sample-level probability proxy，可用于 data curation、outlier detection、domain readiness assessment。第三，它把 JEPA 和 score-based methods 联系起来：如果能从 encoder Jacobian 取出 density / score 信息，那么原本看似“不生成”的模型也可能支持类似 Langevin sampling 的密度驱动操作。

Figure 1 解读：这张图把 ImageNet class 21 中 JEPA-SCORE 最低的 5 张样本放在左侧、最高的 5 张放在右侧，并跨 DINOv2、I-JEPA-1k、I-JEPA-22k、MetaCLIP 四个预训练 backbone 对齐展示。关键现象不是单个模型能排序，而是不同 JEPA-like 模型倾向于挑出相似的低/高概率样本：低分样本更像非典型姿态、背景或质量异常，高分样本更接近训练分布中常见的“原型”。这正是作者提出 JEPA-SCORE 的直观动机：样本概率似乎已经存在于 pretrained encoder 内部。

2. Idea (核心思想)

核心 insight：如果一个 encoder $f_{\theta }$ 把真实数据 $X\sim p_X$ 映射成 Gaussian Embeddings，那么它必须学会一个把 $p_X$ “推平”为 Gaussian / hypersphere 上近似均匀分布的几何变换；这个变换的局部体积因子由 Jacobian singular values 给出，因此 density 可以从 $J_f(x)$ 闭式读出。

本文的关键创新不是新训练一个 JEPA，而是把已有 JEPA 的 anti-collapse / diversity 目标重新解释为 density learning，并定义 JEPA-SCORE：对输入 $x$，计算 encoder Jacobian $J_f(x)$ 的奇异值 ${\sigma }_k$，把 ${\sum }_k\log {\sigma }_k(J_f(x))$ 作为 log-scale 的 learned density proxy。与 diffusion/score matching 不同，它不需要训练 denoiser 或显式 score network；与 reconstruction model 不同，它也不需要 decoder 回到像素空间。

和典型 generative density estimator 的根本差异在于，JEPA 不直接最大化 $\log p_X(x)$，也不参数化一个可归一化的输入空间密度。它只约束 latent distribution 和 view-level predictive invariance。论文证明：在这些约束最优时，输入 density 会通过 encoder 的 local stretching/compression 隐含进入 Jacobian。也就是说，JEPA-SCORE 是一个“从已经训练好的 representation 中提取 density”的方法，而不是一个新的 density-training objective。

3. Method (方法)

3.1 总体框架：Gaussian Embeddings 迫使 encoder 学 density

论文方法可以拆成三层逻辑。第一层是高维 Gaussian 的几何事实：$K$ 维标准 Gaussian 经过 $1/\sqrt{K}$ 缩放后，norm 会集中到半径 1 附近，而方向近似均匀，因此可近似看成 hypersphere surface 上的 uniform distribution。第二层是 change-of-variable / coarea 公式：若 $f$ 把 $X$ 映射到某个指定的输出密度 $p_{f(X)}$，那么 $p_{f(X)}$ 与 $p_X$ 之间由 $J_f$ 的奇异值乘积联系。第三层才是 JEPA：anti-collapse/diversity 项要求 embedding 分布高熵、接近 Gaussian；view prediction/invariance 项则让同一 generator 的 augmentation 落在相同或相近 level-set 上。两者合起来，使 generator density $p_{\mu }$ 可由 encoder 的 Jacobian 恢复。

Figure 2 解读：左上给出几何直觉：输入空间里高密度区域若要被映射到 hypersphere 上近似 uniform 的 embedding 分布，encoder 必须在这些区域做更强的局部压缩/扩张调节；Jacobian determinant 或 singular values 的乘积就记录了这种体积变换。右上表格是 controlled GMM 实验，展示 JEPA-SCORE 与真实 $\log p(x)$ 的 Pearson correlation；底部用 Langevin dynamics 显示，即使 JEPA-SCORE 只是近似 score，也能逐步从粗糙分布恢复到接近真实的 2D density 形状。

直觉上，Gaussian reference distribution 扮演了一个“标准尺”。如果数据空间某个区域样本很多，直接把它们映射到 embedding 空间会造成那一块 latent density 过高；为了让输出看起来像 Gaussian/hypersphere uniform，encoder 必须把高密度区域在 embedding 中拉开，或把低密度区域压缩/调整。这样，encoder 虽然没有显式输出 $p_X(x)$，但它的 Jacobian 已经记录了“从输入密度到参考密度”的局部体积校正。JEPA-SCORE 就是读取这把局部尺子的数值。

3.2 Lemma 1：Gaussian Embeddings 在高维接近 hypersphere uniform

令 $Z\sim \mathcal{N}(0,I)$，$X=Z/\sqrt{K}$。论文使用的第一个事实是，当 $K$ 增大时，$\Vert Z{\Vert }_2^2/K$ 的均值为 1、方差为 $2/K$，因此 $X$ 的 norm 快速集中到 1；又因为 Gaussian 是 isotropic 的，方向分布均匀，所以 $X$ 在极限上等价于半径 1 的 hypersphere surface 上的 uniform density： $X=\frac{Z}{\sqrt{K}},Z\sim \mathcal{N}(0,I),X\Rightarrow \mathcal{U}(\mathbb{S}(0,1,K)).$ 这一步重要，是因为许多 JEPA/SSL 方法并不一定显式写成“输出 Gaussian”，但它们通过 covariance、whitening、teacher-student entropy 或 non-parametric uniformity 机制，都会让 embedding 接近高维 isotropic Gaussian / hypersphere uniform 的几何状态。论文把这类表示称为 Gaussian Embeddings (GE)。

3.3 Lemma 2：指定 embedding density 等价于学习输入 energy

对任意非线性 encoder $f:{\mathbb{R}}^D\to {\mathbb{R}}^K$，如果输出 density $p_{f(X)}$ 被约束为某个固定形式，那么输入 density $p_X$ 必须由 $f$ 的 level-set 积分和 Jacobian singular values 来补偿。按 level-set 变量 $u$ 重写 coarea 公式，可表达为： $p_{f(X)}(f(x))={\int }_{\{u\in {\mathbb{R}}^D\mid f(u)=f(x)\}}\frac{p_X(u)}{{\prod }_{k=1}^{\mathrm{rank}(J_f(u))}{\sigma }_k(J_f(u))}\mathrm{d}{\mathcal{H}}^r(u),$ 其中 ${\mathcal{H}}^r$ 是 level-set 上的 $r$ 维 Hausdorff measure。若 $f$ 在 support 上近似 bijective，公式退化成更直接的形式： $p_X(x)\propto {\prod }_{k=1}^{\mathrm{rank}(J_f(x))}{\sigma }_k(J_f(x)).$ 论文强调一般 JEPA encoder 未必 bijective，所以 density 只能在 level-set 内达到 mean-preserving rescaling 的可识别性：不同 $u$ 若有相同 embedding，$f$ 可以在 level-set 内重新分配局部 scale，只要积分保持常数。这是理论上的一个重要 caveat：JEPA-SCORE 是 learned density proxy，而不是无条件、无歧义的真实 $p_X$。

3.4 JEPA 目标：prediction + diversity 如何落到 density

论文把多数 JEPA 写成两个目标的组合： $\mathcal{L}\triangleq {\sum }_{n=1}^N{\mathbb{E}}_{(x_n^{(1)},x_n^{(2)})\sim \mathcal{G}(x_n)}\left[\mathrm{dist}\left(\mathrm{Pred}(\mathrm{Enc}(x_n^{(1)})),\mathrm{Enc}(x_n^{(2)})\right)\right]+\mathrm{diversity}\left((\mathrm{Enc}(x_n){)}_{n\in [N]}\right).$ 第一项是 predictive invariance：同一样本经 stochastic operator $\mathcal{G}$ 产生两个 views 后，一个 view 的 embedding 经 predictor 应能预测另一个 view 的 embedding。第二项是 anti-collapse/diversity：所有样本不能映射到同一表示；在很多方法中它等价或近似等价于推动 embedding 分布高熵、白化、uniform 或 Gaussian。

为了把 augmentation 纳入 density，作者设 generator density 为 $p_{\mu }$，transformation density 为 $p_T$，输入数据 density 写成： $p_X\triangleq p_{\mu }\otimes p_T.$ 实际图像训练中，$\mu$ 可理解为原始训练样本，$T$ 是 masking/crop/color 等 view transformation。估计 $p_{\mu }$ 就相当于估计训练数据本身的密度，而不是某个 augmentation 后随机 view 的密度。论文证明在 JEPA loss 最优时： $p_{\mu }(\mu )\propto {\mathbb{E}}_{p_T}{\left[\frac{1}{{\prod }_{k=1}^{\mathrm{rank}(J_f(x))}{\sigma }_k(J_f(\mu ,T))}\right]}^{-1}.$ 单样本、log-scale 的 Monte Carlo estimator 即 JEPA-SCORE： $\mathrm{JEPA}\text{mbox}-\mathrm{SCORE}(x)\triangleq {\sum }_{k=1}^{\mathrm{rank}(J_f(x))}\log \left({\sigma }_k(J_f(x))\right).$ 论文公式与 released code 实现差异：论文式 (5) 是纯数学定义，没有写数值稳定项；论文 Listing 和非官方 jepa-score-eval 都在实现中把 singular values clamp/clip 到 eps=1e-6 后再取 log，避免 $\log 0$。非官方 repo 还提供 randomized VJP approximation，但该近似不是论文方法本身，笔记只把它作为工程 caveat。

3.5 PyTorch 风格伪代码：JEPA training abstraction

import torch
import torch.nn.functional as F
 
 
def jepa_loss(encoder, predictor, batch, make_views, diversity_regularizer):
    x1, x2 = make_views(batch)  # stochastic operator G(x)
    z1 = encoder(x1)
    z2 = encoder(x2).detach()
    pred = predictor(z1)
    invariance = F.mse_loss(pred, z2)
    diversity = diversity_regularizer(encoder(batch))
    return invariance + diversity

这段伪代码不是论文发布的训练脚本，而是式 (3) 的抽象：make_views 对应 $\mathcal{G}$，invariance 对应 view prediction，diversity_regularizer 对应 anti-collapse。论文没有重新训练 I-JEPA/DINOv2/MetaCLIP，因此没有给出这些 backbone 的 launch config、GPU 数或 LR；它关注的是训练完成后如何从 encoder 中读取 density。

3.6 PyTorch 风格伪代码：精确 JEPA-SCORE

import torch
 
EPS = 1e-6
 
 
def jepa_score_full(model, images, eps=EPS):
    # model(images) -> (num_samples, features_dim)
    def summed_features(x):
        return model(x).sum(0)
 
    jac = torch.autograd.functional.jacobian(summed_features, images)
    with torch.inference_mode():
        jac = jac.flatten(2).permute(1, 0, 2)  # (B, feature_dim, input_dim)
        singular_values = torch.linalg.svdvals(jac)
        scores = singular_values.clamp(min=eps).log().sum(dim=1)
    return scores, singular_values

这与论文 Listing 和非官方 repo 的 jepa_score_full 一致：对每个输入样本计算 encoder 输出对输入像素的 full Jacobian，再对每个样本的 Jacobian 做 SVD。复杂度的瓶颈在 full Jacobian 和 SVD，因此论文说实现只有几行，但实际在 ImageNet 尺度上并不便宜。

3.7 PyTorch 风格伪代码：非官方 randomized approximation（不是论文方法）

import torch
 
 
def jepa_score_randomized(model, image, n_proj=64, eps=1e-6, generator=None):
    with torch.no_grad():
        embedding = model(image.unsqueeze(0))
    feature_dim = embedding.shape[1]
    omega = torch.randn(n_proj, feature_dim, device=image.device, generator=generator)
 
    rows = []
    for j in range(n_proj):
        x = image.unsqueeze(0).detach().requires_grad_(True)
        out = model(x)
        (grad_x,) = torch.autograd.grad(out, x, grad_outputs=omega[j].unsqueeze(0))
        rows.append(grad_x.flatten().detach())
 
    projected_jacobian = torch.stack(rows, dim=0)
    singular_values = torch.linalg.svdvals(projected_jacobian)
    score = singular_values.clamp(min=eps).log().sum()
    return score, singular_values

这段只来自非官方 itsanderz/jepa-score-eval，用于降低 exact Jacobian 的成本；它计算的是随机投影后的 Jacobian spectrum，不等价于论文式 (5)。因此如果用它做数据筛选或 OOD detection，应把结果标记为 approximation，而不能声称复现了论文的 exact JEPA-SCORE。

Code reference: main @ 40de6866 (2026-03-20) — pseudocode and mapping based on this non-official commit; official author code was not found in title/arXiv/author GitHub searches.

Paper Concept	Source File	Key Class/Function
论文 Listing 的 exact JEPA-SCORE	jepa_score.py	jepa_score_full(model, x, eps=1e-6)
批量样本打分与 spectrum 保存	jepa_score.py, run_extended.py	jepa_score_batch, run_full_jacobian_batch
非论文方法的 randomized VJP approximation	jepa_score.py, run_extended.py	jepa_score_randomized, randomized projection sweep
OOD evaluation wrapper	run_extended.py	get_ood_pairs, compute_metrics
Baseline comparison in non-official repo	run_extended.py	knn_scores, maha_scores, iforest_scores

4. Experimental Setup (实验设置)

论文实验分为 controlled synthetic validation 和 pretrained backbone qualitative validation 两类。

Datasets used and scale. Controlled 部分使用 Gaussian Mixture Model (GMM) synthetic data。Figure 2 的相关系数表覆盖 input dimension $64,128,256,512$，样本数 $512,1024,2048,4096$。另外，作者在 2D density 上用 Langevin sampling 展示 JEPA-SCORE approximation 是否能恢复真实 data distribution。真实图像部分使用 ImageNet-1k、ImageNet-A/R、MNIST 和 Galaxy；Figure 3 明确对不同 dataset 各取 $5,000$ samples 计算 JEPA-SCORE 分布。Figure 1 在 ImageNet class 21 的 $1,000$ samples 中展示最低/最高概率样本；附加 Figure 5 和 Figure 6 对 ImageNet class 141/147 做同样的 top/bottom visualization，并展示 random samples 作为参照。

Baseline methods compared. 论文主体不是 benchmark paper，因此没有传统表格化 baseline（如 AUROC 对比）。它比较的是不同 pretrained JEPA-like backbone 对同一 scoring rule 的一致性：DINOv2、I-JEPA-1k、I-JEPA-22k、MetaCLIP。引言中把 JEPA 家族分为 moment-matching (VICReg, W-MSE)、non-parametric estimators (SimCLR, MoCo, CLIP) 和 teacher-student (DINO, I-JEPA)，但实验图主要展示 DINOv2/I-JEPA/MetaCLIP。

Evaluation metrics. Controlled GMM 使用 Pearson correlation 衡量 JEPA-SCORE 与真实 $\log p(x)$ 的相关性；2D Langevin 图用视觉方式检查 score approximation 能否引导样本从初始分布移动到真实 density；ImageNet / OOD 图使用 qualitative ranking：低 JEPA-SCORE 代表低 learned probability，高 JEPA-SCORE 代表高 learned probability。Figure 3 使用 score distribution histogram 检查 OOD datasets（MNIST/Galaxy/ImageNet-A/R）相对 ImageNet-1k 是否整体左移。

Training config / hardware / hyperparameters. 论文没有训练新的 JEPA，也没有提供 DINOv2、I-JEPA、MetaCLIP 的 launch scripts、GPU 数、LR、batch size 或 training steps；这些 backbone 被当作 pretrained models 使用。论文唯一明确的 implementation hyperparameter 是 PyTorch Listing 中的 eps=1e-6，用于对 singular values 做 clip/clamp 后取 log。非官方 repo 的 run_extended.py 使用 torch.hub 加载 facebookresearch/dinov2 的 dinov2_vits14、dinov2_vitb14、dinov2_vitl14，并提供 CIFAR10_vs_SVHN / CIFAR10_vs_CIFAR100 / CIFAR10_vs_Textures 的 OOD evaluation，但这不是论文原始实验设置。

5. Experimental Results (实验结果)

5.1 Controlled GMM：JEPA-SCORE 与真实 log-density 高相关

Figure 2 的核心数字如下，列为样本数、行为 input dimension：

dim	512	1024	2048	4096
64	0.64	0.69	0.84	0.90
128	0.75	0.85	0.90	0.94
256	0.82	0.83	0.69	0.76
512	0.72	0.75	0.84	0.88

这些结果说明，在 synthetic GMM 中，JEPA-SCORE 与真实 $\log p(x)$ 通常呈强正相关，尤其在 $128$ 维、$4096$ samples 时达到 $0.94$。它不是严格单调随 dimension 或 sample count 增大而提升：例如 $256$ 维在 $2048$ samples 时为 $0.69$，低于同维度其他设置；这提醒我们 JEPA-SCORE 是依赖训练最优性、Jacobian 数值估计和数据几何的近似，而不是一个无噪声 oracle。

5.2 OOD / domain readiness：未见过的数据整体低分

Figure 3 解读：这张图对 ImageNet-1k/a/r、MNIST、Galaxy 等不同 dataset 各取 $5,000$ samples 并画出 JEPA-SCORE 分布。作者观察到，随着 pretraining dataset size 增大（图中以 I-JEPA-1k 作为参照），MNIST 和 Galaxy 这类训练分布外图像会被赋予更低 probability，即 score distribution 更偏低。这个结果支持一个 practical use case：在 zero-shot 下游任务前，用 JEPA-SCORE 粗略评估目标域是否被当前 pretrained model “看作正常”。

Figure 3 的意义不只是异常检测。它说明 JEPA-SCORE 可以成为 model assessment 的内部信号：同一张图像是否属于模型的 learned pretraining density，并不只由 classifier head 或 downstream loss 决定，而可以直接从 encoder geometry 读出。对数据清洗来说，低分样本可被视为 atypical / low-density candidates；对迁移学习来说，目标数据整体低分意味着当前 backbone 可能缺少该 domain 的支持。

5.3 ImageNet class-level ranking：不同 backbone 的极端样本一致

Figure 4 解读：这张图给出 ImageNet-1k class 21 的随机样本，用作 Figure 1 的参照。随机样本包含多种姿态、背景和拍摄条件；与 Figure 1 中被 JEPA-SCORE 选出的 low/high extremes 对比后，可以看出 score 排序并非简单按类别标签，而是在类内进一步区分“典型/非典型”的 visual modes。

Figure 5 解读：class 141 的低分样本更多是静止、背景复杂、主体小或姿态不典型的鸟；高分侧则更常出现飞行姿态、清晰主体或模型在预训练中更可能频繁见到的 visual pattern。底部 random samples 说明该类别本身存在很大外观多样性，因此 JEPA-SCORE 不是简单把所有同类样本打成一样，而是在同一 label 内估计 learned density。

Figure 6 解读：class 147 重复了同样现象。跨 DINOv2、I-JEPA-1k、I-JEPA-22k、MetaCLIP，低/高 probability 样本的视觉模式有明显一致性；这种跨模型一致性是论文的重要 qualitative evidence，因为这些模型训练机制不同，却都在 encoder geometry 中形成了类似的 density ranking。

5.4 Jacobian singular vectors：score 来自局部敏感方向

Figure 7 解读：这张附加图展示了给定输入图像的 top Jacobian singular vectors，每列对应一个局部高敏感方向，并在标题中标出对应 singular value。虽然 JEPA-SCORE 本身只用 singular values 的 log-sum，不直接使用 singular vectors，但 singular vectors 帮助理解 score 的来源：encoder 对某些局部方向的拉伸/压缩越强，局部体积变化越大，density proxy 也随之改变。

5.5 结论、消融与局限

论文没有给出传统 component ablation 表，也没有报告 AUROC / accuracy 这类下游 benchmark 数字；它的 evidence 主要是理论推导、controlled GMM correlation、Langevin 视觉恢复、以及 ImageNet/OOD qualitative ranking。因此最强结论是“JEPA-style Gaussian embedding objective 在最优时隐式学习 density，并且该 density 可通过 Jacobian spectrum 提取”，而不是“JEPA-SCORE 已在所有 OOD benchmark 上优于专门方法”。

主要局限包括：第一，理论依赖 JEPA loss 接近最优、embedding 确实接近 Gaussian / hypersphere uniform；实际 pretrained model 只能近似满足。第二，full Jacobian 对高分辨率图像和大 encoder 成本高，论文 Listing 简洁但计算并不便宜。第三，非 bijective encoder 存在 level-set ambiguity，density 只能在 level-set 积分意义上恢复。第四，真实图像实验偏 qualitative，缺少系统 benchmark、显著性分析和训练/硬件细节。第五，作者官方代码未找到；公开可检索到的 itsanderz/jepa-score-eval 是非官方独立评估，实现了 exact Jacobian 与 randomized approximation，但后者不是论文方法。

整体上，这篇论文的价值在于把 JEPA 的 anti-collapse 从“防塌缩技巧”提升为“implicit density learning mechanism”：Gaussian Embeddings 不是无害的 latent 形状约束，而是给 encoder 施加了一个必须吸收 $p_X$ 的几何任务。JEPA-SCORE 则把这个理论变成了可操作的诊断工具：对任意已训练的 JEPA-like encoder，只要能计算或近似 Jacobian singular values，就可以得到样本级 learned probability proxy，用于 outlier detection、data curation 和 model-domain readiness analysis。