Work-in-Progress: Cooperative MLP-Mixer Networks Inference On Heterogeneous Edge Devices through Partition and Fusion

Abstract

As a newly proposed DNN architecture, MLP-Mixer is attracting increasing attention due to its competitive results compared to CNNs and attention-base networks in various tasks. Although MLP-Mixer only contains MLP layers, it still suffers from high communication costs in edge computing scenarios, resulting in long inference time. To improve the inference performance of an MLP-Mixer model on correlated resource-constrained heterogeneous edge devices, this paper proposes a novel partition and fusion method specific for MLP-Mixer layers, which can significantly reduce the communication costs. Experimental results show that, when the number of devices increases from 2 to 6, our partition and fusion method can archive 1.01-1.27x and 1.54-3.12x speedup in scenarios with heterogeneous and homogeneous devices, respectively.