何时停止：掌握用于提前退出网络的路由选择与风险控制方法

《Knowledge-Based Systems》：When to stop: Learned routing and risk control for early exit networks

　　摘要：早期退出网络通过让简单输入在较浅层就输出，将复杂的计算留到更深层来加速推理过程。现有方法依赖单一的置信度分数来决定何时退出，混淆了路由与校准功能，且无法保证每个样本的错误率。我们解决了这三个问题。我们训练了一个轻量级的预测器，该预测器结合三种信号——预测置信度、特征空间与训练分布的距离，以及不确定性在多次退出过程中的变化情况——来估算每次退出时的错误概率。路由与校准被分开处理：预测器负责决定何时退出，而每次退出后的温度缩放则用于校准输出置信度。随后通过一种符合论方法为每个样本的错误率提供基于数据的正式上限。一个重要发现是熵超鞅性质：对于正确分类的输入，预测不确定性应随深度增加而降低。这一规律在卷积网络中表现明显，但在变换器网络中却不成立，这也解释了为何基于一致性的路由方法在变换器架构上效果不佳。我们的预测器在两种架构中均有效，在相同校准误差下可节省35%的浮点运算次数，而最佳一致性基线仅能节省12.6%。我们还发现并纠正了一种训练偏差，这种偏差会导致深层退出的预测器高估难度，从而使路由策略简化为两种退出方式。分阶段条件训练可解决这一问题。在三个数据集和八种基线方法的对比中，我们的方法有55%–61%的时间能与最优路由决策相匹配，而最佳启发式方法仅为24%–26%；在相同校准误差下，其计算效率是最强一致性基线的2.8倍；当所有启发式方法都失效时，它的性能仍能保持稳定。