# FasTrCaps: An Integrated Framework for Fast yet Accurate Training of Capsule Networks

## TL;DR

FasTrCaps is a framework that optimizes the training process of Capsule Networks by integrating lightweight techniques and a novel learning rate policy, significantly reducing training time while maintaining or improving accuracy.

## Contribution

The paper introduces FasTrCaps, a new framework combining multiple optimizations and a novel learning rate policy to accelerate Capsule Network training with minimal accuracy loss.

## Key findings

- Achieved 58.6% reduction in training time.
- Preserved or slightly improved accuracy on MNIST.
- Demonstrated Pareto-optimal trade-offs between training speed and accuracy.

## Abstract

Recently, Capsule Networks (CapsNets) have shown improved performance compared to the traditional Convolutional Neural Networks (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 their gigantic training time, which is primarily due to the relatively higher complexity of their new constituting elements that are different from CNNs. In this paper, we implement different optimizations in the training loop of the CapsNets, and investigate how these optimizations affect their training speed and the accuracy. Towards this, we propose a novel framework FasTrCaps that integrates multiple lightweight optimizations and a novel learning rate policy called WarmAdaBatch (that jointly performs warm restarts and adaptive batch size), and steers them in an appropriate way to provide high training-loop speedup at minimal accuracy loss. We also propose weight sharing for capsule layers. The goal is 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 the FasTrCaps framework can achieve 58.6% reduction in the training time, while preserving the accuracy (even 0.12% accuracy improvement for the MNIST dataset), compared to the CapsNet by Google Brain. The Pareto-optimal solutions generated by FasTrCaps can be leveraged to realize trade-offs between training time and achieved accuracy. We have open-sourced our framework on https://github.com/Alexei95/FasTrCaps.