Segmentation-Aware Convolutional Networks Using Local Attention Masks

Adam W. Harley; Konstantinos G. Derpanis; Iasonas Kokkinos

Source-linked AI summary

Segmentation-Aware Convolutional Networks Using Local Attention Masks

Adam W. Harley, Konstantinos G. Derpanis, Iasonas Kokkinos

arXiv:1708.04607v1Aug 15, 2017cs.CV

TL;DR

CNNs tend to smooth information across regions, limiting spatial precision in dense prediction. This paper learns segmentation embeddings, converts them into local attention masks, and integrates them into convolution. The resulting segmentation-aware CNN improves semantic segmentation and optical flow, matching DenseCRF performance while being faster and simpler and producing sharper flow responses.

Problem
CNNs smooth information across regions, creating spatially imprecise outputs for dense prediction tasks.
Method
The method learns embedding-based segmentation cues, constructs local foreground masks, and uses them to make convolution selectively attend to each neuron’s region.
Results
Segmentation-aware convolution yields systematic improvements in semantic segmentation and optical flow, matching dense-CRF performance while being faster and simpler and producing sharper flow responses.
Takeaways & Limitations
A single end-to-end CNN framework can incorporate segmentation information across its architecture for both discrete labelling and continuous regression.

Abstract

from arXiv · show

We introduce an approach to integrate segmentation information within a convolutional neural network (CNN). This counter-acts the tendency of CNNs to smooth information across regions and increases their spatial precision. To obtain segmentation information, we set up a CNN to provide an embedding space where region co-membership can be estimated based on Euclidean distance. We use these embeddings to compute a local attention mask relative to every neuron position. We incorporate such masks in CNNs and replace the convolution operation with a "segmentation-aware" variant that allows a neuron to selectively attend to inputs coming from its own region. We call the resulting network a segmentation-aware CNN because it adapts its filters at each image point according to local segmentation cues. We demonstrate the merit of our method on two widely different dense prediction tasks, that involve classification (semantic segmentation) and regression (optical flow). Our results show that in semantic segmentation we can match the performance of DenseCRFs while being faster and simpler, and in optical flow we obtain clearly sharper responses than networks that do not use local attention masks. In both cases, segmentation-aware convolution yields systematic improvements over strong baselines. Source code for this work is available online at http://cs.cmu.edu/~aharley/segaware.

1. Introduction

Segmentation-aware CNNs address feature smoothness by using learned segmentation cues to focus convolution on each neuron's local region. The approach is differentiable, GPU-efficient, and improves dense classification and regression tasks.

1. Introduction: CNNs’ repeated pooling and subsampling produce smooth outputs by mixing information across regions, while prior work mainly addressed lost resolution.The paper targets smoothness as a distinct spatial-precision problem.
1. Introduction: Segmentation-aware convolution selectively attends to a neuron’s region instead of mixing foreground and contextual signals, reducing feature blurring before output post-processing.Local foreground-background masks act as gating mechanisms and can be integrated at early and later CNN stages.
1. Introduction: Learned embeddings provide segmentation information by bringing pixels from the same object closer and pixels from different objects farther apart in feature space.A contrastive side loss trains the embedding branch, whose distances are then used to construct masks.
1. Introduction: The method learns segmentation cues, constructs local foreground masks, and combines them with convolution in a unified, differentiable, GPU-optimized pipeline.This design is intended for both discrete and continuous prediction tasks.
1. Introduction: The approach yields substantial gains on semantic segmentation and optical flow while requiring only minimal modifications to existing CNN architectures.These experiments cover dense discrete labelling and dense regression.

2. Related work

The work connects segmentation-aware filtering to metric learning, locally gated convolutions, and methods for improving CNN spatial precision. Its distinguishing choice is to learn embeddings inside the network and use them to measure context relevance.

2. Related work: The embeddings target same-object invariance across pixels, requiring invariance not only to viewpoint changes but also to interior appearance details.This differs from metric-learning work focused primarily on one-to-one correspondences across views.
2. Related work: Unlike handcrafted descriptors and generic gated convolutions, this approach uses convolutional embeddings to construct interpretable local attention masks and support task-agnostic pre-training.Masks are efficiently constructed in batch mode around each neuron.
2. Related work: Segmentation-aware convolution adjusts a neuron’s field of view toward its local foreground, whereas atrous convolution widens the field of view generally.Both are preventative strategies for spatial imprecision.
2. Related work: The method belongs to a broader distinction between preventing spatial imprecision within CNNs and curing it afterward with upsampling, multiresolution features, edges, or CRFs.Its design performs sharpening inside the network rather than solely as post-processing.

3. Technical approach

The approach learns semantic pixel embeddings, converts embedding distances into local masks, and uses those masks for adaptive smoothing, CRF filtering, and convolution. These operations selectively emphasize information from the reference pixel’s region while remaining differentiable and learnable within CNNs.

3.1. Learning segmentation cues: The method learns pixel embeddings in which same-object pairs are nearby and different-object pairs are separated by a margin-based loss.The embedding function maps RGB pixels into a 64-dimensional feature space; the loss uses α = 0.5 and β = 2.
3.2. Segmentation-aware bilateral filtering: Embedding distances are converted into local foreground-background masks, with the reference pixel assigned foreground weight 1.The learned semantic masks are reported as more reliable than masks based on color distances.
3.2. Segmentation-aware bilateral filtering: Segmentation-aware smoothing multiplies neighboring features by normalized mask weights, generalizing bilateral filtering and adaptively preserving region structure.The formulation can use learned embeddings rather than fixed color vectors, and the embedding and mask parameters can be trained end-to-end.
3.3. Segmentation-aware CRFs: Replacing color features with learned embeddings lets a CRF use small standard convolutions instead of wide permutohedral-lattice filters while retaining comparable or better performance.The learned embedding signal supports filters such as 13 × 13, reducing computation time by half and enabling backpropagation through the CRF.
3.4. Segmentation-aware convolution: Segmentation-aware convolution applies a learned filter after pointwise masking, interpolating between bilateral filtering and standard convolution.An all-ones learned filter recovers bilateral filtering, while an all-ones mask recovers standard convolution.

4. Implementation details

The implementation integrates an embedding network with a task-specific CNN, using learned embedding distances to create masks that modify convolution. The pipeline is implemented efficiently through convolution-like matrix operations and supports both bilateral filtering and segmentation-aware convolution.

4.1. Network architecture: The embedding network generates representations at multiple scales, combines intermediate embeddings into a final embedding, and uses pairwise-distance losses to guide learning.Intermediate-layer losses provide earlier layers with stronger training signals.
4.1. Network architecture: Embedding-derived distances are converted into local attention masks that intercept convolutions in a task-specific network, with joint backpropagation fine-tuning both networks.The embedding network and task-specific network are trained through a final objective.
4.1. Network architecture: The masks can be applied either to final-layer outputs through segmentation-aware bilateral filtering or throughout the network by converting convolutions into segmentation-aware convolutions.The more intrusive convolutional modification can be initialized without detrimental effect using λ = 0.
4.2. Efficient implementation: The pipeline reduces mask creation and masking to image-to-column transformations, pointwise matrix multiplication, summation, and final matrix multiplication, enabling GPU-efficient implementation.Segmentation-aware convolution multiplies masked H · W × K values by a K×F filter matrix.
4.2. Efficient implementation: Mask creation computes central-to-neighbor embedding distances for each local feature patch, arranging each H × W feature map as an H · W × K matrix.The distance values are passed through an exponential function with learned hardness λ to realize the mask term.

5. Evaluation

The evaluation tests minimally modified CNNs with segmentation awareness on semantic segmentation and optical flow. Across both tasks, the approach improves strong baselines, with semantic segmentation matching dense CRF quality more simply and quickly and optical flow producing sharper predictions.

The method yields substantial gains over strong baselines on both semantic segmentation and optical flow.The experiments cover dense discrete labelling and dense regression with minimally modified CNN architectures.
5.1. Semantic segmentation: 1.3% IOU improvement is obtained by applying the learned masking process four times, although this adds approximately 500 ms of computation.A dense CRF performs slightly worse and takes approximately 1 second.
5.1. Semantic segmentation: 1.6% IOU improvement results when all DeepLab layers are segmentation-aware, compared with approximately 1.0% from making only FC6 segmentation-aware.The all-layer modification costs +200 ms, while the FC6-only modification costs +100 ms.
5.1. Semantic segmentation: Segmentation-aware convolution improves accuracy beyond 10 pixels from object boundaries, whereas bilateral filtering helps most within 5 pixels.Trimaps measure mean IOU in narrow bands around object boundaries across half-widths 1 to 40.
5.1. Semantic segmentation: 2.0% improvement over DeepLab is achieved on the PASCAL VOC test server, while segmentation-aware modifications are 0.1% superior to dense CRF post-processing and twice as fast.The segmentation-aware method runs in 0.5s rather than 1s and is described as simpler.
5.2. Optical flow: 2.78 to 2.26 average end-point error is achieved on FlyingChairs, an 18% reduction, while average angular error falls from 15.58 to 9.54.The resulting flow fields are visibly sharper, especially near object boundaries, without aggressive data augmentation.

6. Conclusion

Segmentation-aware convolution integrates segmentation information throughout CNNs to prevent feature blurring before it occurs. The end-to-end architecture competes with structured prediction methods while extending to continuous prediction.

Segmentation-aware convolution integrates segmentation information throughout a deep architecture and avoids feature blurring before it happens.
The full segmentation-aware architecture can be trained end-to-end and extends beyond segmentation-specific structured prediction algorithms to continuous prediction.

Supplementary Material

The supplementary material includes author information and an implementation comparison between standard and segmentation-aware convolution. The latter adds distance-based processing to the usual image-to-column workflow.

The supplementary material lists Adam W. Harley from Carnegie Mellon University and Iasonas Kokkinos from Facebook AI Research.
Segmentation-aware convolution augments the standard im2col-and-matrix-multiplication workflow with an image-to-distance transformation on embeddings.