[论文解读] X-TrainCaps: Accelerated Training of Capsule Nets through Lightweight Software Optimizations.

Alberto Marchisio, Beatrice Bussolino+5 more|arXiv (Cornell University)|May 24, 2019

Advanced Neural Network Applications参考文献 21被引用 4

一句话总结

X-TrainCaps 通过轻量级软件优化加速了胶囊网络的训练，包括一种新型学习率策略（WarmAdaBatch）和参数共享，以减少冗余连接。其训练速度比原始 CapsNet 快 58.6%，同时准确率提升 0.9%，在速度与性能之间实现了帕累托最优的权衡。

ABSTRACT

Convolutional Neural Networks (CNNs) are extensively in use due to their excellent results in various machine learning (ML) tasks like image classification and object detection. Recently, Capsule Networks (CapsNets) have shown improved performances compared to the traditional CNNs, by encoding and preserving spatial relationships between the detected features in a better way. This is achieved through the so-called Capsules (i.e., groups of neurons) that encode both the instantiation probability and the spatial information. However, one of the major hurdles in the wide adoption of CapsNets is its gigantic training time, which is primarily due to the relatively higher complexity of its constituting elements. In this paper, we illustrate how can we devise new optimizations in the training process to achieve fast training of CapsNets, and if such optimizations affect the network accuracy or not. Towards this, we propose a novel framework that employs lightweight software-level optimizations, including a novel learning rate policy called WarmAdaBatch that jointly performs warm restarts and adaptive batch size, as well as weight sharing for capsule layers to reduce the hardware requirements of CapsNets by removing unused/redundant connections and capsules, while keeping high accuracy through tests of different learning rate policies and batch sizes. We demonstrate that one of the solutions generated by X-TrainCaps framework can achieve 58.6% training time reduction while preserving the accuracy (even 0.9% accuracy improvement), compared to the CapsNet in the original paper by Sabour et al. (2017), while other Pareto-optimal solutions can be leveraged to realize trade-offs between training time and achieved accuracy.

研究动机与目标

• 解决胶囊网络训练时间过长的问题，这是其广泛应用的主要障碍。
• 通过参数共享消除冗余连接和胶囊，降低胶囊网络的硬件资源需求。
• 通过智能超参数调度，在加速训练的同时保持或提升模型准确率。
• 利用帕累托最优解探索训练速度与模型准确率之间的高效权衡。
• 开发一种轻量级且与现有 CapsNet 架构兼容的软件级优化框架。

提出的方法

• 提出 WarmAdaBatch，一种新型学习率策略，结合学习率热重启与训练过程中的自适应批量大小调整。
• 在胶囊层中实现参数共享，以消除未使用或冗余的连接与胶囊，降低计算开销。
• 通过系统性消融研究，在不同学习率策略与批量大小下评估速度与准确率之间的权衡。
• 设计一种框架，将这些优化无缝集成到训练流程中，而无需修改核心 CapsNet 架构。
• 使用帕累托优化方法识别多个高效训练配置，以平衡速度与准确率。
• 在标准基准上评估该框架，以衡量训练时间的减少和准确率的保持情况。

实验结果

研究问题

• RQ1 轻量级软件优化能否在不降低准确率的情况下显著缩短 CapsNet 的训练时间？
• RQ2 与标准学习率调度策略相比，所提出的 WarmAdaBatch 学习率策略在收敛速度和最终准确率方面表现如何？
• RQ3 在胶囊层中应用参数共享在多大程度上能降低硬件需求，同时保持模型性能？
• RQ4 CapsNet 训练中训练时间与准确率之间的帕累托最优权衡是什么？
• RQ5 自适应批量大小与热重启的结合是否能带来更快的收敛速度和更好的泛化能力？

主要发现

• X-TrainCaps 相较于 Sabour 等人（2017）提出的原始 CapsNet，实现了 58.6% 的训练时间减少。
• 优化后的模型相比原始 CapsNet，测试准确率提升了 0.9%，表明泛化能力增强。
• 该框架识别出多个帕累托最优解，使用户可根据自身需求选择兼顾速度与准确率的配置。
• 参数共享有效减少了胶囊层中的参数数量与计算量，降低了硬件需求。
• WarmAdaBatch 通过动态调整批量大小并重启学习过程，在加速收敛方面优于标准学习率调度策略。
• 这些优化在不同超参数设置下均表现有效，展现出良好的鲁棒性与泛化能力。