안녕하세요. 지난 포스팅의 [Transformer] Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions (ICCV2021)에서는 Feature Pyramid를 활용한 PVT에 대해서 알아보았습니다. 오늘도 Transformer에 Feature Pyramid를 발전시킨 모델 중 하나인 Pyramid Pooling Transformer (P2T)에 대해서 알아보도록 하겠습니다.
P2T: Pyramid Pooling Transformer for Scene Understanding
Recently, the vision transformer has achieved great success by pushing the state-of-the-art of various vision tasks. One of the most challenging problems in the vision transformer is that the large sequence length of image tokens leads to high computationa
arxiv.org
Background
CNN의 발전은 심층신경망 모델을 다양한 영상 분류, 의미론적 분할, 객체 탐지, 객체 분할 등과 같은 다양한 Computer Vision task에 활용되는 계기가 되었습니다. 대표적인 모델로는 VGG, ResNet, Res2Net과 같은 모델들이 있었죠. 그러나, 자연어 처리 (NLP) 분야에서 Transformer를 활용하여 전체 문맥 (Global Context)를 추출하는 과정을 Computer Vision에 적용하려는 시도가 증가하게 되었습니다. 대표적으로 ViT가 있었죠.
하지만, ViT의 가장 큰 문제는 Transformer 구조를 Computer Vision에 적용할 때 sequence의 길이가 NLP에 비해 너무 길어져 계산 복잡도가 quadratic 만큼 커진다는 문제점입니다. 그 예시로 NLP 분야에서 자주 사용되는 데이터셋은 WMT2014 English-German dataset은 평균적인 sequence의 길이가 25인데 반해 $224 \times 224$를 $4 \times 4$ 크기의 패치로 분할하게 되면 총 3136개의 sequence가 나오게 되죠. 따라서, ViT에서는 이러한 문제를 해결하기 위해 패치의 크기를 $16 \times 16$이나 $32 \times 32$까지 늘려서 사용할 수 밖에 없었습니다.
하지만, 작은 패치 크기를 사용한다는 것은 더 많은 sequence를 얻음으로써 영상 내의 global context를 픽셀 단위로 얻을 수 있다는 것을 의미하기 때문에 더 작은 패치 크기를 사용하기 위해 연산량을 줄이기 위한 다양한 모델이 등장하였습니다. 가장 대표적으로 Swin Transformer와 PVT, 그리고 MViT가 있었습니다. (MViT 역시 향후에 블로그 포스팅에 추가하도록 하겠습니다.)
본 논문은 PVT와 MViT에서 사용하는 Single Pooling Operation에 집중합니다. 만약, Single Pooling Operation이 아닌 CNN 구조에서 자주 활용되는 Pyramid Pooling Operation을 Transformer에 적용할 수 있다면 성능이 향상될 것으로 저자들은 동기로 삼았습니다. 심지어 당시에는 다른 모델들이 Pyramid Pooling을 Transformer에 적용하지 않았다고 하죠. 따라서, 기본적으로 여러 pooling operation을 적용하여 MHSA을 통해 어텐션 시 추가적인 multi-scale 정보를 고려한다면 Transformer가 global context와 함께 더욱 강력한 feature representation을 얻을 수 있을 것으로 예상됩니다. 이러한 목적에서 Pyramid Pooling Transformer (P2T)가 제안되었습니다.
본 논문의 핵심 기여는 다음과 같이 정리할 수 있겠네요.
- MHSA에 pyramid pooling 구조를 처음으로 적용하여 sequence legnth를 크게 줄이고 강력한 feature representation을 얻을 수 있음
- Pyramid Pooling 기반의 MHSA (P-MHSA)를 통해 얻은 P2T는 시각적 인지 능력이 CNN뿐만 아니라 Transformer 기반 모델들인 PVT와 MViT보다도 뛰어남
- 다양한 영상 분류, 의미론적 분할, 객체 탐지, 객체 분할에서 기존 모델들보다 적은 파라미터로 높은 성능을 달성함 (그림 1은 의미론적 분할 성능 vs 파라미터 개수에 따른 성능 비교 의미).
Methology
1). Overview
그림 2는 P2T의 전체적인 구조를 블록 다이어그램으로 시각화하여 보여주고 있습니다. 전체적으로 보시면 흔히 저희가 보아왔던 Transformer 구조와 유사해보입니다. 가장 큰 차이점은 Patch Embedding 후 Pyramid Pooling Transformer라는 블록입니다. 저 부분이 P2T의 핵심이죠. 일단 전체적인 Forwarding 과정은 PVT와 거의 동일하기 때문에 다른 점만 보도록 하겠습니다.
일단, Positional Encoding 과정이 생략되었습니다. 이는 나중에 논문 리뷰로 추가할 예정이지만 PVT의 후속 논문인 PVT V2라는 모델에서 Positional Encoding을 하지 않고 Depth-wise Convolution (MobileNet 참고)를 통해 relative positional encoding을 수행하는 방식을 제안해었습니다. P2T는 이러한 방식을 따라 positional encoding을 적용하지 않습니다.
다음으로 각 스테이지 사이에 $2 \times 2$의 패치 그룹을 하나로 병합하여 $4 \times C_{i}$에서 $C_{i + 1}$로 linear projection을 적용합니다.
이 과정을 통해 각 스테이지에서 특징 맵 $\{ B_{1}, B_{2}, B_{3}, B_{4} \}$를 얻을 수 있습니다. 이때, Classification task를 수행하기 위해 $B_{4}$만 사용하게 되고 다른 Downstream scen understanding task (의미론적 영상 분할, 객체 탐지, 객체 영상 분할)을 수행할 때는 4개의 특징 맵 $\{ B_{1}, B_{2}, B_{3}, B_{4} \}$을 모두 사용하게 됩니다.
2). Pyramid Pooling Transformer
P2T의 핵심 기여 중 하나는 기존에 CNN에서만 사용되어 왔던 Pyramid Pooling operation을 Transformer에 성공적으로 적용하였다는 점 입니다. 이를 통해 연산되는 P2T는 다음과 같은 연산을 하게 됩니다.
$$\begin{cases} \mathbf{X}_{att} &= \text{LayerNorm} (\mathbf{X} + \text{P-MHSA} (\mathbf{X})) \\ \mathbf{X}_{out} &= \text{LayerNorm} (\mathbf{X}_{att} + \text{FFN} (\mathbf{X}_{att})) \end{cases}$$
여기서 P2T의 핵심은 $\text{P-MHSA}$와 $\text{FFN}$입니다. 이제 두 모듈의 동작 과정에 대해서 좀 더 자세히 분석해보도록 하겠습니다.
2-1). Pooling-based MHSA
그림 3에서는 본 논문에서 제안하는 P-MHSA의 모습입니다. 그림에서 왼쪽이 처음에 설명한 수식을 표현한 것이고 오른쪽이 바로 P-MHSA입니다. 각 단계를 좀 더 자세히 보도록 하겠습니다.
STEP1. Pyramid feature map을 얻기 위해 $\mathbf{X}$를 얻기 위해 다양한 비율의 pooling ratio를 적용
$$\begin{cases} \mathbf{P}_{1} &= \text{AvgPool}_{1} (\mathbf{X}) \\ \mathbf{P}_{2} &= \text{AvgPool}_{2} (\mathbf{X}) \\ &\vdots \\ \mathbf{P}_{n} &= \text{AvgPool}_{n} (\mathbf{X}) \end{cases}$$
STEP2. Relative positional encoding을 적용하기 위해 각 pyramid feature map에 대해 Depthwise Convolution을 적용
$$\begin{align*} \mathbf{P}_{i}^{\text{enc}} = \text{DWConv} (\mathbf{P}_{i}) + \mathbf{P}_{i} \end{align*}$$
STEP3. 각각의 pyramid pooling feature $\mathbf{P}_{i}^{\text{enc}}$를 vector로 flatten하고 하나로 concatenate
$$\begin{align*} \mathbf{P} = \text{LayerNorm} ( \text{Concat} (\mathbf{P}_{1}^{\text{enc}}, \mathbf{P}_{2}^{\text{enc}}, \dots, \mathbf{P}_{n}^{\text{enc}}) ) \end{align*}$$
이 과정을 통해 얻은 $\mathbf{P}$는 pooling ratio가 충분히 크다면 입력 특징 맵 $\mathbf{X}$보다 더 짧은 sequence length를 가지는 얻을 수 있습니다. 또한, 다양한 크기의 pooling ratio를 적용하였기 때문에 multi-scale 정보가 임베딩되어 입력 특징 맵 $\mathbf{X}$의 contextual abstraction을 얻을 수 있기 때문에 MHSA를 계산할 때 높은 정보 상관성을 얻을 수 있습니다.
STEP4. Query $\mathbf{Q}$, Key $\mathbf{K}$, Value $\mathbf{V}$를 정의
$$\begin{align*} (\mathbf{Q}, \mathbf{K}, \mathbf{V}) = (\mathbf{X}\mathbf{W}^{q}, \mathbf{X}\mathbf{W}^{k}, \mathbf{X}\mathbf{W}^{v}) \end{align*}$$
STEP5. Query $\mathbf{Q}$, Key $\mathbf{K}$, Value $\mathbf{V}$를 이용하여 MHSA를 적용
$$\begin{align*} \mathbf{A} = \text{SoftMax} \left( \frac{\mathbf{Q} \times \mathbf{K}^{T}}{\sqrt{d_{K}}} \right) \mathbf{V} \end{align*}$$
이때, Query $\mathbf{Q}$와 Key $\mathbf{K}$는 모두 pyramid pooled feature $\mathbf{P}$로부터 얻기 때문에 입력 특징 맵 $\mathbf{X}$보다 더 짧은 sequence를 가지므로 P-MHSA는 기존 MHSA보다 연산이 적게 소요됩니다. 또한, multi-scale 정보가 임베딩된 $\mathbf{P}$로인해 Query $\mathbf{Q}$와 Key $\mathbf{K}$ 기반의 P-MHSA는 global context를 모델링하는 데 더 강력한 기능을 가지게 됩니다.
$N$이 입력 특징 맵 $\mathbf{X}$의 sequence 길이 그리고 $C$를 채널의 개수라고 하면 일반적으로 Pyramid Pooling Operation는 $O(NC)$의 복잡도를 가지게 됩니다. 따라서 최종적으로 P-MHSA의 계산복잡도는 다음과 같습니다.
$$\begin{align*} O(\text{P-MHSA}) = (N + 2M)C^{2} + 2NMC \end{align*}$$
여기서 $M$는 pyramid pooled feature $\mathbf{P}$의 sequence 길이를 의미합니다. 본 논문에서는 default pooling ratio로 $\{ 12, 16, 20, 24 \}$를 사용했기 때문에 입력 특징 맵의 길이보다 약 64배 더 짧게 됩니다 ($M \approx \frac{N}{66.3} \approx \frac{N}{8^{2}}$).
2-2). Feed-Forward Network
마지막으로 본 논문에서 사용하는 Feed-Forward는 1D CNN 기반의 FFN (Local ViT)와 Inverted Bottleneck Block (MobileNet V2)를 결합하여 2D Locality와 Efficient를 동시에 얻게 됩니다. 각 단계는 다음과 같이 진행됩니다.
STEP1. 입력 sequence $\mathbf{X}_{att}$를 2D 특징 맵 $\mathbf{X}_{att}^{I}$으로 변환
$$\begin{align*} \mathbf{X}_{att}^{I} = \text{Seq2Image} (\mathbf{X}_{att}) \end{align*}$$
STEP2. Inverted Residual Block을 적용
$$\begin{cases} \mathbf{X}_{\text{IRB}}^{1} &= \text{Act} (\mathbf{X}^{I}_{att} \mathbf{W}_{\text{IRB}}^{1}) \\ \mathbf{X}_{\text{IRB}}^{\text{out}} &= \text{Act} (\text{DWConv} (\mathbf{X}_{\text{IRB}}^{1})) \mathbf{W}_{\text{IRB}}^{2} \end{cases}$$
STEP3. 출력 feature map $\mathbf{X}^{\text{out}}_{\text{IRB}}$를 sequence $\mathbf{X}^{\text{S}}_{\text{IRB}}$로 변환
$$\begin{align*} \mathbf{X}_{IRB}^{\text{S}} = \text{Image2Seq} (\mathbf{X}_{\text{IRB}}^{\text{out}}) \end{align*}$$
3). Model Variants
ViT와 PVT와 마찬가지로 각 Stage의 반복 횟수를 기반으로 P2T의 4가지 변형 모델 P2T-Tiny, P2T-Small, P2T-Base, P2T-Large을 함께 제안합니다. 이때, 연산량을 주로 Stage 3에 집중하여 설계하였습니다. 각각의 모델은 ResNet-18, ResNet-50, ResNet-101, PVT-Large 모델과 거의 유사한 개수의 파라미터를 가지고 있습니다.
4). Implementation Details
- Pyramid Pooling Operations
- Number of Parallel Pooling Operation: 4
- Pooling ratios of P2T are different
$$\begin{align*} \{ 12, 16, 20, 24 \} \rightarrow \{ 6, 8, 10, 12 \} \rightarrow \{ 3, 4, 5, 6 \} \rightarrow \{ 1, 2, 3, 4 \} \end{align*}$$
- Other settings
- All DWConv in P-MHSA share the same parameters
- The kernel size of DWConv is set to $3 \times 3$ for efficiency
- Choose Hardswish as non-linear activation $\text{Act}$ since it is more memory-efficient than GELU
- Overlapped Patch Embedding (Same as PVT V2)
- $7 \times 7$ Convolution with a stride 4 for patch embedding in the first stage
- $3 \times 3$ Convolution with a stride 2 for patch embedding from the second to the last stage
Experiment Results
1). ImageNet Classification
- Dataset
- ImageNet-1K: 1.28 million training images & 50K validation images with 1,000 classes
- Data Augmentation (Same as DeiT & PVT): random cropping, random horizontal flipping, label-smooting regularization, MixUp, CutMix, Random Erasing
- Optimization: AdamW
- momentum: -
- weight decay: 0.05
- learing rate: 0.001
- batch size: 1,024
- epochs: 300
- GPU Type 언급 X
- 표 2는 영상 분류에서 실험 결과를 요약하고 있습니다. 전체적으로 PVT 계열보다 성능이 소폭 향상 (0.1% ~ 1.1%) 향상된 것을 볼 수 있습니다.
- P2T와 CNN 기반 모델과 비교해보았을 때 Tiny에서 ResNet18과 거의 동일한 파라미터를 가지지만 성능으 약 11% 향상되었습니다. 또한, 더 대규모 모델인 Small, Base, Large에서는 ResNet-50, ResNet-101, ResNeXt-101보다 파라미터 개수가 현저히 적으면서 성능도 크게 향상되 었습니다. 이러한 현상은 다른 CNN 기반 모델인 Res2Net에서도 관찰됩니다.
- P2T와 Transformer 기반 모델들과 비교해보았을 때 Tiny에서 DeiT나 ViT 보다는 파라미터 개수가 소폭 늘어난 것을 볼 수 있습니다. 하지만 성능은 각각 3% 그리고 7% 향상된 것을 볼 수 있죠. PVT V2와 비교해보면 모델이 커질수록 차이가 점점 줄어드는 것을 볼 수 있습니다. 하지만, Large 모델을 기준으로 비교해보았을 때 P2T와 PVT V2는 거의 동일한 성능을 보이지만 파라미터 개수로는 훨씬 적은 것을 볼 수 있죠.
2). Semantic Segmentation
- Dataset: ADE20K
- 20,120 training images & 2,000 validation images & 3,352 testing images with 150 fine-grained semantic categories
- Semantic Segmentation model: Semantic FPN (Same as PVT & Twins)
- Randomly Resize & Cropping the image to $512 \times 512$
- Optimization: AdamW
- momentum: -
- weight decay: $1 \times 10^{-4}$
- learing rate: $1 \times 10^{-4}$ with polynomial decay schedule with a power of 0.9
- batch size: 16
- epochs: - (80k iteration)
- MMSegmentation Toolbox
- GPU Type 언급 X
- 표3는 Semantic FPN을 모델로 했을 때 전체적인 성능을 요약해서 보여주고 있습니다. 결과적으로 CNN 계열을 backbone으로 사용했을 때 보다 더 높은 성능을 달성하였습니다. 이 역시 영상 분류와 거의 유사하게 모델의 규모가 커질수록 PVT V2와 성능 차이는 나지 않지만 파라미터 개수는 훨씬 적은 것을 볼 수 있습니다.
3). Object Detection
- Dataset: COCO train2017 (train dataset) & COCO val2017 (evaluation)
- 118k training images & 5k validation images
- Object Detection model: RetinaNet
- Optimizer: AdamW
- momentum: -
- weight decay: $1 \times 10^{-4}$
- learing rate: $1 \times 10^{-4}$ and divide by 10 after 8 and 11 epochs
- batch size: 16
- epochs: 12
- MMDetection Toolbox
- GPU Type 언급 X
4). Instance Detection
- Dataset: COCO train2017 (train dataset) & COCO val2017 (evaluation)
- 118k training images & 5k validation images
- Object Detection model: Mask RCNN
- Optimizer: AdamW
- momentum: -
- weight decay: $1 \times 10^{-4}$
- learing rate: $1 \times 10^{-4}$ and divide by 10 after 8 and 11 epochs
- batch size: 16
- epochs: 12
- GPU Type 언급 X
Ablation Studies
1). Multiple Pyramid Pooling Ratios
- 표 5는 pooling ratio의 다양성에 따른 성능을 보여주고 있습니다.
- 단일 pooling ratio로 했을 때 12보다 큰 경우에는 성능이 크게 떨어지는 것을 볼 수 있습니다. 하지만 12보다 작게 설정하면 성능이 더 이상 향상되지는 않죠.
- 다중 pooling ratio로 하게 되면 포화되던 성능이 mIoU에서 크게 향상되는 것을 볼 수 있습니다.
2). Significance of Pyramid Pooling for Different Stages
- 표 6은 P2T의 P-MHSA를 Stage 단위로 추가했을 때 성능을 보여주고 있습니다.
3). Pooling Operation Choices
- 표 7은 Pooling 연산을 Average Pooling, Max Pooling으로 했을 때와 Convolution으로 바꾸었을 때 성능을 보여주고 있습니다. Max Pooling을 하게 되면 성능이 크게 하락하네요.
4). Fixed Pooled Sizes
Code Analysis
PyTorch 기반의 P2T 전체 코드는 공식 깃허브 사이에서 확인할 수 있습니다.
GitHub - yuhuan-wu/P2T: [TPAMI22] Pyramid Pooling Transformer for Scene Understanding
[TPAMI22] Pyramid Pooling Transformer for Scene Understanding - GitHub - yuhuan-wu/P2T: [TPAMI22] Pyramid Pooling Transformer for Scene Understanding
github.com
본 논문의 핵심인 Pyramid Pooling Attention만 확인해보도록 하겠습니다.
class PoolingAttention(nn.Module):
"""
Pyramid Pooling Attention
Args:
dim (int): input dimension
num_heads (int): number of attention heads
qkv_bias (bool): enable bias for qkv if True
qk_scale (Optional[float]): override default qk scale of head_dim ** -0.5 if set
attn_drop_rate (float): attention dropout rate
proj_drop_rate (float): projection dropout rate
pool_ratios (List[int]): pyramid pooling ratios
"""
def __init__(self,
dim: int,
num_heads: int=2,
qkv_bias: bool=False,
qk_scale: Optional[float]=None,
attn_drop_rate: float=0.,
proj_drop_rate: float=0.,
pool_ratios: List[int]=[1, 2, 3, 6]) -> None:
super(PoolingAttention, self).__init__()
assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."
self.dim = dim
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = qk_scale or head_dim ** -0.5
self.q = nn.Linear(dim, dim, bias=qkv_bias)
self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop_rate)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop_rate)
self.pool_ratios = pool_ratios
self.norm = nn.LayerNorm(dim)
def forward(self, x, H, W, d_convs=None):
B, N, C = x.shape
q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
pools = []
x_ = x.permute(0, 2, 1).reshape(B, C, H, W)
for (pool_ratio, l) in zip(self.pool_ratios, d_convs):
pool = F.adaptive_avg_pool2d(x_, output_size=(round(H / pool_ratio), round(W / pool_ratio)))
pool = pool + l(pool)
pools.append(pool.view(B, C, -1))
pools = torch.cat(pools, dim=-1)
pools = self.norm(pools.permute(0, 2, 1))
kv = self.kv(pools).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
k, v = kv[0], kv[1]
attn = (q @ k.transpose(-2, -1)) * self.scale
attn = attn.softmax(dim=-1)
attn = self.attn_drop(attn)
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
x = self.proj(x)
x = self.proj_drop(x)
return x
PVT와 다르게 sr_ratio가 아닌 pool_ratios를 입력받는 것을 볼 수 있습니다. forwarding 시에는 입력 특징맵에 대해 pooling을 다양한 크기로 적용하여 relative positional embedding도 함께 적용하게 되죠. 이후 과정은 모두 동일하게 진행됩니다.
