这幅图展示了Axial Transformer模型在处理二维张量(如图像)时的采样过程。Axial Transformer是一种基于自注意力机制的自回归模型,专门设计用于处理高维张量数据,如图像和视频。图中描述了模型在不同阶段的操作,具体如下:
编码阶段(编码所有先前的通道和帧):
在采样一个新的通道之前,模型首先使用8个块(blocks)的无掩码(unmasked)行注意力和列注意力来编码所有先前的通道和帧。这一步骤在图的左侧展示,其中G代表生成(Generate)操作,用于处理先前的通道和帧信息。
整合先前采样的行(针对每一行):
对于图像的每一行,模型应用4个块的无掩码行注意力和掩码(masked)列注意力,以整合先前为活动通道采样的行到编码表示中。这一步骤在图的中间部分展示,其中B代表基础(Base)操作,用于整合先前采样的行。
采样新行:
最后,模型将编码表示向上移动,以确保条件信息满足因果关系(causality),并运行由4个块组成的内部解码器(inner decoder),这些块由掩码行注意力组成,用于采样图像中的新行。这一步骤在图的右侧展示。
图中的流程图和对应的操作说明展示了Axial Transformer在采样过程中的三个主要阶段:
每个通道运行一次:在采样新通道之前,对所有先前的通道和帧进行编码。
每行运行一次:对于每一行,整合先前采样的行。
每个位置运行一次:对每个位置进行采样。
这种结构允许Axial Transformer在解码过程中大部分上下文可以并行计算,而不引入任何独立性假设。这种半并行结构使得即使非常大的Axial Transformer模型也能广泛适用。图中的网格图示进一步说明了在不同阶段,模型是如何通过注意力机制处理和整合信息的。
无掩码行注意力和掩码列注意力有什么区别?
无掩码行注意力(unmasked row attention)和掩码列注意力(masked column attention)是Axial Transformer中使用的两种不同的注意力机制,它们在处理数据时有不同的侧重点和应用场景:
无掩码行注意力(Unmasked Row Attention):
这种注意力机制不使用掩码,意味着在处理行时,每个元素可以关注到该行中的所有其他元素。
它允许模型在处理行内数据时,考虑到行内所有先前的信息,从而捕捉行内的长期依赖关系。
在Axial Transformer中,无掩码行注意力用于编码先前的通道和帧,以及整合先前采样的行到编码表示中。
掩码列注意力(Masked Column Attention):
这种注意力机制使用掩码,意味着在处理列时,每个元素只能关注到该列中当前位置之前(或之后,取决于掩码的方向)的元素。
掩码用于确保模型在生成序列时遵循因果关系,即每个元素只能依赖于它之前的时间步(或空间位置),而不能“看到”未来的信息。
在Axial Transformer中,掩码列注意力用于确保在生成新行时,每个元素只能依赖于它左边和上面的元素,从而保持自回归模型的正确性。
在Axial Transformer的上下文中,这两种注意力机制的组合使用允许模型在保持计算效率的同时,有效地捕捉和利用多维数据中的长期依赖关系。无掩码行注意力用于整合行内信息,而掩码列注意力则确保了在生成过程中遵循因果关系,从而使得模型能够生成符合数据分布的高质量样本。
这幅图(Figure 2)展示了Axial Transformer中使用的不同类型的轴向注意力(axial attention)层,这些层是构建Axial Transformer的基础。图中展示了四种类型的注意力层,每种层都对应于输出红点位置的感受野(receptive field),即影响输出位置的输入位置集合。蓝色位置表示这些感受野。
Full Row(完整行注意力):
在这种注意力层中,输出位置(红点)可以接收到同一行中所有其他位置(蓝点)的信息。
这意味着在处理时,该行的所有元素都可以直接相互影响。
Masked Row(掩码行注意力):
与完整行注意力不同,掩码行注意力使用掩码来限制信息流。
输出位置只能接收到它左侧的所有位置(蓝点)的信息,而不能接收到右侧的位置信息。
这种掩码确保了在自回归模型中,当前位置的生成不会依赖于未来位置的信息,从而保持了因果关系。
Full Column(完整列注意力):
在这种注意力层中,输出位置可以接收到同一列中所有其他位置的信息。
这允许模型在处理时考虑列内的所有元素。
Masked Column(掩码列注意力):
掩码列注意力也使用掩码来限制信息流。
输出位置只能接收到它上方的所有位置(蓝点)的信息,而不能接收到下方的位置信息。
这同样是为了保持自回归模型中的因果关系,确保当前位置的生成不会依赖于未来位置的信息。
这些不同类型的轴向注意力层使得Axial Transformer能够有效地处理高维数据,如图像和视频,同时保持计算效率。通过在不同维度上应用不同类型的注意力,模型能够在保持对数据的全面理解的同时,有效地捕捉和利用数据中的长期依赖关系。
这张表格(Table 1)比较了最近提出的多维Transformer架构之间的权衡。表格中列出了五种不同的模型,并根据四个特性对它们进行了评估:
Full receptive field(全感受野)
:
表示模型是否能够捕获数据的所有可能的上下文信息。
“yes”意味着模型可以访问到完整的数据上下文,而“no”则表示模型的上下文受到限制。
Attention faster than O(N^2)(注意力计算速度快于O(N^2))
:
指的是模型在计算注意力时是否比传统的O(N^2)复杂度更快。
“yes”表示模型在处理长序列时更加高效,而“no”则表示模型可能在处理大数据时面临计算上的挑战。
Needs no custom kernels(不需要自定义内核)
:
表示模型是否需要特定的硬件加速或自定义的内核来实现。
“yes”意味着模型可以在标准的深度学习框架中实现,而“no”则可能需要额外的硬件支持或优化。
Semi-parallel context aggregation(半并行上下文聚合):
指的是模型是否能够在解码时并行处理部分上下文信息。
“yes”表示模型可以在不牺牲表达能力的情况下加速解码过程,而“no”则意味着模型可能需要顺序处理所有信息。
具体分析如下:
Transformer (Vaswani et al., 2017):
具有全感受野,但注意力计算复杂度为O(N^2),需要自定义内核,不支持半并行上下文聚合。
Image Transformer (Parmar et al., 2018):
不具有全感受野,但注意力计算速度快于O(N^2),不需要自定义内核,不支持半并行上下文聚合。
Block Transformer (Weissenborn et al., 2019):
不具有全感受野,注意力计算速度快于O(N^2),不需要自定义内核,不支持半并行上下文聚合。
Strided Sparse Transformer (Child et al., 2019):
具有全感受野,注意力计算速度快于O(N^2),但需要自定义内核,不支持半并行上下文聚合。
Axial Transformer (ours):
具有全感受野,注意力计算速度快于O(N^2),不需要自定义内核,支持半并行上下文聚合。
从表格中可以看出,Axial Transformer在所有四个特性上都表现优异,它不仅能够捕获完整的上下文信息,而且计算效率高,易于实现,并且支持半并行处理,这使得它在处理高维数据时具有显著的优势。
这幅图(Figure 3)描述了Axial Transformer编码网络的输入排列方式。图中展示了如何将图像或视频的先前可用或已生成的通道顺序堆叠在输入中。具体来说:
Channel index plane(通道索引平面):最上面的黑色层表示通道索引平面,用于指示当前正在处理的通道索引。
Variable number of padding planes(可变数量的填充平面):在通道索引平面下方,有多个灰色层,这些是填充平面,用作未来生成通道的占位符。填充平面的数量是可变的,可以根据需要生成的通道数量进行调整。
RGB channels for frame T-3, T-2, T-1, current frame T(帧T-3、T-2、T-1和当前帧T的RGB通道):在填充平面下方,依次是不同帧的RGB通道。从下到上,分别是当前帧T之前的三个帧(T-3、T-2、T-1)的RGB通道,然后是当前帧T的先前RGB通道。
这种输入排列方式允许Axial Transformer在生成新的通道时,能够考虑到先前生成的通道信息。通过这种方式,模型可以在生成每一帧或每个通道时,利用之前的数据来提高生成的质量和一致性。
图中的说明进一步解释了这种输入排列的目的:
先前可用或生成的图像或视频通道被顺序堆叠在输入中。
使用可变数量的填充平面作为未来生成通道的占位符。
最终的整数平面向Axial Transformer发出信号,指示当前步骤正在生成的通道。
这种设计使得Axial Transformer能够在处理多维数据(如图像和视频)时,有效地利用先前生成的信息,从而提高生成结果的质量和一致性。