Wrapper-Aware Rate-Distortion Optimization in Feature Coding for Machines

TL;DR

The paper tackles wrapper-aware rate-distortion optimization for feature coding for machines (FCM), where a non-differentiable inner codec is sandwiched between a feature encoder and a reconstruction wrapper. It introduces WA-RDO, replacing the standard SSE distortion with a wrapper-aware weighted distortion derived from the Jacobian of the restoration wrapper: $\| \mathbf J_g(\mathbf z)(\hat{\mathbf z}(\boldsymbol{\theta}) - \mathbf z) \|_2^2$, while keeping the rate term. To make this practical, it sketches the Jacobian with a random matrix $\mathbf S$ and computes an importance map $\mathbf h(\mathbf z)$ as the diagonal of $\mathbf H_s(\mathbf z) = \mathbf J_s(\mathbf z)^{\top} \mathbf J_s(\mathbf z)$, using per-block optimization with $\text{diag}(\mathbf h)$ and defining $\tau$ and $\lambda$ to balance terms. The authors further propose two simplifications: IWA-RDO (reusing the same importance map across GOPs) and FWA-RDO (freezing the wrapper to obtain a fixed $\mathbf h_a$), yielding substantial reductions in encoder complexity. Experiments on MPEG FCTM show WA-RDO with HEVC inner codec matching the VVC-based anchor under SSE-RDO and closing the codec-generation gap for AVC with negligible runtime overhead, demonstrating practical viability for wrapper-aware compression in FCM.

Abstract

Feature coding for machines (FCM) is a lossy compression paradigm for split-inference. The transmitter encodes the outputs of the first part of a neural network before sending them to the receiver for completing the inference. Practical FCM methods ``sandwich'' a traditional codec between pre- and post-processing neural networks, called wrappers, to make features easier to compress using video codecs. Since traditional codecs are non-differentiable, the wrappers are trained using a proxy codec, which is later replaced by a standard codec after training. These codecs perform rate-distortion optimization (RDO) based on the sum of squared errors (SSE). Because the RDO does not consider the post-processing wrapper, the inner codec can invest bits in preserving information that the post-processing later discards. In this paper, we modify the bit-allocation in the inner codec via a wrapper-aware weighted SSE metric. To make wrapper-aware RDO (WA-RDO) practical for FCM, we propose: 1) temporal reuse of weights across a group of pictures and 2) fixed, architecture- and task-dependent weights trained offline. Under MPEG test conditions, our methods implemented on HEVC match the VVC-based FCM state-of-the-art, effectively bridging a codec generation gap with minimal runtime overhead relative to SSE-RDO HEVC.

Figures (6)

• Figure 1: Sandwich-based FCM setup. The feature reduction and restoration blocks are the wrappers. Our additions appear in blue.
• Figure 2: (a) Bit difference per macroblock using AVC with SSE-RDO and WA-RDO (feature channels arranged as an image, larger values mean SSE-RDO invests more bits than WA-RDO), and (b) WA-RDO importance map, defined for each entry of the features to compress. Since some feature channels have low importance for the restoration wrapper, WA-RDO invests fewer bits to reconstruct them.
• Figure 3: Correlation coefficient (CC) between importance maps. (a) Average CC across the SFU dataset as a function of the separation to the I frame with $95\%$ confidence intervals, and (b) distribution of the CC for a subset of the pairs of I frames in SFU and images in MPEG-OIV6. We observe temporal and architectural consistencies.
• Figure 4: RD curves for (a-b) class D of SFU choi2021dataset and (c-d) TVD gao2022open, using FCTM-HEVC and FCTM-AVC with WA-RDO and SSE-RDO. We also show the FCTM-VVC anchor, remote inference using VVC, and local inference (dashed line). Our methods improve over SSE-RDO.
• Figure 5: Average BD-accuracy (%) against FCTM-AVC using SSE-RDO across (a) all the datasets in our experimental setup, and (b) only the video datasets. WA-RDO with HEVC achieves the same rate–accuracy trade-off as the FCTM-VVC anchor. WA-RDO with AVC matches the performance of SSE-RDO FCTM-HEVC.