Topic #33: Slim Attention, KArAt, XAttention - 혁신적 어텐션 기법 3가지

글을 시작하며

AI 모델을 이야기하면서 ‘어텐션’만큼 많이 언급되는 토픽도 많지는 않겠죠. 특히 현재 생성형 AI 모델의 근간이라고 할 수 있는 ‘트랜스포머’ 같은 아키텍처와 관련해서는, 기본 중의 기본 (?)이라고 할 만큼 중요하다고 하겠습니다.

간단히 말하자면, ‘어텐션 메커니즘 (Attention Mechanism)’은 모델이 결정을 내릴 때 어떤 부분에 집중하고 있는지를 살짝 엿볼 수 있게 해 주는 기법입니다.

‘어텐션’이라는 구조로 모델이 입력값의 특정한 부분에 다이나믹하게 집중을 할 수 있고, 연구자들은 어텐션의 가중치를 ‘해석 가능성’, 그리고 모델이 어떤 결과값을 낼 때 왜 그런 선택을 했는지 파악하는데 활용하려고 하기도 하죠.

어텐션 메커니즘 중에 특히 두 가지 유형의 어텐션이 유명하죠:

1. 셀프 어텐션 (Self-Attention): 셀프 어텐션은 각각의 토큰이 맥락을 이해하기 위해서 입력 시퀀스의 ‘모든 다른 토큰’을 ‘볼’ 수 있게 합니다.

2. 멀티 헤드 어텐션 (MHA, Multi-Head Attention): 여러 개의 어텐션 메커니즘을 병렬적으로 실행해서, 입력 시퀀스에 숨어있는 다양한 유형의 관계와 패턴을 포착하게 합니다.

이 어텐션들이 지금 GPT, BERT, T5, Llama와 같은 거의 모든 주요 LLM의 기반이 되었죠.

또 다른 주목할 만한 어텐션 메커니즘은, 올 초 화제가 되었던 DeepSeek에서 만들어 낸 MLA (Multi-Head Latent Attention; 멀티 헤드 잠재 어텐션)가 있습니다. 튜링 포스트 코리아의 이전 에피소드에서도 다룬 적이 있는데요:

MLA는 MHA보다 한 단계 더 나아가 MHA 메커니즘을 수정해서, KV 캐시를 훨씬 더 작게 압축, 메모리 사용량을 줄일 수 있게 해 줍니다.

위의 몇 가지 예시로만 봐도, ‘새로운 어텐션 기술 = 새로운 가능성과 기능’이라는 표현이 과언이 아니라는 걸 알 수 있죠. 게다가, 어텐션을 통해서 모델의 ‘생성 과정’을 조정하거나 가이드할 수도 있다는 걸 발견하면서, 최근 연구자들은 점점 더 AI에서의 어텐션에 큰 관심을 가지는 걸 볼 수 있습니다. 다시 말해서, AI 커뮤니티가 지금의 주종인 LLM 모델을 다음 단계의 무언가로 끌어올릴 새로운 메커니즘을 찾고 있다는 뜻이겠죠?

자, 오늘은 세 가지의 최신 어텐션 메커니즘에 대해 알아보겠습니다:

1. 긴 컨텍스트를 더 빠르게 처리하고 메모리 사용량을 최대 32배까지 줄일 수 있게 해 주는 슬림 어텐션(Slim Attention)

2. 텍스트와 비디오를 포함한 긴 시퀀스에서 희소 어텐션 (Sparse Attention)의 효율성을 향상시켜주는 엑스어텐션(XAttention)

3. 기존의 어텐션과 완전히 다른 접근 방식으로, ‘학습을 하고’ ‘변화하는 상황에 적응하는’ 어텐션을 만드는 Kolmogorov-Arnold Attention (KArAt과 Fourier-KArAt)

이런 새로운 어텐션들은 어떤 것이고 어떻게 작동할까요? 그리고 이 새로운 어텐션을 적용했을 때 혜택을 볼 수 있는 모델은 어떤 것들이 있을까요? 관심이 가시나요? ^.^

그럼 한 번 알아보러 가죠!

슬림 어텐션 (Slim Attention)에 대해 알아봅시다

LLM. 지난 2년간 엄청난 발전이 있었지만, 반면에 ‘긴 컨텍스트를 다루는 작업’은 여전히 모든 LLM에게 어려운 과제로 남아 있습니다. 최근 대세를 이루고 있는 ‘추론 모델’에서는, 특히 추론 중에 새로운 토큰을 생성하면서 많은 메모리를 차지하고 작업이 느리게 되는 경우가 많죠.

‘슬림 어텐션’은, OpenMachine의 연구자들이 ‘어텐션 메커니즘에 초점을 맞춰서 이런 문제를 해결해 보자’는 목적으로 연구하고 제안한 기법입니다.

이 기법은 MHA (Multi-Head Attention)를 사용할 때와 동일한 결과를 더 빠르게, 더 적은 메모리로 얻을 수 있게 해 주는 기법으로, 거대 모델을 스케일링할 때 아주 유용합니다. 예를 들어, Whisper 같은 모델에서는 슬림 어텐션이 메모리 사용량을 8분의 1로 줄여주고, 대규모의 배치를 실행할 때 텍스트 생성 속도를 최대 5배 빠르게 해 줄 수 있습니다. T5-11B 모델 같은 경우에는 메모리 사용량을 자그마치 32분의 1로 줄여준다고 하네요!

슬림 어텐션 기법이 어떻게 이렇게 엄청난 결과를 가져다주는지, 한 번 살펴보도록 하죠.

슬림 어텐션의 작동 방식

슬림 어텐션이 비밀은, 바로 MHA(Multi-Head Attention)와 같은 개념으로 접근하되, 메모리 사용을 더 효율적으로 하게 만드는게 핵심입니다. 모델이 긴 텍스트 청크를 처리할 때, 수많은 정보를 기억해야 하죠 - 특히 K(keys)와 V(values)를 ‘KV 캐시’라고 불리는 공간에 저장하는데, 이 공간이 어마어마하게 커질 수가 있는데, 이 부분에서 효율성을 극대화하는 거죠.

보통 트랜스포머 모델이 텍스트를 처리할 때는, 메모리에 K와 V를 모두 저장합니다. 그런데 슬림 어텐션은 더 똑똑한 방법을 찾아낸 겁니다. K와 V는 같은 입력값으로부터 만들어지고, 나오고 정방행렬 (Square Matrices)을 사용하니까, 수학적인 계산을 뒤집어서 V를 따로 저장하는 대신 K에서 V를 재구성할 수 있는데 이렇게 해도 정확히 같은 결과를 만들얼 냅니다. 이 때, K에서 V를 재구성하는 두 가지 방법이 있습니다:

1. K에서 V를 재계산한 다음에 어텐션 계산을 수행합니다. 이 방법은 간단하지만 더 연산을 많이 해야 합니다.

2. 어텐션 계산을 먼저 수행한 다음에 V에서 K로 변환을 합니다. 이 방법은, 특히 한 번에 하나의 토큰을 생성할 때 더 효율적이고 빠르지만, 구현이 약간 더 복잡합니다.

두 번째 방법은 보통 ‘생성 (Generation)’ 과정에서 유리한데, 그건 메모리를 적게 사용하면서 연산량도 절약할 수 있기 때문이죠.

앞에서 이야기한 대로, Slim Attention은 K와 V를 모두 저장하지 않고, K만 저장하고 필요할 때 V를 재계산합니다. 그래서 메모리 사용량을 절반으로 줄이거나 모델이 한 번에 "볼 수 있는" 텍스트의 양을 두 배로 늘릴 수 있게 됩니다.

이 기법을 연구한 팀은 RoPE(Rotary Positional Embeddings)를 적용할 때 나타날 수 있는 잠재적인 문제도 해결했는데요. RoPE는 모델이 ‘각 토큰이 시퀀스 내에서 어디에 위치하는지’ 알 수 있도록 해 줍니다. 문제는 RoPE가 Q(queries)와 K(keys)에만 적용되고 V(values)에는 적용되지 않는다는 점인데, 그래서 슬림 어텐션으로 K에서 V를 재구성하려면 먼저 RoPE를 제거해야 할 수도 있습니다. 희소성(Sparsity) 기법을 사용하는지 여부에 따라 이 작업을 하는 두 가지 방법이 있습니다:

• Sparsity 기법을 사용하지 않는다면, 원시 (원래의) K를 캐시에 저장하고 메모리에서 K를 읽은 후 나중에 RoPE를 적용합니다. 여기서 원시 K를 사용해서 V를 계산하면 되겠죠.

• Sparsity 기법을 사용한다면, RoPE가 적용된 K를 저장하지만, 실제로 사용되는 몇 개의 K-벡터에 대해서만 이걸 제거합니다. 이런 과정을 RoPE 디코딩이라고 하는데, 다행히 (^.^;) 인코딩과 같은 사인/코사인 연산 정도를 사용하는 간단한 방법입니다.

그렇다면, 슬림 어텐션은 KV를 다루는, 이런 똑똑하고도 간단한 접근 방식을 고안해서 어떤 성과를 이뤘을까요?

슬림 어텐션이 가져다주는 장점

슬림 어텐션이 잘 적용된 경우 아래와 같은 장점이 있습니다:

• 정확도에 영향을 주지 않으면서 모델이 더 빠르게 작동하고 메모리 요구량도 절반으로 줄어듭니다 - 특히 입력값이 긴 경우, 배치 크기가 큰 경우에 도움이 된다고 해요.

  • 한 개의 입력(배치 크기 1)을 처리하는 경우라면, 처리 속도를 최대 1.8배 빠르게 만들 수 있습니다.

  • 여러 개의 입력(배치 크기 16)을 처리하는 경우라면, 2배까지 속도를 향상시킬 수 있습니다.

  • 그리고 이미 이야기했듯이, Whisper와 같은 모델에서는 메모리 사용량을 8배나 줄이고 처리 속도를 최대 5배 빠르게 할 수 있습니다. T5-11B의 경우, 메모리 사용량을 32배까지 줄일 수 있습니다.

• 실행 비용이 더 저렴합니다.

• 슬림 어텐션은 CodeLlama, Aya, Phi-3-mini, SmolLM2, LLAVA 등의 비전-언어 모델, Qwen2-Audio 같은 Audio-Text 모델, Whisper나 T5 같은 인코더-디코더 모델 등 MHA(Multi-Head Attention)를 사용하는 기존의 모델에 쉽게 적용할 수 있습니다. 재훈련이 필요없어요.

이런 여러가지 인상적인 장점이 있지만, 슬림 어텐션에는 잘 알아차리기 힘든 몇 가지 약점도 있습니다.

슬림 어텐션의 한계점

슬림 어텐션의 잠재적인 약점은 다음과 같은 것들입니다:

• 메모리를 절약하기는 하지만, 슬림 어텐션은 순간적으로 V값을 재구성하기 위해서 약간 더 많은 연산량을 필요로 합니다. 어떤 시스템의 경우에, 메모리보다 연산량이 병목이라면, 슬림 어텐션을 사용하는게 오히려 안 좋을 수 있겠죠.

• 최신 LLM들 중에 MQA(Multi-Query Attention)나 GQA(Grouped Query Attention) 같은 어텐션의 변형 형태를 사용하는 경우 - Qwen 모델이 그렇죠 - 슬림 어텐션이 작동하지 않습니다.

• RoPE를 처리하는 로직 때문에 조금 더 복잡해집니다. 특히 Sparsity 기법을 사용한다면 더욱 그렇습니다.

• 작은 사이즈의 모델인 경우, 또 입력값의 사이즈가 작은 경우는 큰 이점이 없죠. 이미 메모리 사용량이 작으니까요.

• 슬림 어텐션은 비정방 가중치 행렬 (Non-Square Weight Matrices)과는 잘 맞지 않습니다 - 예를 들어, T5-11B 같은 모델이 출력 크기 > 입력 크기인 Projection Weights를 가지게 될 때죠. 이런 큰 행렬을 역변환하는 것은 더 어렵기도 하고 비용도 많이 듭니다.

• 모든 위치 인코딩 (Positional Encoding) 기법과 잘 작동하지는 않습니다.

요약하자면, 슬림 어텐션은 메모리가 주요 병목이고 토큰 처리를 더 빠르게 해야 하는 경우에 MHA를 보완, 대체할 수 있는 유망한 방법입니다.

이렇게, 현대적인 LLM에게는 길이가 긴 컨텍스트라든가 다양한 유형의 - 비디오, 이미지, 오디오 등 - 데이터를 잘 처리하는 것이 아주 중요한데요. 이를 해결하기 위해서 새로운 어텐션 메커니즘이 등장하고 있습니다.

다음에 살펴볼 방법은, 슬림 어텐션과는 또 다른 접근 방식입니다.

엑스어텐션 (XAttention)에 대해서 알아봅시다

엑스어텐션의 핵심 아이디어

AI 모델이 데이터를 처리할 때, 그 데이터의 유형이 어떤 것이든 간에 처리를 빠르게 하도록 해 주는 한 가지 방법이 바로 ‘Block Sparse Attention’ 기법인데요. 이 기법은, 토큰들을 블록 단위로 묶고, 모든 블록 간의 관계를 계산하는 게 아니라 ‘중요한 블록 간의 관계’만 계산하는 방식으로 효율화를 추구합니다. 실제 구현은 아래와 같은 패턴으로 이루어진다고 개념적으로 이해하시면 될 것 같습니다.

• Global Attention: 특정한 글로벌 토큰 (또는 블록)이 모든 다른 토큰 (또는 블록)과 어텐션을 계산합니다.

• Sliding Window Attention: 각 토큰 (또는 블록)이 자기 주변 일정 범위 내의 토큰 (또는 블록)들과 어텐션을 계산합니다.

• Random Attention: 각 토큰 (또는 블록)이 무작위로 선택된 일부 토큰 (또는 블록)들과 어텐션을 계산합니다.

이렇게, Block Sparse Attention은 시퀀스의 길이가 길어져도 계산의 복잡도를 줄이게 설계되어 있지만, 어떤 부분이 중요한지 파악하는 것도 시간과 노력이 필요한 일이어서, 결국 속도 향상의 이점을 상쇄하는 경우가 많습니다.

여기서 바로 ‘XAttention’이 등장합니다. 엑스어텐션은, Block Sparse Attention을 훨씬 더 효율적으로 만들어주는데요, 칭화대학교, MIT, SJTU 및 NVIDIA의 연구자들이 개발했습니다. 엑스어텐션은, 중요한 부분을 찾기 위해 복잡한 계산을 수행하지 않고, 대신 아주 ‘간단한 무언가’를 "살펴봅니다": 바로 이 ‘간단한 무언가’가 바로 어텐션 행렬에서 ‘반 대각선’을 따라가는 값들의 합인데, 이게 ‘가장 중요한 패턴’을 찾아내는 좋은 지름길이라고 밝혀졌습니다.

상식적으로 생각해 보면, 지역적인 맥락 (주변 단어)과 텍스트 전체를 가로지르는 먼 관계를 모두 포착하는게 바로 이 ‘반 대각선 (Antidiagonal)’ 패턴이죠. 추리소설을 읽는다고 생각해 보세요 - 처음 나오는 단서, 그리고 그게 마지막의 반전과 어떻게 관련되어 있는지 집중하면서도, 그 사이에 연결되는 내러티브를 이해하는 셈이죠. 이런 접근 방식으로, 거의 동일한 수준의 이해력을 유지하면서 모델을 최소 2~3배 더 빠르게 실행할 수 있습니다.

반복해서 말씀드리자면, 엑스어텐션이 특별한 건, 고정된 규칙이 아니라 어떤 관계가 가장 중요한지를 학습한다는 점입니다 - 이런 ‘적응성’ 덕분에 사람의 지도 없이도 다양한 유형의 데이터에서 잘 작동합니다.

엑스어텐션의 작동 방식

그럼, 엑스어텐션이 어떻게 작동하는지 단계별로 한 번 살펴보겠습니다.

• 우선, 시작은 ‘중요도를 예측’하는 겁니다. 이야기했다시피, 핵심적 아이디어는 ‘반 대각선(Antidiagonals)’ - 각 블록 내의 왼쪽 하단에서 오른쪽 상단으로 이어지는 선 - 에 집중하는 겁니다. 이 선을 따라 값들을 합산해서, 모델은 블록이 얼마나 중요한지 빠르게 추정합니다.

• 위 이미지를 예로 들면: 입력이 8×8 블록으로 나뉘어 있죠. 빨간 선은 더 높은 합산 값 (더 중요함)을 나타내고, 파란 선은 더 낮은 값 (덜 중요함)을 나타냅니다. 그런데 왜 이런 방식이 효과가 있을까요? 그건, ‘모든 토큰이 최소한 하나의 반 대각선에 기여하기 때문’입니다. 게다가, 수직선이나 사선과 같은 높은 영향을 미치는 패턴들은 이런 선들과 교차하는 경향이 있습니다. 결국, 이런 계산 방식이 엑스어텐션에게 우선순위를 정하는 ‘신뢰할 수 있고 빠른 시그널’을 제공하는 셈입니다.

• 그 다음으로, ‘유지할 블록’을 선택합니다. 반 대각선을 사용해서 블록에 점수를 매긴 후:

  • 매겨진 점수는 소프트맥스(Softmax)를 사용해서 정규화, 확률로 변환합니다.

  • 이 점수를 바탕으로, 모델은 높은 점수를 받은 블록만 "선택"하고 이렇게 선택된 블록들에 대해서만 어텐션을 수행합니다 (위 첫 번째 이미지에서 빨간색으로 표시). 이게 ‘결합된 중요도 점수’가 특정 임계값을 넘는 가장 작은 블록 집합입니다.

• 이제, 각 어텐션 헤드에 대한 임계값을 조정합니다. 모델의 다른 부분들, 즉 어텐션 헤드들은 서로 다르게 작동합니다 - 일부는 더 많은 세부 정보가 필요하고, 다른 것들은 더 Sparse할 수 있습니다. 엑스어텐션은 각 헤드별로 임계값을 개별적으로 조정하기 위한 다이나믹 프로그래밍 전략을 사용합니다. 속도와 정확도 사이의 적절한 균형을 찾기 위해서 각 단계마다 임계값을 10%씩 천천히 낮춥니다.

이 마지막 부분은 성능을 더욱 미세하게 조정하는 데 도움이 되지만, 옵션이라고 볼 수 있습니다 - 엑스어텐션은 꼭 이 과정이 없어도 잘 작동하기는 합니다.

자, 이제 이 엑스어텐션의 성능, 그리고 주요 장점을 볼까요?

엑스어텐션의 성능 상 이점

엑스어텐션은 다른 희소 어텐션 (Sparse Attention) 메커니즘에 비해서 꽤나 인상적으로 성능을 향상시키는데, 어떤 경우에는 풀 어텐션을 하는 경우와 비슷한 성능을 보여주기까지 합니다.

• 효율성과 속도: 긴 시퀀스의 경우 최대 13.5배 더 빠른 어텐션 계산을 할 수 있고, 비교적 더 가벼운 패턴 덕분에 짧은 입력값이어도 다른 기법보다는 더 성능이 좋습니다. 메모리와 계산 비용을 크게 줄여주죠.

• 경량의 설계: 엑스어텐션의 패턴 선택은 최대 24.9배 더 빠릅니다.

• 향상된 텍스트 이해 능력: 긴 텍스트를 다룰 때 다른 희소 어텐션 기법보다 더 나은 성능을 보여주고, 시퀀스의 길이가 아주 길어도 정확도를 잘 유지합니다.

  • RULER에서는 풀 어텐션 (Full Attention)을 능가하는 성능을 보여주기도 합니다.

  • LongBench에서는 전체적인 작업에서 가장 평균적으로 높은 정확도를 보여줍니다.

• 최고 수준의 비디오 이해 능력: 엑스어텐션은 초당 1프레임으로 최대 1시간 길이의 비디오를 처리할 수 있고, 다른 희소 어텐션 기법, 심지어 FlashAttention도 능가하는 비디오 이해 능력을 보여줍니다.

• 높은 비디오 생성 효율성: 풀 어텐션 (Full Attention)에 아주 근접한 성능을 보여주는데도 50% 이상의 계산량이 절약됩니다.

• 플러그-앤-플레이 (Plug-and-Play): 엑스어텐션은 큰 변경이 없이도 기존의 트랜스포머 모델에 쉽게 추가할 수 있습니다.

• 우수한 확장성: LCTM (Long-Context Transformers)를 위해 설계되어서, 모델이 256k+에 이르는 긴 토큰이나 한 시간 길이의 비디오를 처리하게끔 스케일링하는 데 도움을 줍니다.

전반적으로, 엑스어텐션은 계산 비용에 많은 돈을 쓰지 않고도 실제 환경에서 강력한 모델을 만들 수 있게 해 줍니다. 또한, 경량의 설계, 플러그-앤-플레이 등의 특성 때문에 특히 멀티모달 AI를 포함하는 실제 AI 시스템을 만들 때 실용적입니다.

엑스어텐션의 한계점

물론, 여전히 고려해 봐야 할 몇 가지 이슈가 있죠.

• 반 대각선 (Antidiagonals)을 따라 스코어링을 하는 건 영리한 기법이지만, 여전히 근사치라는 한계는 있죠. 드물게, 중요한 패턴을 놓칠 수도 있습니다.

• 희소성과 정확도 사이의 아슬아슬한 균형: 너무 높은 임계값이나 너무 큰 보폭(Stride)을 선택하면, 어텐션이 너무 희소해져 정확도에 손실 (Loss)이 있을 수 있습니다.

• 비디오를 생성할 때 웜업 (Warm Up)이 필요: 비디오를 생성하는 과정 초반에 엑스어텐션을 너무 일찍 적용하면 레이아웃 문제가 발생한다고 합니다. 이 문제는 처음 5단계 동안 Full Attention을 사용하는 "웜업 단계"를 추가해서 해결할 수는 있습니다만, 어쨌든 복잡성을 더하는 것이죠.

• 블록 구조에 맞게 설계되어 있어, 엑스어텐션 구조는 불규칙하거나 아주 다이나믹한 입력에는 최적화가 잘 되지 않을 수도 있습니다.

• 많은 모델에 즉시 사용 가능하지만, 모든 종류의 변형된 트랜스포머 모델들을 대상으로 직접 테스트하지는 못했습니다. 그래서 비 표준적인 어텐션 형식을 사용하는 모델은 재검토와 조정이 필요할 수도 있습니다.

이런 여러가지 한계가 있긴 하지만, 엑스어텐션은 ‘긴 시퀀스에서 약한 모습’을 보이는 Sparse Attention 메커니즘을 위한 ‘강력한 업그레이드 기법’이라는 건 입증되었습니다.

자, 그럼 이제 오늘의 마지막, 세 번째 어텐션 기법, Kolmogorov-Arnold Attention으로 넘어가서, 어텐션 메커니즘을 향상시키기 위해서 이 기법은 뭘 하는 건지 살펴보겠습니다.

KArAt (Kolmogorov-Arnold Attention)에 대해 알아봅시다

어텐션 메커니즘에 KAN을 적용하겠다는 아이디어는 어떻게 생겼을까?

KAN (Kolmorogov-Arnold Network)에 대해서 이전에 튜링 포스트 코리아의 AI 101 시리즈에서 커버한 적이 한 번 있었는데요:

그 때도 상당히 개념적으로 이해하는데 애를 먹었던 것으로 기억합니다 ^.^;

아주 간단히 리캡 (Recap)을 해 보면:

여기까지만 하고, 다시 돌아가서 KArAt 이야기를 해 보죠.

위에서 이야기한 두 가지의 독특한 어텐션 메커니즘과는 완전히 다른 접근 방식을 센트럴 플로리다 대학교의 연구자들이 고안한 겁니다. 다소 철학적이라고까지 할 수 있을 이 생각이, 바로 콜모고로프-아놀드 네트워크(KANs, Kolmogorov-Arnold Networks)가 어떻게 비전 - 시각 데이터 처리 - 작업의 영역에서 트랜스포머의 성능를 향상시킬 수 있는가 였습니다. 기존의 구조에서 효율을 추구하는 것이 아니라, 학습 방식 자체를 더 스마트하고 유연하게 하는 방법을 추구한 것이죠.

그래서 이 친구들의 핵심 질문은, “어텐션 자체가 학습을 할 수 있다면 어떨까?”하는 거였습니다.

사실, KAN은 수학 공식을 발견하거나 시간의 변화에 따른 단순한 1차원 패턴을 학습하는 등의 작업에서는 훌륭한 결과를 보여주었습니다만, KAN을 트랜스포머 같은 더 강력한 아키텍처와 결합하는 시험은 많이 이루어지지 않았는데요.

과연 KAN을 기반으로 한 ‘학습하는’ 어텐션 메커니즘이, 비전 트랜스포머(ViTs, Vision Transformers)에서 전통적으로 사용되는 ‘소프트맥스’ 함수를 대체하고 더 우수한 결과를 보일 수 있는지 확인해 보기 위해서, 이 연구자들은 콜모고로프-아놀드 어텐션, KArAt이라는 버전을 만들었고, 나중에 푸리에 특징을 사용해서 더 실용적이고 모듈로 구성된 Fourier-KArAt을 만들었습니다.

자, 쉽지가 않죠? 어쨌든, 이 메커니즘이 어떻게 작동하는지 계속해서 살펴보겠습니다.

KArAt는 어떻게 작동하나?

비전 트랜스포머(ViT, Vision Transformer)에서, 각각의 어텐션 헤드는 입력값을 "살펴보고" 어떤 부분에 주의를 기울일지 "결정"해서 어텐션 행렬을 구축하게 되죠. 이 행렬에서 각 행은 입력의 한 부분이 다른 모든 부분에 얼마나 주의 (어텐션)를 기울이는지 나타냅니다. 각 행은 소프트맥스 함수를 통과하게 되는데, 이 과정을 거쳐서 원시 숫자의 행을 확률로 변환합니다 (모든 숫자가 양수가 되고 합이 1이 됩니다).

KArAt은, 위 아키텍처에서 ‘소프트맥스 함수를 KAN을 사용한 맞춤형의, 학습 가능한 활성화 함수’로 대체합니다. 그 결과로, 모델은 데이터로부터 ‘자체적인’ 변환 구조를 학습하게 됩니다.

이런 접근 방식 자체가 콜모고로프-아놀드 정리에서 영감을 받은 건데, 이 정리는 ‘복잡한 함수가 어떻게 더 단순한 함수로 분해될 수 있는지’ 설명하는 것이구요. 어쨌든 여기서 KArAt은 기저 함수(Basis Functions)를 사용해서 함수를 구축하게 되는데, 이 기저 함수는 다음과 같은 수학적 구성 요소와 같습니다 (같다고 합니다 ^.^;):

• 푸리에 기저 (Fourier Basis): 사인과 코사인 파동의 조합

• B-스플라인: 부드러운, 곡선의 맞춤형 함수

• 더 특수한 경우를 위한 웨이블릿이나 프랙탈

각각의 어텐션 행은 이렇게 구축된 ‘사용자 정의 함수’ 중에 하나를 통과하고, 다시 이 함수는 다양한 기저 함수를 혼합해서 만들어진다는 겁니다. 모델이 학습하는 건, ‘각각의 기저 함수를 얼마나 사용할지’ 하는 것이고, 이게 바로 훈련 중에 조정하게 되는 파라미터입니다.

하지만 출력값이 ‘확률’로 동작하지 않을 수 있기 때문에, 특별한 ℓ1 투영(Projection)을 통해서 값을 양수로 만들고 합이 1이 되도록 해서 이 문제를 해결합니다.

그런데, 여기 현실적인 문제 한 가지가 생깁니다: 바로 ‘계산 비용이 너무 많이 든다’는 겁니다. 그래서 어떻게 할까요? 아키텍처를 더 작고, 더 스마트하게 만들겠다는 방향이었고, 그 결과 나오게 된 게 바로 ‘Fourier-KArAt’입니다.

푸리에 (Fourier)-KArAt는 그럼 어떻게 작동하나?

어텐션 행렬은 많은 경우에 ‘저차원 구조’로 되어 있어서, 중요한 정보를 표현한다고 하더라도 모든 차원을 사용할 필요가 없습니다. 그래서, 연구자들은 학습 가능한 연산자 (Learnable Operator)의 크기를 축소하는 방법을 제안했습니다:

• N × N 값을 가진 전체 행렬을 사용하는 대신, 훨씬 작은 r × N 행렬을 사용할 수 있습니다 (여기서 r은 N보다 훨씬 작음)

• 이 작은 행렬은 어텐션 데이터를 더 낮은 차원의 공간으로 압축합니다

• 그 다음, 다른 행렬이 이걸 원래 크기로 다시 투영 (Projection)합니다. 일종의 트릭이라고 할 수 있는 이 방법은, 메모리 사용량을 드라마틱하게 줄여 줍니다.

이 새로운 경량의 버전을 Fourier-KArAt (푸리에 기저를 가진 콜모고로프-아놀드 어텐션)이라고 합니다.

왜 푸리에 기저 함수가 다른 옵션보다 더 좋을까요? 간단합니다. 이게 모든 부드러운 주기적 함수를 표현할 수 있고, GPU 친화적이고, 더 잘 확장됩니다.

Fourier-KArAt은 실제의 ViT 모델에서 사용할 수 있습니다. 그리고 학습 가능한 어텐션 모듈은 두 가지 방식으로 구성할 수 있습니다:

1. 블록별 (Blockwise): 트랜스포머의 각 층(블록)은 자체적인 고유한 어텐션 함수를 가집니다. 이 구조는 더 상세한 패턴을 포착하지만 (당연하게도) 더 많은 파라미터를 사용합니다.

2. 범용 (Universal): 모든 층이 동일한 어텐션 함수를 공유하는데, 이렇게 해서 모델을 더 작게 유지하고 더 단순한 작업이라면 더 잘 작동합니다.

결국 이 두 가지 다른 설정은 ‘유연성 vs. 효율성’의 트레이드오프 사이의 결정이네요. 그렇다면, Fourier-KArAt가 어떤 성능 상의 이점이 있을까요?

푸리에 (Fourier)-KArAt의 장점

Fourier-KArAt이 명확한 장점을 보여주는 시나리오는 아래와 같습니다:

• 작은 모델에서 더 높은 정확도: ViT-Tiny에서 Fourier-KArAt은 추가적인 기법을 적용하지 않아도 정확도를 높여줍니다:

  • CIFAR-10에서 +5.4%

  • CIFAR-100에서 +7.4%

• 더 선명한 어텐션 맵: 시각화를 통해서 살펴보면, Fourier-KArAt이 이미지의 중요한 부분에 더 정확하게 집중하는 경향이 있음을 확인할 수 있습니다 - 모델이 더 명확한 결정을 내리는데 도움이 된다는 이야기겠죠.

전반적인 이점을 요약해 본다면:

• 학습 가능한 어텐션: Fourier-KArAt은 모델이 입력 데이터에 집중하는 방식을 훈련, 적응하는데 있어서 높은 유연성을 보여줍니다.

• 푸리에를 사용한 부드러운 근사: 사인과 코사인 함수를 사용하면 모델이 학습 가능하고 GPU 친화적인 상태를 유지하면서 소프트맥스와 유사한 동작을 부드럽게 근사할 수 있습니다.

• 저차원 설계를 통한 계산량 감소: 무거운 전체 크기의 연산 대신, 어텐션 메커니즘의 더 낮은 차원 버전을 사용해서 메모리와 계산 비용을 줄입니다.

• 다양한 아키텍처에 유연하게 적용: 이 방식의 설계라면 기본적으로 어떤 수학적 기저 함수와도 매치해서 작동시킬 수 있기 때문에, 이후 다양한 어텐션 동작을 실험하기 쉽습니다.

그런데 말입니다. Fourier-KArAt은, 우리가 원하는 그 만큼 잘 작동하지는 않는 것 같습니다.

푸리에 (Fourier)-KArAt의 한계점

Fourier-KArAt는 이런 한계점을 드러냅니다. 이 한계점은 생각보다 현재 AI 모델 개발의 트렌드라는 관점에서 큰 임팩트가 있는 것 같습니다:

• 큰 모델에서는 분명한 이점이 없습니다: ViT-Small과 ViT-Base에서는 정확도 향상이 미미하거나 일관성이 없고, 어쩔 때는 베이스 모델보다도 더 안 좋은 결과를 보여줍니다.

• Fourier-KArAt은 일반적인 어텐션보다 구현하기가 더 복잡합니다: 어텐션 맵을 해석하기가 더 어렵고, 시각화를 하려면 추가적인 조정이 필요합니다.

• 많은 파라미터를 추가해야 합니다: 그래서, 모델을 훈련하고 조정하기가 더 어려워질 수 있는데, 특히 ViT-Base 같은, 더 깊거나 큰 아키텍처라면 더 그렇습니다.

• Fourier-KArAt은 훈련 경로가 덜 안정적인데, 특히 대형 모델에서 그런 모습을 보입니다.

• 처음에는 빠르게 학습하지만, 후반 단계에서는 훈련 속도가 느려집니다.

• 높은 메모리 사용량: Fourier-KArAt으로 훈련하면 GPU 메모리를 많이 소비할 수 있고, 어떤 경우에는 단일한 이미지에 대해서 최대 60GB까지도 사용하게 될 수 있습니다.

어때요? 현대적인, 최근의 AI 모델 발전 방향을 생각해 봤을 때 문제가 좀 많아 보이죠?

Fourier-KArAt은 현재로는 이점보다 제약사항이 더 많지만, 그래도 그 잠재력만은 엄청난, 신선하고 대안적인 어텐션에 대한 접근 방식임에는 틀림없습니다. 다른 많은 새로운 기술과 마찬가지죠. Fourier-KArAt을 더 갈고 닦아서 쓸만하게 만들려면, 더 많은 시간과 실험이 필요한 듯 합니다.

세 가지 어텐선 기법의 비교

자, 이렇게 해서, 오늘 세 가지 새로운 어텐션 메커니즘을 나름대로 깊이 살펴봤습니다. 세 가지 기법 모두, 트랜스포머 기반 모델의 성능을 향상시키기 위해서 어텐션의 다양한 측면을 개선하고자 시도합니다 - 물론, 어떤 기법이 절대적으로 좋다는 건 없죠. 목표가 뭐냐, 환경이 뭐냐에 따라서 어떤 기법이 좋은지 갈라질 겁니다:

• 아주 긴 컨텍스트를 처리하면서 메모리 사용량을 크게 줄여야 한다면 - 슬림 어텐션이 최선의 선택입니다.

• 비디오를 포함한 긴 시퀀스를 효과적으로 정말 빠르게 처리해야 할 때 - 전통적인 희소 어텐션보다 더 단순하고 13.5배까지 더 빨라질 수 있는 엑스어텐션이 훌륭한 선택입니다.

• 학습하고, 적응하고, 유연한 어텐션 메커니즘을 꼭 원한다면 - KAN을 사용한 Fourier-KArAt을 검토해 볼 필요가 있습니다.

오늘 알아본, 이런 새로운 메커니즘들이, 셀프 어텐션이 끌고 온 현재의 성과 위에 미래의 LLM을 더 뛰어난 모델로 만드는데 큰 기여를 할 거라고 생각합니다.

보너스: 참고자료

1. Slim attention: cut your context memory in half without loss of accuracy – K-cache is all you need for MHA

2. Kolmogorov-Arnold Attention: Is Learnable Attention Better For Vision Transformers?

3. XAttention: Block Sparse Attention with Antidiagonal Scoring

4. Attention Is All You Need

5. Effective Approaches to Attention-based Neural Machine Translation

6. Neural Machine Translation by Jointly Learning to Align and Translate

7. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model

읽어주셔서 감사합니다. 친구와 동료 분들에게도 뉴스레터 추천해 주세요!