https://zhuanlan.zhihu.com/p/2020115822945777267https://github.com/lucas-maes/le-wm核心训练逻辑deflejepa_forward(self,batch,stage,cfg):encode observations, predict next states, compute losses.ctx_lencfg.wm.history_size n_predscfg.wm.num_preds lambdcfg.loss.sigreg.weight# Replace NaN values with 0 (occurs at sequence boundaries)batch[action]torch.nan_to_num(batch[action],0.0)outputself.model.encode(batch)emboutput[emb]# (B, T, D)act_emboutput[act_emb]ctx_embemb[:,:ctx_len]ctx_actact_emb[:,:ctx_len]tgt_embemb[:,n_preds:]# labelpred_embself.model.predict(ctx_emb,ctx_act)# pred# LeWM lossoutput[pred_loss](pred_emb-tgt_emb).pow(2).mean()output[sigreg_loss]self.sigreg(emb.transpose(0,1))output[loss]output[pred_loss]lambd*output[sigreg_loss]losses_dict{f{stage}/{k}:v.detach()fork,vinoutput.items()iflossink}self.log_dict(losses_dict,on_stepTrue,sync_distTrue)returnoutput优雅的 SIGReg 损失SIGRegSketched-Isotropic-Gaussian Regularizer来自同一团队2025年的工作LeJEPA其核心思想是强制 latent embeddings 的分布逼近标准各向同性高斯分布 N(0, I)。这看似只是一个正则化方法的选择但它解决的问题非常根本——如果 embedding 的分布是 N(0, I)意味着每个维度的方差都是1不会有维度死掉→ 防止维度坍缩不同维度之间不相关没有冗余 → 防止信息冗余坍缩一个正则项同时杀死了坍缩的两种形态。而且它有 Cramér-Wold 定理提供的数学保证不是 heuristic。SIGReg 原理图解SIGReg 的工作原理可以用三个步骤来理解。第一步问题——编码器的偷懒想象编码器把每张图片编码成一个小球放到一个高维操场latent space上。预测器的工作是根据当前球的位置和动作猜下一个球在哪。如果编码器把所有球都堆到同一个点预测器的工作就变得极其简单——永远猜还在原地就行。预测误差为零但什么也没学到。这就是坍缩。坍缩状态 所有图片 → 编码器 → 同一个点 → 预测器永远猜对 → 没有意义 正常状态 不同图片 → 编码器 → 不同位置 → 预测器必须真正学习动态规律第二步检查员的巧妙方法——随机投影问题是我们的 latent space 是 192 维的——人类根本没法直接检查球有没有均匀散开。SIGReg 的做法是随机选 1024 个方向把所有球的影子投到每个方向上。就像用手电筒从不同角度照过去看影子的分布。这个方法之所以在数学上成立是因为 Cramér-Wold 定理如果从每一个方向看投影都是标准正态分布那么原始的高维分布就一定是各向同性高斯。代码实现极其紧凑module.py 中的 SIGReg.forward()classSIGReg(torch.nn.Module):Sketch Isotropic Gaussian Regularizer (single-GPU!)def__init__(self,knots17,num_proj1024):super().__init__()self.num_projnum_proj ttorch.linspace(0,3,knots,dtypetorch.float32)dt3/(knots-1)weightstorch.full((knots,),2*dt,dtypetorch.float32)weights[[0,-1]]dt windowtorch.exp(-t.square()/2.0)self.register_buffer(t,t)self.register_buffer(phi,window)self.register_buffer(weights,weights*window)defforward(self,proj): proj: (T, B, D) # sample random projectionsAtorch.randn(proj.size(-1),self.num_proj,deviceproj.device)AA.div_(A.norm(p2,dim0))# compute the epps-pulley statisticx_t(proj A).unsqueeze(-1)*self.t err(x_t.cos().mean(-3)-self.phi).square()x_t.sin().mean(-3).square()statistic(err self.weights)*proj.size(-2)returnstatistic.mean()# average over projections and time其中 x_t.cos().mean() 是经验特征函数的实部self.phi 是标准正态的特征函数一个高斯形状的衰减函数两者差距的加权平方和就是 Epps-Pulley 统计量。第三步打分——影子像钟形曲线吗对每个投影方向SIGReg 用 Epps-Pulley 检验来评估这个一维投影的分布有多接近标准正态。如果接近得分低好如果影子挤成一堆得分高差需要惩罚编码器。好的状态投影后的分布 ≈ 钟形曲线 → SIGReg loss 低 → 编码器没有偷懒 坏的状态投影后的分布 一个尖峰 → SIGReg loss 高 → 编码器在偷懒需要惩罚三步汇总起来SIGReg 的完整流程就是从 1024 个随机方向检查投影是否像标准正态 → 对所有方向取平均 → 得到一个标量 loss。整个过程可微分可以直接放进梯度优化。
