Exclusive Self-Attention

이번에 소개해 드릴 논문은 아카이브에 공개된 지 2주 좀 안 된 논문입니다. 정확히는 테크니컬 리포트이고 애플에서 쓴 논문이네요.

Intro

기존 트랜스포머는 문맥(Context)을 파악하는 ‘Self Attention(SA)’과 각 위치의 특징을 업데이트하는 Feed Forward Network(FFN) 구조로 구성되어 있습니다.

self-attention 과정은 잘 아시겠지만 한번 정리하고 가면, 입력 x에 대하여 각각의 fc layer를 통과해 query, key, value 값을 계산한 뒤 query와 key 값을 내적하고 softmax를 취해서 attention map을 생성합니다.

그리고 이렇게 만들어진 attention map이 value 값에 내적합으로 적용됨으로써 attention이 적용된 y값이 생성됩니다. 여기서 저자들은 SA의 출력값이 자기 자신의 value 값과 불필요하게 높은 유사도를 높다는 문제점을 발견하고 이를 어텐션 유사도 편향(attention similarity bias) 문제라고 정의합니다. 쉽게 이야기해서 attention 과정을 거친 y값과 value v값이 코사인 유사도가 높게 계산된다는 점입니다.

저는 처음에 value값에다가 attention map을 적용해서 만들어진 것이 y이니 당연히 y와 v가 유사도가 높은 것이 이상하지는 않겠다고 생각은 했습니다만, 보통 self-attention 연산을 통해서 저희가 바라는 점은 주변 문맥을 잘 파악하는 데 집중해야 하는 것인데 그러지 못하고 자기 자신의 정보들에 집중하고 있다는 의미가 됩니다.

위의 그래프는 저자들이 1.3B LLM 모델에 대해 transformer layer 별로 특징들을 분석한 결과입니다. 가장 좌측 결과는 동일한 head와 sequence 안에서 value vector들에 대한 코사인 유사도를 측정한 것이며, 가운데 결과값은 attention map의 대각성분의 평균 score 값을 나타낸 것입니다. 우측 결과는 y값과 그에 대응되는 위치의 value 값 사이의 코사인 유사도의 평균값을 나타낸 것입니다.

위의 결과에 대해서 저자들은 다음과 같이 분석했는데, 첫째는 value 값과 attention score가 양의 상관관계를 지녔다는 점입니다. 세가지 표가 모두 비슷한 양상을 보여주고 있다는 점이며 둘째로 동일 위치에서의 value와 y 사이의 코사인 유사도가 높다는 점에서 자기 자신에 대한 특징 모델링의 역할을 수행하고 있는 것이 아닌지, FFN과의 역할이 중복되어서 수행되고 있는 것이 아닌지에 대한 우려를 저자들은 표하고 있습니다.

정리하면, 자기 자신에 대한 정보를 잘 파악하고 특징을 모델링하는 부분은 FFN의 역할이지만 그 과정을 SA도 불필요하게 반복하고 있다는 점들이 저자들이 꼽는 문제점에 해당합니다. 결과적으로 문맥을 모델링하는 것과 자기 자신의 특징을 모델링하는 것 사이에 경쟁이 발생해 모델의 효율성과 성능을 갉아먹는다고 저자들은 주장합니다.

Method

그럼 해결 방법에 대해서 알아보겠습니다. 사실 저자들이 제시한 해결책은 상당히 간단합니다. 핵심 아이디어는 SA의 최종 출력값 y에서 ‘자기 자신의 가치 벡터 v의 방향’을 명시적으로 빼버리는 것입니다. 즉, 출력값에 자신의 정보가 섞이지 않도록 직교하게 만듦으로써, 어텐션이 오로지 ‘주변 문맥 정보’만 온전히 담아내도록 강제했습니다.

아까 봤던 수식1에서 파란색 내용이 추가된 모습입니다. 저 파란색 과정이 추가되는 것이 저자들이 제안하는 XSA 방법론입니다. 쉽게 생각해서 y와 v가 서로 유사도가 높다면 y 벡터에 대하여 v 방향에 대한 성분들을 제거하자는 것입니다. 너무 간단하죠? 실제 파이토치 코드에서도 매우 간단하게 구현이 되어 있습니다.

저희가 아는 기존의 self-attention 과정을 다 마친 뒤 파란색 코드 부분만 추가하면 됩니다. value 값에 대하여 정규화하여 방향에 대한 정보만 남긴 Vn을 구한 뒤 Y와 내적하여 해당 방향에 대한 Y의 크기 값을 계산합니다. 그리고 다시 Vn 방향값을 곱해준다음에 해당 성분을 Y에 빼주면 끝이죠.

이렇게 그냥 v 방향에 대한 성분을 제거해버리는 방식이 어찌보면 극단적이라는 생각이 들 수도 있습니다. 저자들의 주장대로 attention의 결과값이 y와 value 사이의 연관성이 높은 것이 더 넓은 맥락을 파악해야 하는 관점에서 문제가 될 수도 있지만, 반대로 이 연관성이 높아야만 하는 경우가 있을 수도 있을 텐데 말이죠.

저자들은 이 부분에 대해서 크게 걱정하지는 않는 모습입니다. 설령 attention 연산을 거친 뒤 저자들의 exclusive한 방식을 통해 자기 자신에 대한 정보를 최대한 제거하더라도 결국 기존의 transformer는 attention 처리가 된 output과 input 사이의 residual connection을 수행하기 때문에 자기 자신에 대한 정보가 여전히 살아있다는 것입니다. 그리고 그 이후에 FFN을 태워서 또 모델링을 하니 역할의 중복을 최소화하고 SA가 주변들과의 관계를 더 집중할 수 있도록 도와주겠다는 것이죠.

Experiments

저자들은 Vision과 같은 다른 모달리티는 다루지 않고, text 모달리티만을 실험에서 고려합니다. 실험에 사용한 모델의 구조는 각각 아래 3가지 버전으로 세부사항은 표를 참고하시면 좋을 듯 합니다.

위의 그림2는 일반적인 transformer의 SA와 저자들의 XSA를 적용하였을 때 사용되는 시간 및 메모리 효율성을 나타낸 것으로 시퀀스 길이가 길어지면 길어질수록 exclusive 과정을 위한 내적 연산량이 늘어나긴 하지만 attention 연산처럼 길이의 제곱배만큼 늘어나고 그런 것은 아니다보니 요구되는 메모리와 추론 시간이 미미한 것을 확인할 수 있습니다.

표2의 경우에는 XSA와 baseline 트랜스포머를 각각 8개의 다운스트림 테스크에 적용하였을 때 결과값으로, 저자들의 방법론이 거의 대부분의 task에서 베이스라인 대비 더 좋은 성능을 보여주고 있습니다. 그리고 재밌는 점은 모델의 사이즈가 커지면 커질수록 XSA를 통해서 얻을 수 있는 성능의 향상 이점이 더 크다는 점으로, 모델의 크기와 학습 데이터의 크기가 더 커지는 larger scale training 상황에서 좋은 이점을 누릴 수 있다고 합니다.

다음은 1.3B 사이즈의 모델에 대하여 sequence length에 따른 XSA와 SA의 loss 변화에 대한 결과입니다. 저자들의 XSA는 Seq length가 길면 길수록 더 loss값이 작게 나오는 것을 확인할 수 있는데, 이는 입력의 길이가 더 길어진만큼 context를 파악하는 능력이 더욱 중요해지는 점에서 XSA가 자기자신에 대한 집중을 배제하고 더 좋은 맥락 파악 능력을 가진다는 것을 보여주는 결과라고 합니다.

결론

우선 테크니컬 리포트여서 그런지 글이 짧고 간결하게 작성이 되어있습니다. 저자도 한명이라서 그런지 실험은 NLP 부분에 대해서만 다루어져있으며, vision쪽 모달리티에 대해서는 실험결과가 없다는 점이 아쉽네요. 그렇지만 저자들이 문제 정의한 부분과 이를 해결하는 부분이 간결하고 직관적이라는 점, 게다가 적용하기도 매우 쉬워서 혹시나 Transformer의 context 모델링 능력에 대한 고민이 있으시던 분들은 가볍게 읽어보기 좋은 글이었던 것 같습니다.