[论文解读] POET: Training Neural Networks on Tiny Devices with Integrated Rematerialization and Paging

Shishir G. Patil, Paras Jain+3 more|arXiv (Cornell University)|Jul 15, 2022

Advanced Neural Network Applications被引用 6

一句话总结

POET 提出了一种基于混合整数线性规划（MILP）的框架，通过联合优化重计算与分页机制，在内存受限的边缘设备上实现大型神经网络的能量最优训练。该方法相比独立使用的方法，能量消耗最高降低40%，并可在仅配备32 KB RAM的Cortex-M类设备上训练ResNet-18和BERT，其能量效率与可扩展性优于现有边缘训练方法。

ABSTRACT

Fine-tuning models on edge devices like mobile phones would enable privacy-preserving personalization over sensitive data. However, edge training has historically been limited to relatively small models with simple architectures because training is both memory and energy intensive. We present POET, an algorithm to enable training large neural networks on memory-scarce battery-operated edge devices. POET jointly optimizes the integrated search search spaces of rematerialization and paging, two algorithms to reduce the memory consumption of backpropagation. Given a memory budget and a run-time constraint, we formulate a mixed-integer linear program (MILP) for energy-optimal training. Our approach enables training significantly larger models on embedded devices while reducing energy consumption while not modifying mathematical correctness of backpropagation. We demonstrate that it is possible to fine-tune both ResNet-18 and BERT within the memory constraints of a Cortex-M class embedded device while outperforming current edge training methods in energy efficiency. POET is an open-source project available at https://github.com/ShishirPatil/poet

研究动机与目标

• 在内存稀缺、依赖电池供电的边缘设备（如智能手机和微控制器）上实现大型现代神经网络的训练。
• 解决边缘训练中反向传播带来的高内存与高能耗需求，当前这些限制使得模型仅能支持小型架构（如逻辑回归）。
• 在一个统一的能量最小化框架中，联合优化两种激活内存减少技术——重计算与分页。
• 满足真实边缘部署中典型的硬性运行时约束，例如在空闲时段或充电期间进行训练。
• 提供数学上正确且保持数值精度的解决方案，无需对现有深度学习模型进行架构修改即可兼容。

提出的方法

• 将边缘训练建模为一个混合整数线性规划（MILP）问题，在内存与运行时约束下联合优化重计算与分页决策。
• 在目标硬件上对模型算子进行性能分析，构建细粒度的能量、内存与执行时间成本模型。
• 使用MILP求解器寻找能量最优的调度方案，在满足内存与时间预算的前提下最小化总能量消耗。
• 将分页操作的流水线处理与重计算相结合，隐藏延迟并满足严格的运行时约束。
• 通过泛化算子级别的成本模型，支持线性与非线性模型，包括ResNet-18与BERT。
• 采用统一的搜索空间，结合廉价操作（如ReLU）的重计算与昂贵操作（如卷积）的分页，以降低能量开销。

实验结果

研究问题

• RQ1 重计算与分页能否被联合优化，以实现比单独使用任一技术更低的能量消耗？
• RQ2 在仅配备32 KB RAM的边缘设备上，是否可能满足严格运行时约束的同时训练复杂模型（如BERT与ResNet-18）？
• RQ3 与启发式或顺序方法相比，重计算与分页的集成如何影响能量效率与内存使用？
• RQ4 基于MILP的建模是否能在现实硬件约束下找到边缘训练的可证明最优调度方案？
• RQ5 在调整内存与时间预算时，能量效率与运行时性能之间的权衡关系如何？

主要发现

• 在A72平台上，POET在严苛内存预算下训练ResNet-18时，能量开销相比全内存基线最高降低141%。
• 在相同条件下，POET的能量开销低于全内存训练的1%，而基于启发式的Capuchin方法则产生73%至141%更高的能量消耗。
• 在Jetson TX2上，POET相比先前仅支持线性模型的POFO方法，将峰值内存消耗降低8.3%，吞吐量提升13%。
• POET的集成方案相比单独使用重计算或分页，能量消耗最高降低40%，尤其在严苛运行时约束下（如每轮0.5 ms）表现更优。
• 即使在复杂模型（如BERT与ResNet-18）下，MILP求解器也能在普通硬件上于十分钟内找到最优解。
• POET成功在仅配备32 KB RAM的Cortex-M类设备上训练了ResNet-18与BERT，证明了其在隐私保护型本地个性化中的可行性。