Token导航 LogoToken导航

AI 资讯 · 科技行者

世界模型瘦身:让AI用更少信息做更聪明规划

更新 2026-11-04科技行者正文 6931字阅读约 22分钟

你有没有想过,一个机器人要在脑子里"预演"未来的时候,到底需要记住多少东西?

想象你在玩一个推箱子的游戏,眼前的画面里有成千上万个像素点:箱子的颜色、地板的纹理、墙上的裂缝、光线的角度。可你真正需要关心的,可能只是箱子在哪、墙在哪、你要往哪推。如果你的大脑把所有像素信息都原封不动地记下来去做决策,那不是聪明,那是累赘。

这正是训练AI世界模型时遇到的一个怪现象。研究人员发现,给世界模型配一个更大的"内存"(也就是更宽的潜在表示向量),性能有时候反而会下降。这事挺反直觉的:我们一般以为模型越大越聪明,可现实中它却可能因为"记太多"而变得更难预测未来、更容易在连续推演中累积错误。

这篇来自Arm AI Research团队的论文,就是想解决这个"记忆臃肿症"。他们提出了一套叫做ALeWM(Adaptive LeWorldModel)的方法,核心思路听起来挺朴素:让模型自己学会把最重要的信息,排到表示向量的最前面几个位置去。这样一来,无论你用多宽的向量去训练,模型在做实际规划的时候,可能只需要看前面一小段"精华摘要",就能把事情做好。

问题到底卡在哪里

先说说世界模型是干什么的。

世界模型*:一种能够预测"如果我这样做,接下来会发生什么"的AI系统,常用于机器人规划、游戏策略和自动驾驶等场景。

它的工作原理大致是:给AI看一张当前画面,再给它一个动作(比如"把机械臂往左推5厘米"),它要预测出下一张画面(或者下一张画面的某种压缩表示)大概是什么样子。这个能力让AI可以在脑子里"模拟"很多种可能的未来,然后挑一个最好的方案去执行,这就是所谓的"规划"。

这篇论文关注的这类世界模型,用的是一种叫联合嵌入预测架构的技术路线。

联合嵌入预测架构(JEPA)*:不直接重建原始图像像素,而是在一个压缩后的"嵌入空间"里预测未来状态的架构,好处是不用纠结图像里的无关细节,比如光线噪点、纹理变化。

这个技术路线里有个代表作叫LeWM(LeWorldModel),它用一个紧凑的编码器把图片压成一个向量,再用一个预测器根据当前向量和动作去猜下一步的向量。为了防止模型"偷懒"把所有向量都压成同一个值(这叫"坍缩"),LeWM还加了一个叫SIGReg的正则化手段,逼着嵌入向量的分布看起来像一个标准的高斯分布。

高斯分布*:也就是钟形曲线分布,统计学里最常见的一种概率分布,这里用来保证向量的每个维度都有足够的"变化空间",不会退化成一个固定值。

这套方案跑起来效果不错,可问题恰恰出在这个"宽嵌入"上。研究人员观察到,当嵌入的维度变大,规划效果反而经常变差。他们的猜测是:预测器需要预测的目标变大了,预测就变得更难,而规划的时候模型要反复用这个预测器往前推演很多步,每一步的误差都会往后面传,越传越大,到最后规划出来的动作可能就跑偏了。

想缩小维度来解决这个问题吗?确实可行,但你得每换一个任务就重新调一次维度大小,调不好还会限制表示能力。这就好比你为了防止背包太沉决定扔东西,但扔多少、扔什么,全靠每次出门前反复试错,这显然不是长久之计。

于是研究者想了另一个办法:既然SIGReg这类常规的"防坍缩"手段只会让所有维度都变得"平等重要"(不分主次),那能不能专门设计一种机制,让模型自己学会把关键信息排到前面,次要信息放到后面,同时又保留足够宽的向量用于学习?

这就是ALeWM要解决的核心矛盾:既要宽,又要能瘦身使用。

前缀机制:让模型学会"排座次"

ALeWM的第一个核心设计,叫做容量网络(capacity network),它的任务是给每一段观察序列,预测出一个"该用多宽的前缀"的概率分布。

这里要理解一个关键概念。

前缀(prefix)*:指从嵌入向量的第一个维度开始,取连续的前k个维度,后面的维度全部清零。比如一个192维的向量,前缀长度选32,就是只保留前32维,剩下160维全部归零。

训练过程中,模型会随机抽样一个前缀长度,把嵌入向量按这个长度"砍掉"后面部分,然后让预测器只用这个被砍过的前缀,去预测完整的、没被砍过的下一步嵌入向量。

这个设计的巧妙之处在于:"完整预测"这个目标逼着模型必须把最有用的信息尽量塞进前面几个维度里,因为前面的维度被保留的概率更高,它们"出场"的次数自然就比后面的维度多。

打个比方。假设你是一个每天要打包行李箱出差的人,但有一个古怪的规矩:每次出门前,有人会随机决定你这次只能带行李箱的前半部分、前三分之一,或者全部。你完全不知道这次会被允许带多少。如果你是个聪明人,你会怎么打包?

你一定会把牙刷、钱包、身份证这些"不管怎样都得用上"的东西,塞到箱子最先被翻到的那一层。而把那件可能用不上的备用外套,放到最底下。因为如果你把牙刷塞在箱子最底下,结果这次只被允许带"前三分之一",那你就惨了,牙刷根本拿不出来。

这正是ALeWM里前缀采样机制在干的事。它用随机砍断前缀这种"惩罚机制",逼着模型把最关键的预测性信息往前排。如果没有这个机制,模型完全可以把重要信息随便塞在向量的任何角落,训练出来的表示会是一团乱麻,没有任何结构可言。

那前缀长度到底怎么确定?这就要靠一个叫做容量集合的东西来规定可选范围。

容量集合(K)*:预先设定的一组可选前缀长度,比如{8,16,32,64,96,128,160,192},训练和部署时模型只能从这个集合里挑一个长度,不是任意数字都行。

容量网络会给容量集合里的每个候选长度算一个概率,概率最高的那个就是模型认为"最适合当前这个序列"的前缀长度。到了部署阶段,模型看一眼初始观察和目标观察,就直接选定一个前缀长度,整个任务过程中都用这个长度去做规划,不再变化。

MixSIGReg:给混乱的"坍缩防线"加上秩序

光靠前缀采样这一个机制还不够。研究者发现,如果直接套用原来的SIGReg正则化去约束这些被砍断的嵌入,会出一个大问题:SIGReg根本不管前缀有多长,它只会一视同仁地要求所有维度都保持高斯分布式的方差,这跟"后面维度应该经常是零"的要求是互相矛盾的。

于是他们设计了一个新的正则化方法,叫做MixSIGReg。

MixSIGReg*:一种改进版的正则化手段,它不再要求嵌入向量整体匹配一个标准高斯分布,而是匹配一个"混合分布":按照容量集合里每种前缀长度出现的先验概率,把对应长度的高斯前缀和全零后缀混合起来,作为新的匹配目标。

说白了,MixSIGReg给模型的"标准答案"变成了这样:假设前缀长度是32,那前32维应该像标准高斯分布那样有方差变化,后面160维应该统统是零。再把所有可能的前缀长度按概率权重叠加起来,就得到了一个混合参考分布。模型训练时,要让自己产生的嵌入向量,尽量在统计上逼近这个混合分布。

这里有个数学上挺漂亮的结论:在这个混合参考分布里,某个维度的方差,正好等于"这个维度在随机抽样时被保留下来"的概率。研究者证明了,越靠前的维度,被保留的概率越高,所以方差也必然越大,而且这个方差是阶梯式递减的,只在容量集合规定的边界处才会下降。

这就像是公司里的岗位编制表。假设公司规定,只有排名前10%的员工才能保留"核心项目组"的席位,前30%能保留"重要项目组",剩下70%只能进"边缘项目组"。核心项目组因为门槛最高、人数最少,每个留下来的人获得的"关注度资源"(对应这里的方差)自然也最集中、最高。如果公司不设这个层级制度,而是让所有人平均分配关注度,那"谁更关键"这件事就根本体现不出来,大家一样重要,也就没有重要的人了。

MixSIGReg做的正是这件事:给前面的维度分配更高的"方差预算",逼着它们承载更多可区分的信息,而不是像SIGReg那样把所有维度一视同仁地"摊大饼"。

论文还从理论角度论证了,只要有两个信息块,其中一个带来的预测误差下降幅度更大,那么把这个大收益的信息块排在更靠前的位置(也就是更常被保留的位置),整体的期望预测误差会更低。这个道理其实挺直观:你把最有用的东西放在最容易被拿到的地方,总比放在角落里划算。

不过论文也很诚实地说明了,这套理论只能说明"这样排序是好的",并不能保证联合训练真的会自动排出这个顺序,这是一个合理的边界声明,不是过度承诺。

规划阶段怎么用

训练好了模型,到了实际干活的时候,ALeWM是怎么规划动作的?

流程是这样的:给模型看一下任务开始时的画面和目标画面,容量网络会给出一个它认为最合适的前缀长度,这个长度在整个任务执行期间保持不变。接下来模型做的所有"内心模拟"(也就是递归预测未来几步的嵌入向量),都只用这个固定长度的前缀去算,包括和目标状态做比较的那一步,也只比较前缀部分,后面的维度完全不参与评分。

这跟原来的LeWM有个本质区别:LeWM训练出来是啥维度,规划时就得老老实实地在全部维度上做计算,没有选择权。而ALeWM因为训练时就刻意制造了"不同前缀长度都要能用"的压力,所以它天生就具备了"按需缩小"的能力,一个模型可以适应多种规划场景,不用为每种场景单独训练一个新模型。

四个水下振荡器告诉你的事

光说理论有点空,研究者专门设计了一个简单的玩具实验来验证这个思路。

他们造了一个"受控阻尼振荡器"系统,说白了就是两个弹簧振子,每个振子有位置和速度两个状态变量,加起来正好是四个真实的物理状态变量。但观察到的数据是十维的,里面只有四个维度跟这四个真实变量关系密切,剩下六个维度是混杂了各种噪声的间接测量值。

这个设计很巧妙,因为它提前知道"标准答案"是四维的,可以用来检验模型学到的表示是不是也把关键信息集中到了对应数量的维度里。

实验结果挺漂亮。一个固定八维的LeWM,当它的维度刚好等于真实状态维度(四维)时,线性解码能力能达到0.980的R?,但维度一旦增加到六维、八维,反而跌到0.930和0.904。这说明维度不匹配,的确会损害表示质量,这和前面说的"更宽不一定更好"是吻合的。

而训练出来的八维ALeWM,它的前四个维度拿出来做线性解码,R?能达到0.983,跟那个刚好四维的LeWM打平,甚至还略微超过。更关键的是,它后面那些额外的维度没有变成"无用的负担",而是保留了额外信息供其他用途,这恰恰验证了ALeWM想做的事:一个模型既能保持宽向量的表达能力,又能把预测关键信息收拢到紧凑的前缀里。

论文还检验了嵌入向量的协方差谱,ALeWM展示出四个明显占主导的方向,跟四个真实因子基本吻合;而固定宽度的LeWM始终保持近似"各向同性"的状态,也就是说它没能学出任何结构性的主次之分。

四个真实测试场景的表现

玩具实验说明了原理,真正的考验是在实际的视觉控制任务上。研究者在四个任务上做了测试:TwoRoom(两个房间的导航)、PushT(推T形物块)、Reacher(机械臂到达目标)、OGBench-Cube(操作立方体)。

用的骨干网络是ViT家族的视觉Transformer,最大嵌入维度分别设置成96、192、384三档,分别对应小、中、大三种模型规模。

结果呢,在所有测试配置里,ALeWM的平均成功率都是最高的。相比用满最大维度的LeWM,ALeWM能提升大约1.5到17个百分点的成功率,同时把平均规划所用的维度数砍掉67%到98%。就算把LeWM的维度当成超参数精心调过,ALeWM依然领先0.3到7.2个百分点,而且在12组对比里有11组用的规划维度更少。

最亮眼的是OGBench-Cube这个任务。在192维骨干下,调过参的LeWM最高成功率是72.33%,而ALeWM达到79.00%,规划用的维度平均只要16,相比LeWM用满的96维,直接少了80维。这就好像两个人去考试,一个人带着满满一箱复习资料,结果因为翻找耗时反而答题慢,另一个人只带了几页精炼的重点笔记,却答得又快又准。信息不是越多越好,关键是你能不能快速找到真正有用的那一部分。

有意思的是,在Reacher这个任务上,研究者发现容量网络选的前缀长度会因为起点和目标的距离不同而变化,目标离得越远,模型越倾向于选更大的前缀。这说明容量选择不是瞎蒙的,它跟任务的实际难度产生了某种呼应关系,虽然论文也坦诚承认,这种变化在不同训练随机种子之间也会有波动,说明这套机制还没有做到完全稳定和可预测。

消融实验:拆开看看每个零件到底有没有用

研究者没有满足于"整体效果好",还做了几组对照实验,想弄清楚到底是哪个部件在起作用。

第一组对照是"嵌套丢弃"(Nested Dropout),也就是训练时用一个固定不变的容量概率分布,而不是让模型自己学。结果显示,无论是配SIGReg还是配MixSIGReg,嵌套丢弃训练出来的模型,在所有测试过的前缀长度里都挑不出一个统一的最优解,表现忽高忽低。相比之下,ALeWM在PushT上用53.49维平均规划容量就达到96.00%的成功率,而嵌套丢弃最好的配置得用96维才能到92.50%。说明让模型自己学习容量分配,确实比死板地套用一个固定先验更聪明。

第二组对照是检验"容量选择机制"本身有没有用。研究者把训练好的ALeWM拿出来,强行让它用八种固定容量分别跑一遍,不让容量网络自己选。结果发现没有哪个固定容量是全能冠军,而ALeWM的自动选择能在不同任务上都匹配甚至超过各自的最佳固定容量,还普遍用更小的平均前缀长度。这说明容量网络确实学到了某种"因地制宜"的本事,不是随便凑出来的。

第三组对照是直接比较MixSIGReg和原始SIGReg,在完全相同的架构下换掉正则化方法。结果MixSIGReg在PushT上提升了2.50个百分点的成功率,在OGBench-Cube上提升了8.17个百分点,同时把平均选中的前缀长度从138.67砍到53.49,从128.00砍到16.00。这组数据算是给MixSIGReg的设计价值做了最直接的背书。

这个思路放在更大的图景里

说回开头的问题:世界模型到底该记多少东西?

这篇论文给出的答案是:不该事先固定好"记多少",而应该让模型自己学会"按需分配重要性",同时保留足够的容量供学习使用。这其实跟很多领域的智慧是相通的,比如图像压缩里早就有"按重要性排序像素块"的技术(比如渐进式JPEG),推荐系统里也有类似的多粒度表示学习(Matryoshka表示学习),它们都在解决同一个问题:如何让一份数据,同时满足"粗略看一眼就够用"和"仔细看能看到更多细节"这两种不同精度需求。

ALeWM把这个思路带进了世界模型和规划领域,而且做了理论上的论证,证明这种排序确实能降低预测误差的期望值。这给后续研究留下了不少值得追问的方向:比如容量网络的训练稳定性问题还没完全解决,不同随机种子跑出来的选择会有波动;又比如目前用的是直通估计器(straight-through Gumbel-Softmax),这种梯度估计方法本身是有偏的,换成无偏的估计方式会不会让容量学习更稳?这些都是论文自己提出来、留给未来研究的开放问题。

写在后面

读这篇论文的时候,最让我意外的不是"排序"这个思路本身,毕竟压缩编码领域早就有类似的设计,而是那个"维度越宽,规划反而越差"的现象。这其实是对"规模越大越好"这个直觉的一次小小打击。

更值得琢磨的是论文里那个关于Reacher任务的小观察:当起点和目标离得越远,模型会倾向于选更大的前缀容量。这个细节没有被论文大写特写,但它暗示了一种可能性,就是信息压缩的程度,或许应该跟任务难度动态绑定,而不是一个固定不变的超参数。这跟人类做事的方式挺像,简单任务凭直觉就能应付,复杂任务才需要调动更多的注意力资源。

论文自己也承认,目前的容量选择机制在不同训练种子下会有波动,这不是一个已经打磨完美的成熟方案,更像是一个证明了方向可行、但细节还在打磨中的早期探索。这种诚实的态度,比把结果吹得天衣无缝更让人信任。

如果把这个思路继续往前推一步,一个自然的问题是:能不能让模型在执行任务的过程中,动态调整它用的容量,而不是像现在这样一开始选定就不变?如果机器人在任务执行中遇到了意外情况,需要更精细的信息来应对,它能不能临时"扩容"?这篇论文目前没有回答这个问题,但留下的这个口子,足够让人多想一会儿。

Q&A

Q1:ALeWM和LeWM最大的区别是什么?

A:LeWM训练出来的嵌入维度是固定的,规划时必须用满全部维度。ALeWM则会学一个容量网络,自动判断每个任务该用多少维度的前缀来规划,既能保留宽嵌入的表达能力,又能在实际规划时只用一小段紧凑的信息,平均能省下67%到98%的规划容量。

Q2:MixSIGReg解决了什么问题?

A:原来的SIGReg正则化会让嵌入向量所有维度都变得同等重要,没法体现主次关系。MixSIGReg改成匹配一个混合分布,让越靠前的维度分到的方差预算越高,从而逼着模型把关键预测信息集中在前面几个维度,而不是随意分散在整个向量里。

Q3:为什么更宽的嵌入向量反而会让规划效果变差?

A:研究者发现,嵌入维度变大后,预测器要预测的目标也变大,预测难度上升。规划时模型要反复递归调用预测器往前推演多步,误差会一步步累积放大,最终导致规划出来的动作偏离目标,所以维度越宽不代表规划效果越好。

文章标签学术研究

资讯来源:由AI资讯编辑整理自互联网公开内容,版权归原作者所有,未经许可,不得转载。