ROME团队 投稿
量子位 | 公众号 QbitAI

在人工智能模型规模持续扩大的今天,数据集蒸馏(Dataset Distillation,DD)方法能够通过使用更少的数据,达到接近完整数据的训练效果,提升模型训练效率,降低训练成本。

但是,通过数据集蒸馏训练的模型,要在安全性要求比较高的任务中(如医疗诊断、自动驾驶),实现不受干扰并保持性能效果,还有一定难度。

来自北京航空航天大学、上海人工智能实验室和英国利物浦大学的研究团队,提出了名为ROME的新方法,这是首次将信息瓶颈理论引入数据集蒸馏任务。该方法无需对抗训练,即可显著提升模型的对抗鲁棒性,最大提升近40%。

资讯配图

实验结果显示,在不同数据集上,相较于以往最优方法,ROME的鲁棒性均实现了大幅超越,最高从此前43.97%暴涨至103.09%

资讯配图

目前,相关成果已被国际机器学习顶会ICML 2025正式接收,项目代码与数据已全面开源。

ROME:首个引入信息瓶颈理论的鲁棒数据集蒸馏方法

数据集蒸馏(Dataset Distillation,DD)是一种高效的数据压缩方案,通过将大规模原始数据压缩为一组可训练的合成样本,使模型在使用更少数据的情况下,达到接近完整数据训练的效果。

尽管多种数据集蒸馏方法在准确率方面表现优秀,如DM、MTT、IDM、BACON等,但在实际应用中,使用蒸馏数据训练的模型仍对对抗攻击高度敏感,难以满足安全关键任务的鲁棒性要求,比如自动驾驶、安防识别、医疗诊断等领域。

针对现有数据集蒸馏方法在对抗攻击面前的脆弱性,研究团队提出了ROME(RObust distilled datasets via inforMation bottlenEck)方法。该方法首次将信息瓶颈(Information Bottleneck, IB)理论引入数据集蒸馏任务。

资讯配图

其核心思想是通过最小化输入数据与其中间层潜在表示之间的冗余信息,同时增强该表示对于最终标签信息的有效性,从而从源头上提升合成数据的对抗鲁棒性

此外,ROME还引入了基于条件熵瓶颈(Conditional Entropy Bottleneck, CEB)的鲁棒先验机制(如下图),在蒸馏过程中加入合理的扰动,以提高生成数据本身的鲁棒性。

资讯配图

为了实现这一目标,ROME构建了两个核心训练目标:

性能对齐项(Performance-aligned Term):确保合成数据能够支持模型获得良好的分类准确率,如下图所示。

资讯配图

鲁棒对齐项(Robustness-aligned Term)(下图):对齐合成图像与其对抗扰动图像的特征分布,使模型对输入扰动更不敏感,如下图所示。

资讯配图

此外,团队还引入了改进鲁棒性比率的精确指标——I-RR,用于更好地评估数据集蒸馏鲁棒性。

I-RR:考虑到现有鲁棒性评估指标在数据集蒸馏任务中的局限性,研究团队基于BEARD 框架提出了改进指标I-RR(Improved Robustness Ratio)。该指标综合了模型准确率与对抗攻击成功率,有效避免了异常情况下鲁棒性评估的误判,能够更全面且直观地反映蒸馏方法在面对攻击时的表现。

结果:无需对抗训练,鲁棒性大幅提升

在CIFAR-10和CIFAR-100数据集上的实验表明,ROME在未采用任何对抗训练的情况下,依然实现了显著的鲁棒性提升。无论白盒还是黑盒攻击,ROME均表现出全面领先的防御能力,充分验证了其强大性能。

资讯配图

上图展示了白盒攻击(包括有目标、无目标攻击)。在CIFAR-10和CIFAR-100数据集上,相较于以往最优方法,ROME在有目标和无目标攻击下的鲁棒性均实现了大幅超越。以CIFAR-100为例,在有目标攻击下,其鲁棒性(I-RR指标)从之前最优的43.97%暴涨至103.09%

资讯配图

上图展示了黑盒攻击。在迁移攻击和查询攻击下,ROME同样表现出色,多个指标提升超过5%,查询攻击下的鲁棒性提升高达15.2%

可视化

下图展示了在不同鲁棒先验配置下,由ROME生成的合成数据集。这些图像突显了不同设置如何影响合成数据的分布,进而揭示了ROME在生成鲁棒蒸馏数据集方面的有效性。

资讯配图

最关键的是,ROME的强大鲁棒性完全不需要对抗训练的加持,其训练时间与标准蒸馏方法几乎无异,远低于引入对抗训练后的时间开销。这证明ROME是一条兼顾了性能、鲁棒性与效率的实用技术路线。

论文地址:
https://openreview.net/pdf?id=agtwOsnLUB
代码链接:
https://github.com/zhouzhengqd/ROME
项目主页:
https://zhouzhengqd.github.io/rome.page/
ICML线上主页:
https://icml.cc/virtual/2025/poster/44781

一键三连「点赞」「转发」「小心心」

欢迎在评论区留下你的想法!

—  —

资讯配图


🌟 点亮星标 🌟

科技前沿进展每日见