Lambda Networks 논문 리뷰

1. 개요

Google Research가 ICLR 2021에서 발표한 Lambda Networks에 대한 논문 리뷰입니다. 람다레이어는 입력과 다른 contextual한 정보에 대한 long-range 상호작용을 self-attention 없이 구현했습니다. 이 레이어는 linear attention과 attention map을 사용하지 않는다는 점에서는 비슷하지만 위치정보와 내용정보의 상호작용을 모델링할 수 있습니다. 결과를 보면 LambdaResNets은 EfficientNets보다 3.2 - 4.4x 빠르며 비슷한 정확도를 보였습니다.

 

2. 특징

데이터에서 long-range 정보를 모델링하는 것은 머신러닝에서 중요한 문제입니다. Self-attention은 이러한 문제를 해결했지만 메모리의 높은 사용량은 이미지같은 다차원 데이터에 적용하기 힘들게 만듭니다.

Linear attention은 높은 메모리 사용량을 해결했지만 픽셀간의 상대적인 거리등의 내부적인 데이터 구조를 모델링하지 못합니다.

이 논문에서는 쿼리와 구조화된 Context 요소들의 집합사이의 long-range 상호작용을 낮은 메모리비용으로 가능하게 합니다. 람다 레이어는 람다를 통하여 Context를 선형함수로 변환하게 됩니다. Self-attention은 메모리를 많이 차지하는 쿼리와 Context 사이의 similiarity kernel을 정의하지만 람다레이어는 Context 정보를 고정된 길이의 선형함수로 축약하게 됩니다.

결론적으로 람다레이어는 위치정보와 내용정보에 대한 상호작용 모두 모델링할 수 있으며 매우 적은 메모리로 이를 가능하게하여 고해상도 이미지를 입력으로 받을 수 있습니다.

• 컨볼루션과 어텐션에 비해 이미지넷 분류에서 효율적이며 더 나은 점수를 보여줍니다.
• ResNet-50에서 3x3 컨볼루션을 람다레이어로 바꾸는것만으로도 50퍼센트의 계산효율 상승과 +1.5%의 성능향상을 보입니다.
• LambdaResNets은 컨볼루션-람다의 혼합모델로서 EfficientNets보다 4.4x 빠르고 메모리 효율적입니다.

Modeling Long-Range Interaction

여기서는 쿼리, 컨텍스트, interaction을 정의합니다. 키는 쿼리와 컨텍스트 사이의 상호작용을 잡는 역할을 하게 됩니다.

Notation

스칼라는 소문자, 벡터는 bold 소문자, 텐서는 bold 대문자로 나타내며 |n|은 집합의 크기를 나타냅니다. 

Queries and Contexts

Q와 C는 각각 벡터들의 구조화된 집합으로 쿼리와 컨텍스트를 의미합니다. 각 쿼리 $(\mathbf{q}_n, n) \in \mathbb{R}^{|n|}$은 벡터 $\mathbf{q}_n$과 위치 $n$에 의해 특징화됩니다. 비슷하게 컨텍스트 요소인 $(\mathbf{c}_m, m)$도 벡터 $\mathbf{c}_m$과 위치 $m$에 의해 나타내집니다. $(n, m)$ 쌍은 픽셀간의 위치나 그래프에서 노드간의 엣지와같은 구조요소사이의 관계를 나타냅니다.

Defining Interactions

고려해야하는건 컨텍스트 C와 함수 $F :((\mathbf{q}_n, n), C) \to y_n$이 주어졌을때 쿼리 $(\mathbf{q}_n, n)$을 ouput 벡터 $y_n$으로 맵핑하는 것입니다. 함수 $F$는 신경망에서 입력을 처리할때 레이어처럼 작동하게 됩니다.

• $(\mathbf{q}_n, \mathbf{c}_m)$은 내용기반의 상호작용을 나타냅니다.
• $(\mathbf{q}_n, (n, m))$은 위치기반의 상호작용을 나타냅니다.

Keys to capture long-range interactions

논문에서는 Interactions을 얻기 위해 빠른 배치선형 연산을 우선시하면서 dot-product 연산을 사용했습니다. 따라서 query와 상호작용할 수 있는 벡터의 도입를 도입하였고 이 벡터의 크기는 내적연산을 위해 query와 같은 차원을 갖습니다. 특별히, 내용기반의 상호작용 $(\mathbf{q}_n, \mathbf{c}_m)$은 보통 key라고 알려져있는 $|k|$-차원의 벡터 $\mathbf{k}_m$이 필요합니다. 반대로, 위치기반 상호작용 $(\mathbf{q}_n, (n, m))$은 상대적인 위치 임베딩 $\mathbf{e}_{nm} \in \mathbb{r}^{|k|}$ 이 필요하게 됩니다.

Lambda Interactions

Similarity Kernal(Attention mpa)을 생성하는 Attention과 다르게 Lambda Layer는 심플하게 각 쿼리를 선형 함수인 $\lambda(C, n) : \mathbb{R}^{|k|} \to \mathbb{R}^{|v|}$ 를 이용해 맵핑하게 됩니다. 각 $\lambda_n$은 컨텍스트와 독립적으로 존재하는 작은 선형함수처럼 동작하며 쿼리에 적용된 후에는 버려집니다.

 

3. Lambda Layers

Computational graph of the lambda layer

람다레이어는 입력으로 $\mathbf{X} \in \mathbb{R}^{|n| \times d_{in}}$와 Context $\mathbf{C} \in \mathbb{R}^{|m| \times d_{c}}$를 받게되고 람다라는 선형함수를 만들고 쿼리에 적용해 출력 $\mathbf{Y} \in \mathbb{R}^{|n| \times d_{out}}$을 출력합니다.

$d_{in} = d_c = d_{out} = d.$이고, self-attention처럼 $\mathbf{C} = \mathbf{X}.$를 가진다고 가정합니다. 논문의 나머지 부분에서는 long-range 내용기반과 위치기반의 상호작용을 잡아내는 람다 레이어의 특정경우에 대해 초점을 맞춥니다.

Generating the contextual lambda function

$\mathbb{R}^{|k|} \to \mathbb{R}^{|v|}$로 만드는 선형함수가 필요합니다. 즉 행렬 $\mathbf{\lambda}_n \in \mathbb{R}^{|k| \times |v|}$입니다. 람다 레이어는 첫번째로 Context를 linear projection으로 keys ${\mathbf{K}}$와 values ${\mathbf{V}}$를 생성합니다. 여기서 keys는 소프트맥스를 거쳐 정규화되어 $\mathbf{\bar{K}}$를 생성합니다.

• $content$ $lambda$ $\mathbf{\lambda}^c$는 모든 쿼리에 걸쳐 공유되고 컨텍스트 요소의 순서에 불변합니다. 이것은 쿼리를 컨텍스트 내용에 따라서 변환하게 됩니다.
• $position$ $lambda$ $\mathbf{\lambda}^p_n$는 위치 임베딩 $\mathbf{E}_n$에 통한 쿼리 위치 $n$에 의존합니다. 이것은 쿼리를 컨텍스트 요소와 상대적인 위치 $(n, m)$에 따라서 변환하게 됩니다.

즉 position-based인 위치 임베딩과 V를 내적하고, content-based인 K와 V를 내적한 것을 각각 더하면 람다레이어를 생성하게 됩니다.

Applying lambda to its query

구해준 람다레이어와 쿼리를 곱하면 output $y_n$을 얻을 수 있습니다. 결론적으로 람다레이어는 attention map 없이 위치정보와 내용정보에 관한 특징들을 잘 잡아냅니다.

Multi-Query lambda layers

$|v|=d$를 도입하는 것 대신에 $|h|$개의 쿼리 $\{\mathbf{q}^h_n\}$를 만들어 같은 람다함수를 각 쿼리에 적용한 후 concatenate 시키게 됩니다. 이제 $|v|=d/|h|$를 가지며 $|h|$에 의해 complexity를 낮출 수 있습니다. 위 Pseudo-code는 멀티쿼리를 적용한 람다레이어입니다.

Lambda Convolution

Long-range 상호작용의 이점이 있지만 Locality는 많은 태스크에서 중요합니다. 글로벌한 Context를 이용하는 것은 노이지하거나 계산적으로 불필요할 수 있습니다. 따라서 요구된 범위($m$) 밖에 있는 상대적 임베딩을 0으로 만들어 Locality를 만들 수 있습니다. 하지만 높은 $|m|$에서는 높은 계산량을 요구하게 됩니다. 람다 컨볼루션은 Context가 다차원 그리드에 있을때 기존의 컨볼루션을 이용해 로컬 Context로 부터 positional lambda를 구할 수 있습니다.

4. Experiments

LambdaResNets

기존의 레스넷에서 3x3 convolution을 교체한 LambdaResNets입니다. 컨볼루션을 대체하는 것만으로도 정확도의 향상과 Cost Efficiency가 증가하였습니다.

Computational Benefits

람다레이어, 람다컨볼루션, 기존의 셀프어텐션의 비교 결과입니다. 낮은 메모리양과 높은 생산성을 보유하면서도 정확도는 상승했습니다.

Improve the speed-accuracy tradeoff of ImageNet classification

LmabdaResNets은 정확도를 컨트롤할 수 있으면서도 EfficientsNets보다 3.2-4.4x 빠르며 ResNet-Rs보다 1.6-2.3x 빠릅니다.

5. 결론

람다레이어라고 이름붙여진 새로운 클래스의 레이어는 입력과 그들의 contexts 사이의 상호작용을 잡아낼 수 있는 scalable 프레임워크입니다. 람다레이어는 가능한 contexts를 고정된 사이즈의 함수인 람다로 축약하고 쿼리에 적용하게 됩니다. 이러한 람다레이어를 이용한 LambdaNetworks는 매우 Costly Efficient하고 낮은 메모리 요구량으로 long-range 상호작용을 잡아낼 수 있습니다. 따라서 고해상도의 이미지 또한 입력할 수 있습니다. 특히 LabmdaResNets는 ImageNet에서 EfficientsNets보다 매우 빠르며 거의 동일한 정확도를 보여주었습니다.