LayerShuffle: Enhancing Robustness in Vision Transformers by Randomizing Layer Execution Order
This addresses robustness issues for distributed neural network architectures where execution order is uncertain or parts may fail, though it is incremental as it builds on existing vision transformer methods.
The paper tackles the problem of making vision transformers robust to arbitrary layer execution orders and pruning, achieving this by randomizing attention module order during training, which results in a 20% accuracy reduction at the same model size.
Due to their architecture and how they are trained, artificial neural networks are typically not robust toward pruning or shuffling layers at test time. However, such properties would be desirable for different applications, such as distributed neural network architectures where the order of execution cannot be guaranteed or parts of the network can fail during inference. In this work, we address these issues through a number of training approaches for vision transformers whose most important component is randomizing the execution order of attention modules at training time. With our proposed approaches, vision transformers are capable to adapt to arbitrary layer execution orders at test time assuming one tolerates a reduction (about 20\%) in accuracy at the same model size. We analyse the feature representations of our trained models as well as how each layer contributes to the models prediction based on its position during inference. Our analysis shows that layers learn to contribute differently based on their position in the network. Finally, we layer-prune our models at test time and find that their performance declines gracefully. Code available at https://github.com/matfrei/layershuffle.