[论文解读] Uncertainty Estimation and Calibration with Finite-State Probabilistic RNNs

Cheng Wang, Carolin Lawrence+1 more|arXiv (Cornell University)|Nov 24, 2020

Fault Detection and Control Systems参考文献 55被引用 7

一句话总结

本文提出 ST-τ，一种有限状态概率 RNN，通过使用可学习温度参数的 Gumbel-Softmax 技巧实现随机状态转移，以估计不确定性。通过采样多个推理路径，ST-τ 能够量化预测不确定性，从而实现校准良好的概率输出，提升分布外检测性能，并在强化学习中实现更优的探索-利用权衡，且无需引入贝叶斯神经网络的复杂性。

ABSTRACT

Uncertainty quantification is crucial for building reliable and trustable machine learning systems. We propose to estimate uncertainty in recurrent neural networks (RNNs) via stochastic discrete state transitions over recurrent timesteps. The uncertainty of the model can be quantified by running a prediction several times, each time sampling from the recurrent state transition distribution, leading to potentially different results if the model is uncertain. Alongside uncertainty quantification, our proposed method offers several advantages in different settings. The proposed method can (1) learn deterministic and probabilistic automata from data, (2) learn well-calibrated models on real-world classification tasks, (3) improve the performance of out-of-distribution detection, and (4) control the exploration-exploitation trade-off in reinforcement learning.

研究动机与目标

• 为解决标准 RNN 在分布外输入下概率校准性差的问题。
• 提出一种无需贝叶斯神经网络复杂度的 RNN 不确定性量化方法。
• 在现实世界分类和强化学习任务中，实现校准良好的预测结果，并提升不确定性感知决策能力。
• 支持从数据中学习确定性和概率性自动机。
• 通过内在不确定性建模，控制部分可观察环境中的探索-利用权衡。

提出的方法

• 该模型使用一组 k 个可学习状态，并在每个 RNN 时间步计算这些状态上的概率分布。
• 使用 Gumbel-Softmax 技巧采样下一状态，实现可微分采样，以支持基于梯度的训练。
• Gumbel-Softmax 的温度 τ 在训练过程中被学习，作为熵正则化项，用于控制不确定性的集中程度。
• 通过运行多次前向传播并计算输出概率的均值与方差来估计预测不确定性。
• 该方法在分类、分布外检测和强化学习任务中应用，仅需对架构进行最小改动。
• 该方法兼容 LSTM 和普通 RNN 架构，并支持贪婪策略与随机动作选择。

实验结果

研究问题

• RQ1 有限状态概率 RNN 是否能在不依赖贝叶斯神经网络的前提下，实现校准良好的不确定性估计？
• RQ2 所学习的温度参数 τ 如何影响不确定性量化与模型校准？
• RQ3 与标准 RNN 和基于 Dropout 的方法相比，ST-τ 是否能提升分布外检测性能？
• RQ4 随机状态转移机制是否能增强在部分可观察强化学习环境中的探索能力？
• RQ5 ST-τ 是否能从序列数据中学习确定性和概率性自动机？

主要发现

• 在基于 IMDB 的分布外检测任务中，ST-τ 在 O-AUPR 和 O-AUROC 指标上均达到当前最优性能，优于基线模型。
• 在训练过程中学习温度 τ 显著优于固定 τ=1，尤其在状态数较少（k=2）时效果更明显。
• 在强化学习任务中，ST-τ 在累积奖励方面优于 LSTM、VD 和 BBB 模型，且样本复杂度更低，尤其在随机策略设置下表现更优。
• 在真实世界分类任务中，ST-τ 展现出更优的校准性能，预测概率与实际正确率之间的一致性更高。
• ST-τ 能够成功从序列数据中学习确定性和概率性自动机，证明其具备建模结构化序列行为的能力。
• 消融实验表明，仅使用少量状态（k=2）并结合学习到的 τ 即可实现有效的不确定性建模，避免过拟合并保持性能。