Zeros can be Informative: Masked Binary U-Net for Image Segmentation on Tensor Cores

TL;DR

This work tackles real-time, high-resolution image segmentation on edge devices by combining Masked Binary U‑Net (MBU‑Net) with a Tensor Core–aware deployment. Two empirical insights—the necessity of an explicit zero state and uniform layer sensitivity to quantization—drive a cost-aware masking strategy that selectively masks cheap layers to balance accuracy and compute. A subtractive bit-encoding scheme enables end-to-end, high-throughput inference of masked binary weights with binary activations on commodity GPUs via Tensor Core BMMA/WMMA, achieving an average Dice accuracy within $\sim$3% of full precision while delivering $2.04\times$ speedup and $3.54\times$ energy reductions over 16‑bit FP U‑Net across three segmentation benchmarks. The approach demonstrates practical viability for real-time edge deployments in AR/VR, robotics, and autonomous systems, with potential benefits for ASIC accelerators due to reduced data movement and deterministic latency.

Abstract

Real-time image segmentation is a key enabler for AR/VR, robotics, drones, and autonomous systems, where tight accuracy, latency, and energy budgets must be met on resource-constrained edge devices. While U-Net offers a favorable balance of accuracy and efficiency compared to large transformer-based models, achieving real-time performance on high-resolution input remains challenging due to compute, memory, and power limits. Extreme quantization, particularly binary networks, is appealing for its hardware-friendly operations. However, two obstacles limit practicality: (1) severe accuracy degradation, and (2) a lack of end-to-end implementations that deliver efficiency on general-purpose GPUs. We make two empirical observations that guide our design. (1) An explicit zero state is essential: training with zero masking to binary U-Net weights yields noticeable sparsity. (2) Quantization sensitivity is uniform across layers. Motivated by these findings, we introduce Masked Binary U-Net (MBU-Net), obtained through a cost-aware masking strategy that prioritizes masking where it yields the highest accuracy-per-cost, reconciling accuracy with near-binary efficiency. To realize these gains in practice, we develop a GPU execution framework that maps MBU-Net to Tensor Cores via a subtractive bit-encoding scheme, efficiently implementing masked binary weights with binary activations. This design leverages native binary Tensor Core BMMA instructions, enabling high throughput and energy savings on widely available GPUs. Across 3 segmentation benchmarks, MBU-Net attains near full-precision accuracy (3% average drop) while delivering 2.04x speedup and 3.54x energy reductions over a 16-bit floating point U-Net.

Figures (12)

• Figure 1: Overview. Obtained through a cost-aware masking strategy and optimized for Tensor Core on GPUs, MBU-Net achievs full-precision level accuracy and binary-level efficiency.
• Figure 2: Left: Conventional U-Net architecture exemplified with segmentation task. $c$: number of channels, $W/H$: feature width / height. Right: internal components of each abstracted block.
• Figure 3: (a) Statistics of the weights masked to zero. (b) Accuracy and computing demand distribution for all layer configurations. (c) Shapley analysis: marginal improvement in Dice score from masking each binary layer. Only cases with fewer than five masked layers are included for clarity.
• Figure 4: Per-layer cost profiling. (a) Number of addition and multiplication operations in each layer with a 512$\times$512 image as a representative example. (b) Number of parameters in each layer.
• Figure 5: Latency per image inference (milliseconds). Lower is better.