动态调整退出阈值,让模型更可靠地提前输出结果。
Beyond Greedy Exits: Improved Early Exit Decisions for Risk Control and Reliability
- 用多臂老虎机在线自适应调整退出阈值,无需标注数据。
- 在多个任务上实现1.70~2.10倍加速,性能损失小于2%。
- 适合需要低延迟与高可靠性的实际部署场景。
早期退出深度神经网络通过在中间层进行预测实现自适应推理,显著降低计算开销和延迟。现有方法通常在置信度超过预设静态阈值时立即退出,但该阈值基于验证集设定,可能导致模型对错误类别过度自信,且在部署时面对分布偏移时缺乏鲁棒性,损害模型可信度与准确率。为此,本文提出UAT框架,采用多臂老虎机机制在线、无监督地动态调整退出阈值。UAT基于新设计的奖励函数,综合评估预测确定性及其可靠性,在计算效率与预测质量间取得平衡,并惩罚不必要的晚期退出。我们提供了UAT所实现风险的理论保证,并在视觉-语言理解、文本生成和分类等多样化任务上验证其性能。实验表明,相比完整模型,本框架在保持<2%性能下降的前提下,实现了1.70~2.10倍的速度提升。源代码已开源于https://github.com/Div290/UAT。
原文摘要 · Abstract (English)
Early-Exit Deep Neural Networks enable adaptive inference by allowing prediction at intermediary layers, significantly reducing computational costs and latency. Most of the early exit strategies greedily exit a sample at an intermediary layer if the confidence in class prediction exceeds a predefined threshold that is set using a static validation set. This is problematic as the model might be overconfident in a wrong class. Also, they are not robust to distribution shifts encountered in deployment, which can undermine model trustworthiness and accuracy. To address these challenges, we propose UAT that adapts the threshold for exit decisions using a Multi-Armed Bandit framework, enabling online, unsupervised adjustment of exit decisions. UAT makes decisions based on a new reward function that assesses predictive certainty and its reliability to balance computational efficiency and prediction quality while penalizing unnecessary late exits. We provide guarantees on risk achieved by UAT and validate its performance on diverse tasks spanning vision-language understanding, text generation, and classification. Our framework demonstrates consistent improvements in speedup (1.70-2.10x) with a minimal performance drop (<2%) as compared to full model performance. Our source code is available at https://github.com/Div290/UAT.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。