[论文] 基于上下文赌博机的RL的线性求解器精度自动调优

Source: arXiv - 2601.00728v1

请提供您希望翻译的具体文本内容，我将为您进行简体中文翻译。

概述

本文介绍了一种强化学习（RL）框架，能够自动为线性求解器流水线的每一步挑选数值精度（例如 half、single、double）。通过将问题建模为 contextual bandit，系统从运行时特征（如估计的条件数）中学习，并选择仍能满足目标精度的最低成本精度。作者在迭代细化上演示了该思路，显示混合精度运行可以在保持双精度质量的同时缩短计算时间。

关键贡献

• 首个基于强化学习的数值精度自动调优器，用于科学计算内核，桥接了混合精度计算与强化学习。
• 上下文赌博机（Contextual‑bandit）模型，将求解器状态特征视为上下文，将精度选择视为动作，实现了使用简单 Q 表的快速在线学习。
• ε‑贪婪动作选择，在探索（尝试新精度）与利用（使用已知最佳配置）之间取得平衡。
• 多目标奖励设计，同时优化求解精度和计算成本。
• 对未见线性系统的实证验证，显示在保持双精度准确性的前提下实现了持续的成本降低。
• 可推广的流水线，可扩展至迭代细化之外的其他数值算法。

方法论

1. 特征提取 – 在每次细化迭代之前，收集易于计算的统计信息（例如，近似条件数、矩阵范数、残差大小）。这些构成低维状态向量。
2. 离散化 – 将连续特征空间划分为有限的“上下文”集合，以便使用经典的表格 Q‑learning 方法。
3. 动作空间 – 每个动作对应求解器子步骤的特定精度配置（例如，以单精度计算残差，以半精度求解校正）。
4. 奖励函数 – (i) 负误差（鼓励精度）和 (ii) 负运行时间或 FLOP 计数（鼓励速度）的加权和。权重让用户能够在两者之间进行优先级设置。
5. 学习循环 – 使用 epsilon‑greedy 策略，调优器选择一个动作，运行求解器步骤，观察奖励，并通过增量平均更新 (上下文, 动作) 对应的 Q 值。经过大量求解后，Q 表收敛到每个上下文的最佳精度策略。
6. 部署 – 在推理时，调优器只需查表获取当前上下文的最佳动作，几乎没有开销。

结果与发现

关键要点： RL 调优器在矩阵条件良好时会积极降低精度，而在更困难的情况会自动恢复更高精度。Q 表保持紧凑（仅几千字节），并能在训练期间未见的问题规模和分布上进行泛化。

实际意义

• 性能关键库（例如 PETSc、Trilinos）可以嵌入 bandit 调谐器，以自动利用混合精度硬件（Tensor Cores、bfloat16 单元），无需手动调优。
• 边缘和嵌入式设备 计算资源受限，通过动态降低精度可以在保持功耗预算的同时实现双精度水平的准确性。
• 编译器/运行时框架（LLVM、TVM、OneAPI）可以提供 “precision‑autotune” 过程，自动生成特征提取和 Q 表查找代码。
• 云服务 提供线性求解器 API 时，可降低每个请求的 CPU/GPU 使用率，直接转化为成本节约。
• 该方法足够轻量，可实现 在线学习：生产系统在新问题实例到达时持续优化其 Q 表，以适应硬件升级或工作负载变化。

限制与未来工作

• 状态离散化 对于非常高维的特征集合可能变得过于粗糙，可能会错过细微的精度权衡。
• 当前的奖励仅平衡准确率与运行时间；扩展到内存带宽、能耗或数值稳定性将提升适用范围。
• 实验聚焦于 迭代细化；将相同的 bandit 框架应用于其他核函数（例如特征值求解器、非线性优化）仍有待验证。
• 该方法假设廉价的特征估计是可获得的；对于某些问题（例如高度稀疏或分布式矩阵），提取这些特征可能并非易事。
• 未来工作可以探索 函数逼近强化学习（例如深度 Q‑网络）以处理连续状态空间和更大的动作集合，以及 元学习 来在硬件代际之间转移策略。

作者

• Erin Carson
• Xinye Chen

论文信息

• arXiv ID: 2601.00728v1
• 分类: cs.LG, math.NA
• 出版日期: 2026年1月2日
• PDF: 下载 PDF