Lightweight Neural Networks for Affordance Segmentation: Enhancement of the Decoder Module

AI summary

Overview

Research area: Embedded computer vision and lightweight deep learning for wearable robotics, focused specifically on visual affordance segmentation (VAS) — the task of splitting an object into its functionally usable parts.

Technical level: Intermediate. The paper assumes familiarity with convolutional neural networks, segmentation heads, FLOPs and multitask learning, but its argument (decoder design matters as much as backbone design) is explained in accessible terms.

One-sentence scope: The paper analyzes how the design of the segmentation head (decoder) of a lightweight network affects the accuracy-versus-compute trade-off in affordance segmentation, and shows that carefully designed heads beat a prior baseline on established benchmarks.

What This Paper Is About

Affordance segmentation requires a high level of abstraction, which normally calls for large models, but the hardware embedded in wearable robots such as prostheses and exoskeletons cannot run large models in real time. Prior work has concentrated on designing the backbone (often via Neural Architecture Search), treating the segmentation head as a minor component. This paper asks what happens when the segmentation head is treated as a first-class design variable, and proposes hardware-efficient heads that improve generalization at a comparable compute cost.

Key Contributions

1. A systematic analysis of segmentation head (SH) design under a fixed compute budget. Using a MobileNetV3 backbone and 128x128 input images, admissible SHs were limited to a FLOP count in the range of 700/800 M. Three "base block" instances were compared: depthwise separable convolution with nearest upsampling (U), depthwise separable convolution with a transposed convolution layer (T), and standard convolution with nearest upsampling (B).
2. A hierarchical, multitask reformulation of the learning problem. The network is given two output heads: one performing binary Object Segmentation (OS) and one performing the actual VAS task, with the combined loss L = L_seg + L_aff. This biases feature extraction toward isolating the object before identifying its functional parts.
3. A study of the backbone-to-head interconnection. The paper compares SHs connected to the backbone through four feature-map connections (with the 64x64 connection not implemented) against a variant that uses the 64x64 layer and drops the 8x8 connection, still maintaining four levels of feature maps.
4. Empirical results on two real-world benchmarks. The proposed heads outperform the baseline of Ragusa et al. (2021) while keeping a comparable (and in several cases much smaller) parameter count.

Main Findings

• All tested segmentation heads outperformed the baseline. The baseline of Ragusa et al. (2021) scored TOT 91.2, IIT 88.6, UMD 94.6 with 5.9 M parameters. Every proposed decoder matched or exceeded it, despite the baseline relying on a different subset of connections that prioritized high-level features and led to an increased number of parameters.
• Single-head depthwise designs were highly competitive. Decoder U scored TOT 92.6, IIT 91.0, UMD 94.6 with 1.3 M parameters; Decoder T scored TOT 92.6, IIT 90.5, UMD 95.1 with 2.0 M parameters. The paper states this proves that even without multitasking, depthwise segmentation can support the segmentation process effectively.
• "B" heads (standard convolution) benefited from their expressive power. Decoder B reached TOT 91.8, IIT 88.6, UMD 94.6 with 2.2 M parameters; Decoder BB reached TOT 92.5, IIT 89.9, UMD 95.6 with 3.6 M parameters.
• The best TOT among the single-task decoders was Decoder TB (TOT 92.8, IIT 90.8, UMD 95.2, 3.6 M parameters). Decoder UU scored TOT 92.1, IIT 89.6, UMD 94.9 with 1.6 M parameters; Decoder TT scored TOT 92.5, IIT 90.3, UMD 95.2 with 2.1 M parameters; Decoder UB scored TOT 91.7, IIT 90.3, UMD 95.2 with 3.3 M parameters.
• The auxiliary object-segmentation task was solved well by all multitask heads. On the object segmentation task, Decoder BB reached TOT 92.7, IIT 90.2, UMD 95.4, UMD_B 95.0; Decoder UU reached 93.1, 90.2, 96.0, 95.9; Decoder TT reached 93.7, 91.3, 96.2, 96.1; Decoder TB reached 93.9, 91.8, 96.2, 96.0. The authors report that all SHs discriminated object pixels from background pixels with an accuracy higher than 90% on average.
• Low-level feature connections helped individual cases but not the overall average. The variant using the 64x64 layer instead of the 8x8 layer produced: Decoder UT-X TOT 92.7 (+1.5), IIT 90.5 (+2.4), UMD 95.4 (+0.7) at 1.1 M parameters (-0.6 M); Decoder TU-X TOT 92.6 (+0.5), IIT 90.7 (+1.5), UMD 94.9 (-0.6) at 1.4 M (-0.6 M); Decoder TT-X TOT 92.6 (+0.1), IIT 90.5 (+0.2), UMD 95.1 (-0.1) at 1.5 M (-0.6 M); Decoder TB-X TOT 92.6 (-0.2), IIT 90.5 (-0.3), UMD 95.3 (+0.1) at 1.8 M (-1.8 M); Decoder BB-X TOT 92.5 (+0.0), IIT 90.3 (+0.4), UMD 95.1 (-0.5) at 1.8 M (-1.8 M). The authors state the previous feature set actually yielded the best overall average result, confirming that connections are a crucial factor for automated design strategies.
• On the motivation for lightweight design, the paper notes that large networks which still fit the memory constraints of embedded devices typically feature inference times of many seconds, which is inadequate for affordance segmentation.

Methodology in Plain English

The researchers fixed the backbone (MobileNetV3) and the input size (128x128), then constrained the segmentation head to a FLOP count in the range of 700/800 M so that different heads could be compared fairly. They defined three base block types — depthwise separable convolution with nearest upsampling (U), depthwise separable convolution with transposed convolution (T), and standard convolution with nearest upsampling (B) — and built decoders by combining these blocks into single heads or paired heads.

A key idea is that when a backbone is small, it may not disentangle the features the task needs, so the head has to work harder. To help it, the authors added a second head that solves a simpler binary problem first (is this pixel part of the object?) alongside the harder affordance problem, and trained both jointly with the combined loss L_seg + L_aff. They also varied which backbone feature-map levels feed the head, testing a four-connection configuration that omits the 64x64 layer and an alternative that uses the 64x64 layer and removes the 8x8 connection.

Training used two datasets: the UMD dataset (28,843 RGB-D images across 7 object categories, with 5,135 images forming the test set) and the IIT dataset (8,835 images at different framing and resolution settings). Following Ragusa et al. (2021), all grasping affordances were grouped into a single class and other affordances into a "don't grasp" category. The IIT data was augmented with standard geometrical and color distortions; the UMD data was augmented with a custom procedure replacing the original blue background with a different picture, isolating objects via the label segmentation mask and a color-based procedure exploiting the uniform background, followed by standard augmentation.

Evaluation used class-wise weighted pixel-wise accuracy on three test sets: the IIT test set, the original UMD, and the background-replaced UMD_B. A summary measure "TOT" was computed as the weighted average of the three accuracies, weighting IIT at 0.5 and each UMD version at 0.25.

Why This Matters

Impact on research. Most lightweight-vision work optimizes the backbone, often with Neural Architecture Search. This paper shows that the segmentation head and its connections to the backbone are not minor details: small architectural changes in the head shifted accuracy on the reported benchmarks, and a simple reformulation of the learning problem (adding an auxiliary object-segmentation head) improved results. The authors explicitly argue that connections should be a crucial factor in automated design strategies.

Real-world applications.

• Control of powered prosthetic hands and arms, where semiautonomous control can reduce user effort (Salminger et al. 2022; Tang et al. 2022; Sun et al. 2020).
• Exoskeletons and other wearable robots with multiple degrees of freedom that need onboard perception.
• Camera-based (teleceptive, contactless) sensing for assistive grasping, where the system must identify which part of an object is graspable.
• Any embedded vision system with tight compute budgets that needs fine-grained part-level understanding rather than whole-object labels.

Industry relevance. Designing decoders that reach higher accuracy at 1.3 M parameters (Decoder U) or 1.1 M parameters (Decoder UT-X) rather than the baseline's 5.9 M is directly relevant to edge-AI and robotics suppliers, where memory footprint and inference time determine whether a model can ship. The use of FLOPs as a platform-independent constraint makes the design procedure portable across embedded targets.

Future Directions

• Incorporate head and connection choices into automated architecture search. The authors state that connections should be a crucial factor in automated design strategies, but the paper only evaluates a small manual set of variants.
• Investigate the surprising low-level-connection result further. Lower-level (64x64) connections improved some individual cases and reduced parameters, yet the original feature set gave the best overall average — why this trade-off behaves as it does remains an open question.
• Validate on real embedded hardware. The paper uses FLOPs as a proxy for inference time and does not report measured latency, energy or memory-bandwidth figures on a physical wearable-robot platform.
• Extend beyond the two-task hierarchy. The hierarchical formulation in this work combines object segmentation and affordance segmentation only; whether additional intermediate tasks or other affordance groupings help is not reported.

Target Audience

Researchers and practitioners in embedded computer vision, edge AI, and wearable robotics — particularly those working on prostheses, exoskeletons and assistive devices — as well as graduate students studying lightweight neural network design. Engineers who need to balance model accuracy against strict compute and memory budgets on deployed hardware will find the design criteria and the benchmark comparison most directly useful.

Authors’ abstract

The deployment of deep neural networks for visual affordance segmentation on wearable robots poses may prove critical, due to some conflicting aspects of the problem. On one hand, affordance segmentation requires high-level abstraction capabilities, that typically involve large-size models. On the other hand, computing resources hosted on wearable robots prevent to run large-size models in real-time. The paper presents an analysis of the role of the segmentation head in the trade-off between generalization performance and compute cost. The obtained models outperform modern baseline solutions in well-known, real-world datasets while meeting low computing requirements.