图像分割之Dense Prediction with Attentive Feature Aggregation

Dense Prediction with Attentive Feature Aggregation

原始文档: https://www.yuque.com/lart/papers/xnqoi0

偶然间从 arxiv 上翻到的论文, 可以看做是对之前工作 Hierarchical multi-scale attention for semantic segmentation 的一个扩展.

从摘要读论文

Aggregating information from features across different layers is an essential operation for dense prediction models.

本文重点关注与跨层的特征集成的问题.

Despite its limited expressiveness, feature concatenation dominates the choice of aggregation operations.

虽说是特征拼接, 但是大多数还会跟一些复杂的卷积结构.

In this paper, we introduce Attentive Feature Aggregation (AFA) to fuse different network layers with more expressive non-linear operations. AFA exploits both spatial and channel attention to compute weighted average of the layer activations.

核心模块之 AFA. 使用空间和通道注意力来加权求和不同层的特征. 从而构建一种非线性的集成操作.

Inspired by neural volume rendering, we extend AFA with Scale-Space Rendering (SSR) to perform late fusion of multi-scale predictions.

这里提到一个很有趣的点, 是用来融合多尺度预测的结构. 结构的设计借鉴了神经体渲染的想法(我不太了解这个方面).

AFA is applicable to a wide range of existing network designs.

由于 AFA 本身就是一个模型通用的模块, 所以可以很容易的迁移到不同的模型中来实现特征的跨层集成.

Our experiments show consistent and significant improvements on challenging semantic segmentation benchmarks, including Cityscapes, BDD100K, and Mapillary Vistas, at negligible computational and parameter overhead. In particular, AFA im-proves the performance of the Deep Layer Aggregation (DLA) model by nearly 6% mIoU on Cityscapes. Our experimental analyses show that AFA learns to progressively refine segmentation maps and to improve boundary details, leading to new state-of-the-art results on boundary detection benchmarks on BSDS500 and NYUDv2.

尝试了分割任务和边缘检测任务.

主要内容

• We propose Attentive Feature Aggregation (AFA) as a non-linear feature fusion operation to replace the prevailing tensor concatenation or summation strategies.
  • Our attention module uses both spatial and channel attention to learn and predict the importance of each input signal during fusion. Aggregation is accomplished by computing a linear combination of the input features at each spatial location, weighted by their relevance.
  • Compared to linear fusion operations, our AFA module can take into consideration complex feature interactions and attend to different feature levels depending on their importance.
  • AFA introduces negligible computation and parameter overhead and can be easily used to replace fusion operations in existing methods, such as skip connections.
  • Unlike linear aggregation, our AFA module leverages extracted spatial and channel information to efficiently select the essential features and to increase the receptive field at the same time.
• Inspired by neural volume rendering [Volume rendering, Nerf: Representing scenes as neural radiance fields for view synthesis], we propose Scale-Space Rendering (SSR) as a novel attention computation mechanism to fuse multi-scale predictions.
  • _We treat those predictions as sampled data in scale-space and design a coarse-to-fine attention concept to render final predictions. _(这个想法很有意思. 把最终预测的获取看做是一个从尺度空间中采样不同尺度的预测来渲染最终预测的问题)
  • Repeated use of attention layers may lead to numerical instability or vanishing gradients. We extend the above-mentioned attention mechanism to fuse the dense predictions from multi-scale inputs more effectively.
  • Our solution resembles a volume rendering scheme applied to the scale space. This scheme provides a hierarchical, coarse-to-fine strategy to combine features, leveraging a scale-specific attention mechanism. We will also show that our approach generalizes the hierarchical multi-scale attention method [Hierarchical multi-scale attention for semantic segmentation].

Attentive Feature Aggregation (AFA)

这里设计了两种整合形式, 一种适用于双输入, 另一种适合用于多输入递进式集成. 核心都是基于空间注意力和通道注意力. 注意, 这里的计算都是两两集成的形式, 所以都是算出一个注意力后, 使用 sigmoid 来构造相对权重.

对于双输入形式, 空间注意力由较浅层特征计算, 因为其包含着丰富的空间信息. 而通道注意力由较深层特征计算, 因为其包含着更复杂的通道特征. 对于多输入形式(图中仅仅展示了三层, 实际上可以引入更多层的输入), 通道和空间注意力完全由当前层输入计算, 并且如果有靠前计算的一级的化, 该注意力会用来对当前和之前的输出作加权. 另外集成的顺序原文中如此描述"a feature with higher priority will have gone through a higher number of aggregations", 我的理解是, 应该就是从深到浅层的一个过程.

提出的集成模块可以用于许多结构中, 例如 DLA、UNet、HRNet 和 FCN 中.

Scale-Space Rendering (SSR)

这里提出的 SSR 是一个更加类似于模型集成的策略.

其通过计算针对不同尺度下预测的输出的相对权重来对多尺度推理进行集成. 所以, 这里涉及到两个问题:

1. SSR 如何学习？论文中并没有提到。但是按照上图中的说法，训练使用两个尺度的输入，说明这是可以训练 SSR 的。由于是个会预测参数的可学习的结构，对于每个输入会自动预测一个注意力参数。通过这些不同尺度输入下对应计算得到的参数从而获得最终针对多个尺度的加权比重。
2. 不同大小的预测最后会整合到哪个尺度？这一点论文中没有提。但是按照上图使用基于原始输入的相对尺寸的表述来看，最终应该还是会集成到 1.0 倍原始输入尺度上（与 hierarchical multi-scale attention 中的设计形式应该是一致的）。

表达形式

为了表达对多尺度预测的集成的问题, 作者首先将关注的重点放在单个像素上. 并且假设模型为目标像素在 \(k\) 个不同的尺度上提供了预测.
对于第 \(i\) 个尺度的预测可以表示为 \(P_i \in \mathbb{R}^{d}\). 由此, 在尺度空间中针对目标像素的特征表征可以定义为 \(P \triangleq (P_1, \dots, P_k)\). 进一步, 这里假设 \(i<j\) 表示尺度 \(i\) 比尺度 \(j\) 更加粗糙.

于是目标像素就可以想象成在尺度空间中移动的光线, 从尺度\(1\)朝向尺度\(k\).

基于这样的想法, 重新设计在提出的多特征融合机制中的原始的分层注意力, 并且模拟 volume-rendering equation, 这里的 volume 由尺度空间隐式给出.

为此, 除了位于尺度 \(i\) 的特征表征 \(P_i\), 假设 模型还会针对目标像素预测一个标量 \(y_i \in \mathbb{R}\). 在 volume rendering 的语境下, 粒子将 会穿过尺度\(i\)的概率, 在给定一些非负标量函数 \(\phi: \mathbb{R} \rightarrow \mathbb{R}_{+}\) 时, 就可以表示为 \(e^{-\phi(y_i)}\).

于是可以将尺度注意力 \(\alpha_i\) 表达为粒子到达尺度 \(i\) 并停留在这里的概率(每一次都满足伯努利分布, 非留即走, 前面都走, 就当前次留了下来):

\(\alpha_i(y) \triangleq [1 - e^{-\phi(y_i)}] \prod^{i-1}_{j=1}e^{-\phi(y_j)}, \, y \triangleq (y_1, \dots, y_k)\)

\(y\) 表示针对各个尺度的目标像素预测的标量参数.

\(P_{final} \triangleq \sum^{k}_{i=1}P_i \alpha_i(y)\)

最终, 按照 volume rendering equation, 针对目标像素多尺度预测融合得到的最终预测, 由不同尺度的注意力参数来加权求和获得. 这也反映了对于目标像素获得的最终特征, 是在 \(y\) 驱动下融合所有尺度的特征表达获得的.

综合上下文的分析, 这里的设计中应该是最终是要将所有尺度集成到 1 上的.

提出的 SSR 可以看作是 Hierarchical Multi-Scale Attention (HMA) [Hierarchical multi-scale attention for semantic segmentation, https://github.com/NVIDIA/semantic-segmentation]的一种一般化形式.

通过设置 \(\phi(y_i) \triangleq \log(1 + e^{y_i})\), 并固定 \(\phi(y_k) \triangleq \infty\), 就可以获得后者的形式了. 此时有:

\[\alpha_i = [1-\frac{1}{1+e^{y_i}}] \prod^{i-1}_{j=1}\frac{1}{1+e^{y_j}}, \\ \alpha_1=1-\frac{1}{1+e^{y_1}}, \\ \alpha_k=\prod^{k-1}_{j=1}\frac{1}{1+e^{y_j}}. \]

从这里的形式来看, 这里有两处令人疑惑的地方:

1. 形式不太对. 原本的 hierarchical multi-scale attention 使用的是 sigmoid 来集成不同的尺度. 这里与 sigmoid 并不一致.
2. 按照这里的形式, 并结合空间注意力(sigmoid)的级联关系, 可以看出输出是在 \(i=1\) 的位置, 也就是其它层的信息按照层序号递减的形式逐步集成起来. 这倒是和下图的形式大致类似.

输入是被再次放缩后才送入模型的. 而这里最终输出的尺寸是对应于 1.0 倍原始输入尺寸的. 所以, 假设按照尺度编号从 k 到 1 集成特征, 并在 1 层输出结果.

由于本文中构造的注意力是基于不选择当前层的概率(穿过当前层), 则对应上图的形式, 总的形式为:

\[\alpha_i = [1-p(y_i)]\prod_{j=1}^{i-1} p(y_j), \\ \alpha_1 = 1-p(y_1), \\ \alpha_k = \prod_{j=1}^{k-1} p(y_j), \\ p(y_i) = 1-\text{sigmoid}(y_i), \\ \Rightarrow P = \sum^{k}_{i=1} P_{i}\alpha_i(y). \]

可以看到, 对于第一层的注意力权重就是直接 sigmoid 的输出结果. 而对于第 k 层的输出, 则是对各层 sigmoid 输出取补并类乘而获得.

\(\phi\) 的选择

实验中使用绝对值函数: \(\phi(y_i) \triangleq |y_i|\). 这受启发于更好的保留经过注意力机制的梯度流的分析, 因为作者们发现现存的注意力机制可能会遭受梯度消失的问题.

前面整理的注意力系数的形式:

\[\alpha_i(y) \triangleq [1 - e^{-\phi(y_i)}] \prod^{i-1}_{j=1}e^{-\phi(y_j)} = \prod^{i-1}_{j=1}e^{-\phi(y_j)} - \prod^{i}_{j=1}e^{-\phi(y_j)}, \, y \triangleq (y_1, \dots, y_k) \]

考虑第 \(i\) 层系数 \(\alpha_i(y)\) 关于可学习参数 \(y_l\) 的导数:

\[J_{il} \triangleq \frac{\partial \alpha_i(y))}{\partial y_l} \begin{cases} \frac{\partial [-e^{-\phi(y_i)}]}{\partial y_l}\prod^{i-1}_{j=1}e^{-\phi(y_j)} = \frac{\partial \phi(y_i)}{\partial y_l}\prod^{i}_{j=1}e^{-\phi(y_j)} = \phi '(y_l)\prod^{i}_{j=1}e^{-\phi(y_j)} & \text{ if } l= i\\ 0 & \text{ if } l> i \\ -\phi '(y_l)\prod^{i-1}_{j=1}e^{-\phi(y_j)} + \phi '(y_l)\prod^{i}_{j=1}e^{-\phi(y_j)} = -\phi '(y_l)\alpha_i(y) & \text{ if } l< i \end{cases} \]

当考虑两个尺度的时候, 即 \(k=2\) 时:

\[J = \begin{bmatrix} -\phi '(y_1)a_1 & 0 \\ -\phi '(y_1)a_1(1-a_2) & \phi '(y_2)a_1a_2 \end{bmatrix}, \\ a_i \triangleq e^{-\phi(y_i)}. \]

左上角计算的是第 1 层的注意力系数关于第 1 层的参数的导数, 右上角则是第 1 层关于第 2 层的导数. 可以看到, 如果 \(a_1 \rightarrow 0\) 的时候, 梯度会消失, 不管 \(a_2\) 是多少.

所以为了避免梯度消失的问题, 这里需要对 \(\phi\) 进行仔细的设定. 当选择绝对值函数的时候, 这里的 Jacobian 矩阵将不会在 \(a_1 > 0\) 且 \((y_1, y_2) \neq (0, 0)\) 的情况出现消失的问题.

但是这里如果取了绝对值函数, 求导数是+-1, 这依然会有梯度消失的问题诶?

考虑 HMA 中的情况, 按照作者给出的形式, 此时有:

\[\phi '(y_i) = \frac{\partial \log(1+e^{y_i})}{\partial y_i} = \frac{e^{y_i}}{1+e^{y_i}} = 1 - \frac{1}{1+e^{y_i}} = 1 - e^{-\log(1+e^{y_i})} = 1 - a_i, \\ a_2 = 0. \]

分支 2 不参与注意力计算. 当 \(a_1 \rightarrow 1\) 时会出现梯度消失.

而按照我前面的形式, 则有:

\[\phi '(y_i) = \frac{\partial \log(1+e^{-y_i})}{\partial y_i} = -\frac{e^{-y_i}}{1+e^{-y_i}}, \\ a_i = e^{-\log(1+e^{-y_i})} = \frac{1}{1+e^{-y_i}}, \\ \phi '(y_i) = a_i - 1. \]

同样也会出现消失的问题.

链接

• 论文:https://arxiv.org/abs/2111.00770
• 代码:http://vis.xyz/pub/dla-afa
• 本文的思想来源于 NeRF, 可以看看 NeRF 的介绍再来看着 SSR 的设计.
• 关于体绘制的一些资料:
  • 非常丰富全面的一份中文 CG 学习材料: GPU 编程与 CG 语言之阳春白雪下里巴人(GPU Programming And Cg Language Primer)
  • 21 年出现在知网上的一份小综述: 基于神经辐射场的视点合成算法综述