Memory Efficient Matting with Adaptive Token Routing

TL;DR

MEMatte is a novel memory-efficient image matting framework that uses adaptive token routing to process high-resolution images effectively, reducing memory and computational costs while maintaining high accuracy.

Contribution

The paper introduces MEMatte, which employs a dynamic routing mechanism and a lightweight token refinement module to enable high-resolution image matting with significantly reduced memory and latency.

Findings

Reduces memory usage by approximately 88% on high-resolution images.

Cuts latency by about 50% compared to existing methods.

Outperforms current state-of-the-art in high-resolution and real-world datasets.

Abstract

Transformer-based models have recently achieved outstanding performance in image matting. However, their application to high-resolution images remains challenging due to the quadratic complexity of global self-attention. To address this issue, we propose MEMatte, a \textbf{m}emory-\textbf{e}fficient \textbf{m}atting framework for processing high-resolution images. MEMatte incorporates a router before each global attention block, directing informative tokens to the global attention while routing other tokens to a Lightweight Token Refinement Module (LTRM). Specifically, the router employs a local-global strategy to predict the routing probability of each token, and the LTRM utilizes efficient modules to simulate global attention. Additionally, we introduce a Batch-constrained Adaptive Token Routing (BATR) mechanism, which allows each router to dynamically route tokens based on image content and the stages of attention block in the network. Furthermore, we construct an ultra high-resolution image matting dataset, UHR-395, comprising 35,500 training images and 1,000 test images, with an average resolution of $4872\times6017$. This dataset is created by compositing 395 different alpha mattes across 11 categories onto various backgrounds, all with high-quality manual annotation. Extensive experiments demonstrate that MEMatte outperforms existing methods on both high-resolution and real-world datasets, significantly reducing memory usage by approximately 88% and latency by 50% on the Composition-1K benchmark. Our code is available at https://github.com/linyiheng123/MEMatte.