Sparse Mixture of Experts (MoEs) [1]} replace MLP layers in some Transformer blocks with sparse MoE layers that are partitioned across multiple devices, with each device storing a separate MLP (expert) that processes a subset of patches (tokens). Routers are used to direct subsets of patches to each expert based on their specialization. During training, [2]} proposed batch prioritized routing for sparse MOEs applied to ViTs to not route some patches to any expert (i.e., discard them). They show that this approach can reduce total training Floating Point Operations (FLOPs) by 20%. While our method aims to also drop patches and increase efficiency, we do not rely on routing patches to experts. Compared to sparse MoEs, patch sampling in this work can be viewed as determining weather to route a patch to the network or to discard it entirely.
