Power Attention: Softmax Attention를 개선한 Symmetric Power Transformer 논문의 PyTorch 구현체

읽을거리&정보공유
, , ,
1

Power Attention
Power Attention998×285 10.7 KB

Power Attention 프로젝트 소개

Transformer 모델은 현대 AI 기술의 핵심으로 자리 잡았지만, Softmax를 활용해 쿼리(Q), 키(K), 값(V) 간의 연산을 수행하는 기존 Attention 메커니즘은 시퀀스 길이가 증가할수록 O(n^2) 의 연산 복잡도를 가지며 이에 따라 메모리와 연산 리소스 소모가 매우 큽니다. Power Attention은 이를 해결하기 위한 대안 중 하나로, 효율성과 성능을 동시에 추구하는 선형 Transformer를 제공합니다.

Power Attention 프로젝트는 PyTorch 기반의 확장 라이브러리로, 대칭적 파워 트랜스포머(Symmetric Power Transformers)의 구현체입니다. 이 모델은 Buckman 외 연구진이 발표한 “Symmetric Power Transformers” 논문을 기반으로 하며, 선형적 계산 복잡도를 제공하는 새로운 Transformer 메커니즘입니다.

Power Attention은 기존의 Attention 메커니즘을 Symmetric Power Embedding으로 대체하여 RNN 기반의 선형 연산 복잡도를 달성합니다. 그 결과 긴 시퀀스를 처리하는데 필요한 메모리와 연산 비용이 대폭 감소하여 대규모 데이터를 학습하는 경우 효율성을 크게 향상시킬 수 있습니다. 최신 GPU 환경에 최적화된 CUDA 기반의 구현과 강력한 수학적 이론을 바탕으로 한 이 프로젝트는 연구와 산업 현장에서 모두 유용할 것으로 기대합니다.

주요 기능

Power Attention의 핵심 기능은 다음과 같습니다:

  1. 선형적 Attention 계산: 기존 Softmax Attention의 O(n^2) 복잡도를 O(n) 으로 감소시킴으로써 긴 시퀀스를 처리하는 데 필요한 리소스를 획기적으로 줄였습니다.

  2. 대칭적 Power 임베딩: 대칭적 Power 임베딩(Symmetric Power Embeddings)은 쿼리(Q)와 키(K)를 대칭적 텐서 공간으로 매핑하여 연산을 단순화하고 메모리 효율성을 극대화합니다. 이 때, 파라미터 p를 통해 임베딩 차원을 조정 가능하며, 성능과 상태 크기 사이의 균형을 조정할 수 있습니다.

  3. Rotary Embedding 호환성: 기존 선형 Transformer에서 부족했던 Rotary Positional Encoding과의 호환성을 제공하여 다양한 NLP 및 시계열 작업에 활용될 수 있습니다.

  4. CUDA 최적화: 최신 GPU 환경(A100, H100 등)을 활용할 수 있도록 CUDA 커널이 최적화되어 있어 FP16/BF16 지원을 통해 학습 및 추론 속도를 극대화합니다.

  5. Chunked 알고리즘: 시퀀스를 효율적으로 분할하여 처리하는 Chunked 알고리즘을 지원하며, 이는 O(n) 복잡도에서 선형 Transformer의 성능을 극대화합니다.

사용 예시

Power Attention은 PyTorch 환경에서 쉽게 사용할 수 있습니다. 다음은 기본적인 사용 예제입니다:

import torch
from power_attention.power_full import power_full

batch_size = 2
seq_len = 1024
num_heads = 8
head_dim = 64

# Q, K, V 텐서 생성
Q = torch.randn(batch_size, seq_len, num_heads, head_dim, device='cuda', dtype=torch.float16)
K = torch.randn_like(Q)
V = torch.randn_like(Q)

# Power Attention 수행
output = power_full(Q=Q, K=K, V=V, deg=4, chunk_size=128)
print(output.shape)

또한, 다음과 같이 기존 Transformer의 Attention 레이어를 대체하여 성능과 효율성을 동시에 향상시킬 수 있습니다:

from power_attention.power_full import power_full

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        # ... 초기화 코드 ...

    def forward(self, x):
        # Power Attention 사용
        y = power_full(Q=q, K=k, V=v, deg=4, chunk_size=128)
        return y

"Symmetric Power Transformers"논문에 따르면, Power Attention은 성능과 효율성을 동시에 확보한 결과를 보여줍니다:

파라미터 p 상태 크기 성능 GPU 메모리 제한 (80GB)
p=2 39MB 기준 모델 대비 1.03배 :white_check_mark:
p=4 14GB 기준 모델 대비 0.98배 :white_check_mark:
p=6 2.2TB 기준 모델 대비 0.97배 :x:

라이선스

Power Attention은 Apache 2.0 License로 배포됩니다. 상업적 사용에 대한 제한이 없으며, 자유롭게 수정 및 배포 가능합니다.

:scroll: [Symmetric Power Transformers 논문

:books: Power Attention 공식 문서

https://power-attention.readthedocs.io/

:github: Power Attention GitHub 저장소

https://github.com/m-a-n-i-f-e-s-t/power-attention



이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있습니다. 관심있는 내용이시라면 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다. :hugs:

:pytorch:파이토치 한국 사용자 모임:kr:이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일:love_letter:로 보내드립니다! (기본은 Weekly지만 Daily로 변경도 가능합니다.)

:gift: 아래:arrow_lower_right:쪽에 좋아요:+1:를 눌러주시면 새로운 소식들을 정리하고 공유하는데 힘이 됩니다~ :star_struck: