Expanding SPHERE-JEPA如何实现

SPHERE-JEPA就不看了，应该是和SIGreg的方法差不太多的，用蒙特卡洛投影到一维去处理的。Expanding SPHERE-JEPA继承了SPHERE-JEPA的理念，就是让特征在单位超球面上均匀分布（SIGreg径向上还存在概率分布）。

它首先把编码得到的向量归一化，让所有向量在单位超球面 $\mathbb S^{d-1}$ 上了。

概念：核函数、谱展开

核函数用于测量两个样本在高维（通常是无限维）特征空间中相似度，而无需显式计算该空间中的坐标。这里是定义在球面上的核函数，也就是这个核函数的输入全在球面上：

$$k(x,y),x\in\mathbb S^{d-1},\ y\in\mathbb S^{d-1}$$

应用紧致流形谱定理

对球面配备Laplace-Beltrami算子，高维球谐$\{Y_{l,\alpha}\}$是$\mathbb S^{d-1}$上完备正交本征基底。任意平方可积二元球面函数都可以做双重谱展开：

$$k({x,y})=\sum_{l,\alpha}\sum_{l',\beta} C_{l\alpha l'\beta}\,Y_{l,\alpha}(x)\,Y_{l',\beta}(y)$$

此时没有化简，两套球谐分别对应 $x$ 和 $y$ ，项数极多。 只要输入在球面，一定能这么拆；这由紧致流形谱定理保证，和旋转不变无关。

旋转不变性

我们需要这个核函数具备旋转不变性，也就是对任意正交旋转矩阵$R$（旋转操作），满足 $k(Rx,Ry)=k(x,y)$ 。几何含义：整体转动坐标系，相似度完全不变。

线性代数可证明：满足该条件的函数，必然只依赖向量内积：

$$k({x,y})=\varphi(x^\top y)$$

原本二元函数坍缩为一元函数 $\varphi(c),c=x^\top y\in[-1,1]$。仔细想这个 $c$ 就代表超球面上两个向量之间夹角的余弦。

Schoenberg定理

$\mathbb S^{d-1}$上正定、旋转不变的核，一定可以按Gegenbauer多项式$\tilde C_l^{(\alpha)}$正交级数展开：

$$\begin{equation} k({x,y})=\varphi(c)=\sum_{l=0}^{\infty} w_l(\lambda_l)\tilde C_l^{\alpha}(c) \tag{1} \end{equation}$$

其中$\alpha=(d-2)/2$，$\lambda_l=l(l+d-2)$是拉普拉斯本征值。

超球面上的谱核设计

（RKHS平滑度）对（1）归一化：

$$\begin{equation} k({x,y})=\varphi(c)=\frac{1}{Z}\sum_{l=0}^{\infty} w_l(\lambda_l)\tilde C_l^{\alpha}(c) \tag{2} \end{equation}$$

归一化常数$Z=\sum_{l=0}^{\infty} w_l(\lambda_l)\tilde C_l^{\alpha}(1)$，强制 $\varphi(1)=1$。

每个本征值 $\lambda_l$ 等于其相应特征函数的狄利克雷能量。由于这些核严重惩罚高能量（不光滑）函数，因此从这一通用展开中导出的统计检验在统计学文献中广泛被称为“Sobolev 检验”。

谱权重

• 带限核：$w_l(\lambda_l)=1_{l \le L}$；
• 热核：$w_l(\lambda_l)=e^{-\lambda_l t},t>0$；

谱权重 $w_l(\lambda_l)$ 充当频率滤波器。较低的 $l$ 值对应平滑的全局几何变化，而较高的值则捕捉日益振荡的空间模式。

统计检验的显式形式

这部分的理论会用尽可能好理解的方式写推导，完整过程看论文附录。

MMD

MMD（Maximum Mean Discrepancy），即最大均值差异，对比两组分布的核均值嵌入距离，衡量整体分布差距。

$$D_{\text{MMD}} = \frac{1}{C_{\text{norm/MMD}}} ( \mathbb{E}_{x,y\sim \hat{p}}[\varphi(c)] - C_{\text{bias/MMD}}), \tag{3}$$

$\mathbb{E}_{x,y\sim \hat{p}}[\varphi(c)]$ 是目前所有的超球面上的点俩俩求 $\varphi(c)$ 得到的均值，在所有点都在一起时（$c=1$）取到最大值1。

$C_{\text{bias/MMD}}$ 含义是球面连续（无穷多个点）均匀分布时俩俩求 $\varphi(c)$ 得到的均值。$C_{\text{bias/MMD}} = \int_{-1}^1 \varphi(v)\rho_d(v)dv$ 通过利用超球面边缘密度 $\rho_d(v) \propto (1-v^2)^{\frac{d-3}{2}}$ 的高效高斯-雅可比求积法进行数值计算。

$C_{\text{norm/MMD}} = 1 -C_{\text{bias/MMD}}$ 归一化使得损失 $D_{MMD}$ 最大为1，优化目标是使它尽可能小，趋近于0。

KL

KL（Kullback–Leibler Divergence），即KL-散度/相对熵，衡量两个概率分布的信息差，本文用核密度估计近似球面分布的均匀程度。

$$D_{\text{KL}} = \frac{1}{C_{\text{norm/KL}}} ( \mathbb{E}_{x\sim \hat{p}} [\log \mathbb{E}_{y\sim \hat{p}_{-x}} [\varphi(c)]] - C_{\text{bias/KL}}) \tag{4}$$

$\mathbb{E}_{x\sim \hat{p}} [\log \mathbb{E}_{y\sim \hat{p}_{-x}} [\varphi(c)]]$ 的含义就是我选一个点，然后拿其他所有点和这个点求 $\varphi(c)$ （留一法）的均值，所有点均值的 $log$ 再求均值。当所有点都在一起时（$c=1$）取到最大值0。它用来描述点的密度。

因为这里KL散度的定义是球面均匀KL散度 $KL(p \parallel q) = \int p(x) \log \frac{p(x)}{q(x)} dx$ ，所以我们还需要一个真实密度 $q(x)$ 作为基准。

在球面连续均匀分布时，这个密度 $q(x)=\frac{1}{|\mathbb{S}^{d-1}|}$。那么真实的KL散度表达式为：

$$KL = \mathbb{E}\left[\log \hat{p}(x)\right] + \log \left|\mathbb{S}^{d-1}\right|$$

论文中的（3）的优化方向和-KL散度一致，因为它令

$$C_{\text{bias/KL}}=\log|\mathbb{S}^{d-1}|$$

分子的取值范围即 $[- \infin, -C_{\text{bias/KL}}]$。( $|\mathbb{S}^{d-1}|$ 必然大于1)

$C_{\text{norm/KL}} = -C_{\text{bias/KL}}$ 归一化使得$D_{KL}$ 最小为1，优化目标是使它尽可能大。

KSD

KSD（Kernel Stein Discrepancy），即核斯坦因差异/核斯坦散度，利用斯坦算子检验分布是否匹配目标均匀流形。

KSD的理解更困难一点，一步一步来。

在任何维度下描述旋转都是确定旋转平面。因为圆周运动天然被约束在二维平面内，任何旋转本质都是某二维平面内的圆周转动。$d$ 维空间的旋转自由度就是 $C_d^2$ 。

我们希望球面上所有点均匀分布，等价于 $\varphi(x^{\top}y)$ 沿任意旋转切向挪动一点点 $x$ 或 $y$ ，$\varphi$ 不变。用数学公式刻画就是：

$$k_q = \sum_{i<j} \mathcal{A}_{ij}^x \mathcal{A}_{ij}^y \varphi,\tag{5}$$

其中$i,j$代表选取的两个维度，能确定旋转平面。$k_q$ 的含义就是我先沿 $(i,k)$ 旋转切向对 $y$ 求一阶偏导，再沿同一个切向对 $x$ 再求一个偏导。然后把全部方向上的这个“联合二阶偏导”求和。

某一个旋转平面下 $x$ 处切向量可以利用反对称基础矩阵写出 $E_{ij}x$。

那么切向导数算子（Stein算子，球面黎曼斯坦微分算子） $\mathcal A_{ij}^x= (E_{ij}x)^\top\nabla_x$ ，代入（5）有：

$$\mathcal{A}_{ij}^x \mathcal{A}_{ij}^y \varphi = \varphi''(c)\cdot [(E_{ij}x)^\top y][x^\top E_{ij}y] + \varphi'(c)\cdot (E_{ij}x)^\top (E_{ij}y)$$ $$[(E_{ij}x)^\top y][x^\top E_{ij}y]=y^\top (E_{ij}x)(E_{ij}y)^\top x$$ $$(E_{ij}x)^\top E_{ij}y=\mathrm{tr}\big((E_{ij}x)(E_{ij}y)^\top\big)$$

令 $M=\sum_{{i<j}} (E_{ij}x)(E_{ij}y)^\top=\frac12\big(cI-yx^\top\big)$，（5）可以写成：

$$k_q=\varphi''(c)\cdot y^\top M x +\varphi'(c)\cdot \mathrm{tr}(M)$$

可以对 $M$ 简化得到：

$$M=\frac12\big(cI-yx^\top\big),\quad c=x^\top y$$

1. 计算二次型 $y^\top M x$

$$y^\top M x=\frac12 y^\top(cI-yx^\top)x =\frac12\left(c y^\top x - y^\top y \,x^\top x\right)$$

$x,y$为单位向量：$x^\top x=y^\top y=1,y^\top x=c$

$$y^\top M x=\frac12(c^2-1)$$

2. 计算迹 $\mathrm{tr}(M)$

$$\mathrm{tr}(M)=\frac12\mathrm{tr}(cI-yx^\top) =\frac12(cd-\mathrm{tr}(yx^\top)) =\frac{c(d-1)}{2}$$

回代合并得到最终公式

$$k_q=\frac12\left[(c^2-1)\varphi''(c)+c(d-1)\varphi'(c)\right]$$

再来看论文公式就很好理解了。

$$D_{\text{KSD}} = \frac{1}{C_{\text{norm/KSD}}}\mathbb{E}_{x,y\sim \hat{p}}[ \frac12 \big((c^2 - 1)\varphi''(c) + c(d - 1)\varphi'(c)\big)], \tag{6}$$

很显然均匀状况下二阶偏导都为0，最小值是0；当所有点都在一起时（$c=1$）取到最大值 $\frac{d-1}{2}\varphi'(1)$。

自然归一化常数 $C_{\text{norm/KSD}} = \frac{d-1}{2}\varphi'(1)$ 。$D_{KSD}$ 优化目标是使它尽可能小，趋近于0。

概念补充

• 闭式：能用有限次基础初等运算写出的显式表达式。
• 海森矩阵：二阶偏导矩阵
  • $\nabla_x \nabla_y^\top k(x,y) = \varphi''(c) y x^\top + \varphi'(c) I$。
  • KSD的推导中我们用斯坦算子通过给海森矩阵左右乘旋转切向量，过滤掉所有无关的非旋转分量。

最后一步

以上三种 $D$ 放入JEPA中就是：

$$\mathcal{L}_{\mathrm{reg}} = \frac{1}{V_a}\sum_{v=1}^{V_a} D(\hat{p}_v, q)$$

这就没什么好说了，优化方向和选择的 $D$ 完全一致。$V_a$是数据增广后的全部示图总数量。