Deformable Attention

変形可能アテンションは、各クエリが特徴マップのすべての位置ではなく、参照位置の周囲にある一定数のサンプリング点だけに注目する、コンピュータービジョン向けのアテンション機構です。クエリ自身が、参照位置からのオフセットとして各点の位置を予測し、それぞれの重みも決めます。そのため、計算コストは画像解像度ではなく、クエリ数とサンプリング点数に応じて増加します。SenseTime、中国科学技術大学、香港中文大学のXizhou Zhu、Jifeng DaiらがDeformable DETRで提案し、ICLR 2021で口頭発表されました。この手法は、元のDETR検出Transformerの2つの弱点を解消しました。完全なアテンションでは低解像度の特徴に制限され、トレーニングには500エポックかかっていました。Deformable DETRでは必要なエポック数が10分の1になり、小さい物体の検出精度も向上しました。

変形可能アテンションの仕組み#

この考え方は、Daiも共同開発した変形可能畳み込みを拡張したものです。変形可能畳み込みでは、畳み込みカーネルのタップ位置のオフセットを学習します。参照点を持つ各クエリについて、各アテンションヘッドはクエリ特徴に2つの線形層を適用します。一方はK個のサンプリングオフセットを予測し、もう一方はsoftmaxで正規化したK個のアテンション重みを予測します。オフセットは小数になり得るため、各ずらした位置の値特徴はバイリニア補間で読み取られ、ヘッドはそれらの加重和を返します。論文では8個のヘッドとK = 4個の点を使うため、クエリは各特徴マップから32か所を読み取ります。標準的なセルフアテンションでは、すべてのピクセルをスコアリングします。単一ヘッド、単一スケールのバージョンは、PyTorchで数行のコードとして実装できます。

import torch
import torch.nn.functional as F

torch.manual_seed(0)
dim, n_points, n_queries = 32, 4, 100
value = torch.randn(1, dim, 64, 64)  # a 64x64 feature map: 4,096 positions
query = torch.randn(n_queries, dim)
ref = torch.rand(n_queries, 2) * 2 - 1  # reference points in grid_sample's [-1, 1] coordinates

# Each query predicts where to look (offsets) and how much each sample counts (weights)
to_offsets, to_weights = torch.nn.Linear(dim, n_points * 2), torch.nn.Linear(dim, n_points)
offsets = 0.1 * to_offsets(query).view(n_queries, n_points, 2)
weights = to_weights(query).softmax(dim=-1)

# Bilinearly sample the value map at reference + offset, then take the weighted sum
grid = (ref[:, None] + offsets)[None]  # (1, n_queries, n_points, 2)
sampled = F.grid_sample(value, grid, align_corners=False)[0]  # (dim, n_queries, n_points)
out = (sampled * weights).sum(dim=-1).T  # (n_queries, dim)
print(f"Output {tuple(out.shape)} from {n_queries * n_points} samples, not {n_queries * 64 * 64:,} dense scores")

Output (100, 32) from 400 samples, not 409,600 dense scoresが出力されます。100個のクエリはそれぞれ、4,096か所のうち4か所を読み取ります。トレーニングではオフセット層と重み層を調整し、タスクに役立つ画像領域へサンプリング点を移動させます。

マルチスケール変形可能アテンション#

物体のサイズは大きく異なるため、Deformable DETRは複数の特徴マップ上の点を一度にサンプリングします。L個のレベルがある場合、各ヘッドは各レベルでK個の点を取り、L × K個すべての点にわたって重みを正規化します。これにより、細かい特徴と粗い特徴を一度に混合でき、エンコーダーでは特徴ピラミッドネットワーク (FPN)に代わって使われました。エンコーダーでは各レベルのすべてのピクセルをクエリとするセルフアテンションとして適用し、デコーダーでは物体クエリから画像特徴へのクロスアテンションとして適用します。

Ultralyticsパッケージで利用できるBaiduのリアルタイム検出TransformerであるRT-DETRは、6つのデコーダーレイヤーのクロスアテンションにマルチスケール変形可能アテンションを使用します。300個の物体クエリはそれぞれ、8個のヘッドでストライド8、16、32の特徴マップ上からヘッドごとに4点ずつサンプリングし、合計96か所を参照します。640ピクセル画像のそれらのマップに対する完全なクロスアテンションでは、8,400か所をスコアリングする必要があります。

用途#

• 物体検出: DINOやRT-DETRを含むDeformable DETRの後継モデルは、デコーダー内で変形可能クロスアテンションを維持しているため、小さい物体を検出する際に高解像度の特徴を参照できます。
• 自動運転車: BEVFormerは複数の車載カメラの画像を鳥瞰視点のグリッドに変換し、各グリッドセルが自分を捉えているカメラ画像内の変形可能な点をサンプリングします。
• パノプティックセグメンテーション: Mask2Formerは、マスクを予測する前に変形可能アテンションでマルチスケールのピクセル特徴を構築します。

変形可能アテンションと関連用語の比較#

• スライディングウィンドウアテンション: 位置に基づいて決まる固定近傍に注目し、入力にかかわらず同じ範囲を使います。
• スパースアテンション: 事前に選んだトークンペアの固定サブセットをスコアリングします。
• Flash attention: オンチップGPUメモリーでタイル分割することで通常の密なアテンションを高速に計算し、標準的なアテンションと同じ結果を得ます。
• ビジョントランスフォーマー (ViT)のアテンション: 画像内のすべてのパッチを他のすべてのパッチと比較するため、パッチ数の2乗に比例して計算コストが増加します。