华科大与阿里通义联手:让AI图像生成更“聪明”地分配算力

这项由华中科技大学人工智能与自动化学院和阿里巴巴通义实验室联合完成的研究,于2026年6月25日以预印本形式发布,论文编号为arXiv:2606.26938。有兴趣深入了解的读者可以通过该编号在arXiv平台查询完整论文。

一张照片,从一片噪声中慢慢"浮现"出来——这是当前最先进的AI图像生成系统的工作原理。然而,如何让这个"浮现"的过程更加精准、高效,是研究者们持续攻关的核心难题。这篇论文正是针对这个难题,提出了一套名为SharpMoE的解决方案,让AI在生成图像时能够把"算力"更准确地花在刀刃上。

一、AI生成图像背后的"隐形工厂"

要理解这项研究,先从AI是怎么生成图像说起。当你让AI画一只鹦鹉,它并不是凭空直接画出来的,而是从一片随机的噪点出发,一步一步"去噪",就像是从一张布满雪花的电视屏幕上,逐渐还原出清晰的画面。这个过程叫做"扩散模型",是当前最主流的AI图像生成技术。

驱动扩散模型的神经网络规模越来越大,参数量动辄达到数十亿。更大的模型通常意味着更好的生成质量,但同时也带来了巨大的计算开销——每次生成图像,都要把所有参数全部激活,消耗大量算力。于是,研究者们引入了一种叫做"专家混合"(Mixture-of-Experts,简称MoE)的机制。

可以用一家大型设计公司来打比方。这家公司有很多不同专长的设计师(即"专家"),每当接到一个设计任务时,公司并不会让所有设计师都上阵,而是由一个"项目经理"(即"路由器")来判断:这个任务的哪个部分需要哪几位专家来处理。这样既保住了公司的整体实力,又避免了人力浪费。

在AI图像生成中,MoE机制的工作方式与此类似。图像被切分成许多小块("token"),路由器负责判断每个小块应该交给哪几位"专家"处理。一般来说,图像中细节丰富、结构复杂的区域(比如鹦鹉的羽毛纹理)需要更多专家参与,而背景天空这类单调区域则不需要。这种"按需分配"的机制,既扩展了模型能力,又控制住了计算成本。

二、项目经理的眼睛被蒙上了——噪声路由问题

然而,研究团队发现了一个严重的问题:现有的MoE扩散模型中,那个"项目经理"(路由器)其实是在蒙眼睛工作的。

这是因为,扩散模型在生成图像时,越早期的阶段噪声越重,图像越模糊。而路由器做出"哪个小块应该交给哪些专家"的决策时,依据的正是当前这些充满噪声的模糊信息。这就好比项目经理要在完全看不清任务内容的情况下,就决定派谁去干活——结果自然是乱派一通,细节复杂的区域没有得到足够的算力关注,而背景区域却可能被过度处理。

研究团队用一种叫做"拉普拉斯算子"的数学工具来衡量图像每个小块的"显著程度"——直观来说,就是这个小块的纹理有多复杂、细节有多丰富。然后他们统计了现有方法(如DiffMoE)中,不同显著程度的小块实际被分配了多少专家。结果令人失望:无论一个小块的纹理多复杂还是多单调,分配给它的专家数量几乎没有差别。换句话说,项目经理根本没有在根据任务难度来分配人手,所有任务都被"无差别"地对待了。

这个现象被称为"噪声路由"问题——路由器被噪声干扰,失去了准确识别重要区域的能力,导致算力分配严重失准。

三、SharpMoE的核心思路:让项目经理摘下眼罩

针对这个问题,研究团队提出了SharpMoE框架,核心思路其实相当直觉:既然路由器看不清当前步骤里的图像(因为噪声太重),何不让它参考一张"干净的预览图"来做决策?

在扩散模型的每一步去噪过程中,模型都会预测一个当前对最终干净图像的估计,这个估计被称为"x^0预测"。虽然在早期阶段这个估计还很粗糙,但它已经基本勾勒出了图像的主体结构——比如鹦鹉大概在哪里、背景大概是什么。这就像是设计公司在正式开工前,先拿到了一份草图,项目经理虽然看不到最终的精美成品,但已经能从草图上判断出哪个区域的设计工作最繁重。

SharpMoE的做法是:在当前时间步的路由决策中,不仅参考当前的噪声图像,还额外引入上一时间步预测出的干净图像估计,作为判断哪些区域需要更多专家的依据。这样一来,即使当前图像还充斥着噪声,路由器也能借助这份"干净预览",准确识别出图像中真正需要重点处理的区域,从而合理分配算力。

四、双路由器协同:经验与洞察的结合

具体到实现层面,每个SharpMoE模块内部设置了两个路由器,形成"双轨协作"的机制。

第一个路由器是原来预训练好的路由器,它继续负责读取当前时间步的噪声图像,感知模型在这一步去噪过程中的实时状态,捕捉当前阶段的处理需求。第二个路由器是新引入的"显著性感知路由器",它读取上一时间步预测出的干净图像估计,从中获取图像显著区域的信息。两个路由器各自产生一组得分,将两组得分加在一起,作为最终的路由决策依据。

这两个路由器的分工,就像是一位老员工(预训练路由器)凭借丰富经验处理眼前的状况,同时配合一位拥有"透视镜"的新同事(显著性感知路由器)来看清任务的本质。两人协作,决策自然更加准确。

为了避免新引入的路由器一开始就打乱已经训练好的模型,研究团队将显著性感知路由器的初始权重全部设为零。这样,模型在训练初期依然按照原有方式运行,随着训练的推进,新路由器的影响力才逐步增加,平稳地将显著性信息融入路由决策。

五、全轨迹训练:让项目经理看完整个项目进度

引入干净图像估计带来了一个新挑战:标准的扩散模型训练通常只训练"单步去噪",也就是随机抽取某一个时间步,让模型学会在这一步做好去噪,而不关心前后步骤之间的关联。但显著性感知路由器需要上一步的干净预测结果作为输入,单步训练根本无法提供这个信息。

为此,研究团队设计了"递归全轨迹训练"方案。每次训练时,不再只抽取一个时间步,而是随机抽取一组按时间顺序排列的多个时间步(如10步),让模型完整地模拟一段去噪过程。在这段模拟中,每一步都会产生对干净图像的预测,这个预测会被传递给下一步使用。通过这种方式,显著性感知路由器得到了它所需要的信息,训练得以正常进行。

在训练第一步时,由于还没有上一步的干净预测,研究团队用当前的噪声图像作为替代输入。这个处理方式背后的逻辑是:在生成过程的最初阶段,图像的内容尚未成形,任何区域都谈不上显著不显著,噪声图像在这一刻其实是个合理的起点。

此外,研究团队将训练中第一个时间步固定在接近1.0的位置(具体取0.999)。这是因为如果从纯粹的随机噪声出发,训练目标会变得混乱,无法给模型提供有效的学习信号。0.999这个细微的差别,既保证了训练的有效性,又确保了训练和实际使用时的行为保持一致。

六、轨迹路由损失:从全局视角约束算力分配

光有更聪明的路由器还不够,研究团队进一步提出了"轨迹路由损失",用来从训练目标层面直接驱动模型去学会正确的算力分配方式。

思路是这样的:得益于全轨迹训练,研究团队可以统计在整个去噪过程中,每个图像小块累计被分配了多少专家(即累计算力)。与此同时,他们用前面提到的拉普拉斯算子计算出每个小块的真实显著程度,得到一张"显著性地图"。

理想状态下,一个小块的累计算力分配,应该与它的显著程度成正比。越显著的区域,理应获得越多的算力。轨迹路由损失就是用来衡量实际算力分配和理想显著性分布之间的差距,并在训练中不断缩小这个差距。具体使用的是一种叫做KL散度的统计工具,它能够量化两个分布之间有多"像"或多"不像"。

这个损失项与原本的去噪损失(Flow Matching损失)共同构成总训练目标,两者通过一个权重系数(在实验中设为0.001)进行平衡。这个权重确保算力分配的约束足够有效,同时又不会干扰模型生成图像的基本能力。

七、实验结果:效果显著,适用广泛

研究团队在ImageNet数据集(一个包含超过120万张涵盖1000个类别的标准图像库)上进行了系统性评测,采用256×256分辨率的类别条件图像生成任务。

评测使用的核心指标有两个。FID(Fréchet Inception Distance,弗雷歇初始距离)衡量生成图像与真实图像之间的整体差距,数值越低越好。IS(Inception Score)衡量生成图像的多样性和质量,数值越高越好。

SharpMoE被设计为一个"后训练增强"框架——也就是说,研究团队先拿已经训练了50万步的MoE扩散模型作为起点,然后只需再用SharpMoE继续训练10万步,就能获得显著的性能提升。

在以DiffMoE-L(大号模型)为基础的实验中,SharpMoE将FID从3.86降低到了3.10,IS从203.00提升到了228.88(在引导强度cfg=1.5的设置下)。在以TC-DiT-L为基础时,FID从5.07降到了3.72,IS从174.98提升到了206.93。在以EC-DiT-L为基础时,FID从4.09降到了3.27,IS从195.12提升到了221.36。

这些改进贯穿了小(S)、中(B)、大(L)三种模型规模,也贯穿了TC-DiT、EC-DiT、DiffMoE三种不同的MoE架构,充分说明SharpMoE并不依赖某种特定的架构才能发挥作用,而是具有广泛的适用性。

消融实验(也就是逐步拆除各个组件、观察性能变化的对照实验)进一步验证了每个组件的贡献。以DiffMoE-B为基础,在cfg=1.5的设置下,原始模型的FID为8.03。仅加入显著性感知路由机制,FID就降到了6.95。在此基础上再加入轨迹路由损失,FID进一步降到6.66。两个组件各有独立贡献,组合使用效果最佳。

在预训练阶段影响的实验中,研究团队将SharpMoE分别应用于训练了40万步、50万步、70万步的DiffMoE-B检查点,结果发现:无论从哪个阶段开始后训练,SharpMoE都能在10万步内带来一致的性能提升,哪怕模型已经充分收敛也不例外。这说明SharpMoE作为后训练增强工具,具有很强的鲁棒性。

在轨迹步数的影响实验中,研究团队测试了T=5、10、15、20四种不同的全轨迹训练步数设置。结果显示,四种设置均能带来显著提升,性能差异有限,说明SharpMoE对这个超参数并不敏感,不需要精细调优。实验最终选择T=10作为默认设置,因为在该设置下效果略优。

在专家分配的可视化分析中,研究团队统计了不同显著程度的图像小块在生成过程中各个时间步平均被分配的专家数量。对比图清晰地显示:DiffMoE的分配曲线几乎是水平的,各显著程度的小块获得的专家数量相差无几;而SharpMoE的分配曲线呈明显的单调递增趋势,越显著的小块获得的专家越多。在高噪声阶段(生成早期),这种差异尤为明显,正说明SharpMoE的干净预测指导在最需要的时刻发挥了最大的作用。

说到底,SharpMoE解决的是一个听起来很直觉、但实现起来颇费功夫的问题:让AI在生成图像时,能够"看清楚"哪里需要重点投入,而不是蒙着眼睛随机分配算力。研究团队的方案——借用上一步对干净图像的预测来指导当前步的专家分配,再配合全轨迹训练和显著性对齐的损失函数——在多种架构和规模上都表现出了稳定而显著的改进。

对于普通人来说,这项研究的意义在于:未来的AI图像生成系统,在同等计算资源下,能生成细节更丰富、质量更高的图像。鹦鹉的每一根羽毛、食物的每一层纹理,都能得到更充分的算力关注,而不是被淹没在噪声的混乱分配中。

这也引出了一个值得思考的问题:在AI系统越来越复杂的今天,"怎么分配算力"本身就是一门艺术。SharpMoE的探索揭示,让AI更好地"看清楚"自己面对的问题,往往比一味堆砌更多参数更加有效。有兴趣深入了解技术细节的读者,可以通过arXiv编号2606.26938查阅完整论文。

Q&A

Q1:SharpMoE和普通扩散模型的区别是什么?

A:普通扩散模型的路由器在做"哪些区域需要更多算力"的决策时,只能看当前充满噪声的模糊图像,导致判断失准。SharpMoE额外引入了上一步预测的干净图像作为参考,让路由器能够更准确地识别哪些区域真正需要重点处理,从而把算力花在刀刃上,提升生成图像的细节质量。

Q2:SharpMoE训练起来会不会特别费时?

A:不会特别费时。SharpMoE被设计为"后训练增强"框架,可以直接在已经训练好的模型上继续微调。实验表明,只需在原有50万步训练的基础上再追加10万步后训练,就能获得显著的性能提升,而且无论模型处于哪个训练阶段都适用。

Q3:轨迹路由损失具体是怎么约束算力分配的?

A:轨迹路由损失会统计整个去噪过程中每个图像小块累计获得的算力,同时用拉普拉斯算子计算出每个小块的真实复杂程度(显著性),然后用KL散度来衡量两者分布的差距,并在训练中不断缩小这个差距,从而让模型学会把更多算力分配给纹理复杂、细节丰富的区域。

作者声明:作品含AI生成内容
举报/反馈
分享到: 微博 QQ 空间
对本文内容有合作意向?
我们将在 1 个工作日内与您联系
留言咨询