Hymba : Mamba × 트랜스포머

2024년 11월 20일 · 태그: ai, ssm · 원문 · 한국어 번역(원문에서 변경됨)

지금까지는 어텐션 메커니즘에 기반한 핵심(Core) 딥러닝 아키텍처인 트랜스포머가 있었습니다. 하지만 최근 많은 어텐션(관심)을 받고 있는 다른 실험적 아키텍처들도 있습니다. 그중 유망한 것 하나가 상태 공간 모델(state-space model)인 Mamba입니다. Mamba는 카네기 멜런 대학교와 프린스턴 대학교 연구진이 개발한 아키텍처로, 특히 긴 시퀀스와 큰 컨텍스트 윈도우를 처리할 때 드러나는 트랜스포머 모델의 몇 가지 핵심적인 한계를 해결하는 것을 목표로 합니다.

Mamba는 본질(Core)적으로 세 가지 기본 아이디어 위에 구축되어 있습니다. 첫 번째는 선택적 상태 공간(selective state space)의 사용입니다. Mamba는 트랜스포머처럼 어텐션 메커니즘에 의존하는 대신, 현재 입력에 따라 정보를 선택적으로 처리할 수 있는 SSM을 사용합니다. 이러한 선택적 처리 덕분에 모델은 필요하지 않은 정보는 버리고 관련 있는 정보에 효과적으로 집중할 수 있습니다. 이 이론은 이 블로그 글에 가장 잘 설명되어 있습니다.

Mamba는 기존 트랜스포머 아키텍처에 비해 몇 가지 핵심적인 장점을 제공합니다. 가장 중요한 이점은 선형 시간 처리를 통한 더 나은 확장성으로, 이차적으로 확장되는 트랜스포머에 비해 긴 시퀀스를 처리할 때 훨씬 더 효율적입니다. 또한 이 아키텍처는 메모리 효율성이 뛰어나며, 다른 모델의 속도를 떨어뜨릴 수 있는 메모리 집약적 레이어에서 확장된 상태를 메모리에 구체화(materialization)하는 것을 신중하게 피합니다.

두 번째 아이디어는 Mamba의 단순화된 아키텍처입니다. 어텐션 블록과 MLP 블록을 분리하는 트랜스포머와 달리, Mamba는 이들을 하나의 SSM 블록으로 통합합니다. 이러한 아키텍처 단순화는 계산 복잡도를 줄이고 추론 속도를 향상시킵니다. 그러나 개념적인 가능성에도 불구하고, SSM은 회상(recall) 능력이 제한적이며 인컨텍스트 학습(in-context learning) 작업에서 성능이 떨어지는 것으로 나타났습니다.

요컨대, 기존의 어텐션 헤드는 사진 같은 기억력처럼 특정 정보를 정확하게 기억하고 회상하는 데 뛰어나지만, 계산 비용이 많이 들고 메모리를 많이 소모합니다. 반면 상태 공간 모델은 더 효율적이지만 인간의 기억과 다소 비슷하게 작동하여, 대략적인 요지를 파악하는 데는 능숙하지만 정확한 회상에는 그다지 뛰어나지 않습니다.

Hymba

Nvidia는 서로 다른 연산 패러다임을 결합하여 상호 보완적인 강점을 활용하는 방법을 최근 탐구했습니다. 대표적인 예로 Hymba 아키텍처가 있습니다. 이 아키텍처는 트랜스포머 어텐션 메커니즘과 상태 공간 모델을 같은 레이어 안에서 통합하는 하이브리드 헤드 병렬 설계를 도입합니다. 그 결과 훨씬 적은 메모리를 사용하면서도 더 큰 모델보다 뛰어난 성능을 내는 모델이 만들어졌습니다. 자체 보고한 벤치마크에서 Hymba-1.5B는 메모리 캐시를 11.67배 적게 사용하고 3.49배 빠르게 실행되면서도 Llama-3.2-3B를 능가합니다. 이는 놀라울 정도로 인상적입니다.

Hymba와 같은 하이브리드 아키텍처의 핵심 혁신은 어텐션 헤드와 SSM 헤드의 병렬 융합으로, 이를 통해 서로 다른 연산 경로로 입력을 동시에 처리할 수 있습니다. 서로 다른 유형의 레이어를 쌓는 순차적 방식과 달리, 병렬 융합에서는 각 레이어가 어텐션의 고해상도 회상 능력과 SSM의 효율적인 컨텍스트 요약 능력을 동시에 활용할 수 있습니다. 이는 입력을 어텐션 구성 요소(쿼리, 키, 값)와 SSM 구성 요소(입력 특징과 게이트) 양쪽에 병렬로 투영하는 통합된 공식으로 구현됩니다. 어텐션 헤드와 SSM 헤드의 출력은 정규화된 후 학습 가능한 스케일링 계수 ${\beta }_1$과 ${\beta }_2$를 사용하여 결합됩니다. 구체적으로, 어텐션 출력을 $Y_{attn}$, SSM 출력을 $Y_{ssm}$이라고 하면 최종 레이어 출력은 다음과 같이 계산됩니다.

$$Y=W_{ou{t}_{p}roj}({\beta }_1\cdot \text{norm}(M_{attn}\cdot X)+{\beta }_2\cdot \text{norm}(M_{ssm}\cdot X))$$

여기서 ${\beta }_1$과 ${\beta }_2$는 각 경로의 상대적 기여도를 조절합니다.

메모리 및 효율성 최적화

어텐션 메커니즘의 메모리 요구량 문제를 해결하기 위해 하이브리드 아키텍처는 여러 가지 최적화를 적용합니다. 그중 하나는 전역 어텐션 패턴과 지역 어텐션 패턴을 결합하는 것으로, 일부 선택된 레이어(일반적으로 첫 번째, 중간, 마지막 레이어)만 전체 어텐션을 사용하고 나머지 레이어는 슬라이딩 윈도우 어텐션을 사용합니다. 또한 레이어 간 키-값 캐시 공유는 모델 성능에 큰 영향을 주지 않으면서 메모리 요구량을 줄여 줍니다.

SSM 구성 요소는 어텐션이 국소화되어 있어도 전역 컨텍스트 인식을 유지하도록 돕기 때문에, 순수 트랜스포머 아키텍처보다 슬라이딩 윈도우 어텐션을 더 적극적으로 사용할 수 있습니다. 실증 연구에 따르면 하이브리드 모델은 비슷한 작업 성능을 유지하면서도 메모리 사용량을 크게 줄이고 처리량을 늘릴 수 있습니다.

또한 이 아키텍처에는 입력 시퀀스 앞에 추가되는 학습 가능한 메타 토큰이 포함되어 있습니다. 이 토큰들은 두 가지 목적을 수행합니다. 학습된 캐시 초기화 역할을 하며 어텐션 패턴을 더 효과적으로 재분배하도록 돕고, 압축된 세계 지식을 저장합니다. 메타 토큰은 어텐션 연산과 SSM 연산 모두에 참여하며, 본질적으로 일종의 학습된 메모리 초기화를 제공하여 모델이 관련 정보에 집중하도록 이끄는 데 도움을 줍니다.

하이브리드 아키텍처는 순수 트랜스포머 아키텍처와 순수 SSM 아키텍처 모두에 비해 여러 장점이 있음을 실증적으로 보여 줍니다. 논문에서 이 아키텍처는 적정한 캐시 크기를 유지하면서도 더 나은 유효 수용 영역 측정값을 달성합니다. 어텐션 헤드와 SSM 헤드의 상대적 중요도는 레이어와 작업에 따라 달라지며, 이는 아키텍처가 필요에 따라 서로 다른 연산 모드를 활용하도록 학습한다는 것을 시사합니다. 이러한 유연성 덕분에 경쟁력 있는 효율성 지표를 유지하면서도 일반 언어 작업과 메모리 집약적 응용 모두에서 강력한 성능을 냅니다.

이러한 하이브리드 아키텍처는 매우 흥미롭습니다. 언어 모델의 미래는 아마도 서로 다른 접근 방식 중 하나를 고르는 것이 아니라, 각 방식의 가장 좋은 부분을 섞고 조합하는 데 있다는 것을 보여 주기 때문입니다. 이처럼 서로 다른 아키텍처를 영리하게 결합하면 이론적으로 두 세계의 장점을 모두 얻을 수 있습니다. 즉, 최악의 경우에만 이차 복잡도를 가지면서도 여전히 놀라울 정도로 뛰어난 성능을 내는 빠르고 효율적인 모델을 얻을 수 있습니다. 초기 결과는 매우 유망합니다.

사용하기

코드를 실행하기 전에 CUDA 12.4와 FlexAttention으로 올바른 환경이 설정되어 있는지 확인하십시오.

# Install PyTorch with CUDA 12.4
pip install pytorch==2.5.0
pip install torchvision==0.20.0
pip install torchaudio==2.5.0
pip install pytorch-cuda=12.4

# Install required packages
pip install --upgrade transformers
pip install tiktoken sentencepiece protobuf ninja einops triton packaging

# Install Mamba
git clone https://github.com/state-spaces/mamba.git
cd mamba
pip install -e .
cd ..

# Install causal-conv1d
git clone https://github.com/Dao-AILab/causal-conv1d.git
cd causal-conv1d
export CUDA_HOME=/usr/local/cuda-12.4
export TORCH_CUDA_ARCH_LIST="7.0;7.5;8.0;8.6;8.9;9.0"
python setup.py install
cd ..

# Install attention-gym
git clone https://github.com/pytorch-labs/attention-gym.git
cd attention-gym
pip install .
cd ..

# Install Flash Attention
pip install flash_attn

다음은 Hugging Face를 통해 Hymba를 사용하는 최소한의 예제입니다.

import torch
from transformers import
    AutoModelForCausalLM,
    AutoTokenizer,
    StopStringCriteria,
    StoppingCriteriaList

# Initialize model and tokenizer for chat
repo_name = "nvidia/Hymba-1.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(repo_name, trust_remote_code=True)
model = model.cuda().to(torch.bfloat16)

def chat_with_model(messages, model, tokenizer, max_new_tokens=256):
    # Format the conversation using the model's chat template
    tokenized_chat = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to('cuda')

    # Set up stopping criteria to properly end responses
    stopping_criteria = StoppingCriteriaList([
        StopStringCriteria(tokenizer=tokenizer, stop_strings="</s>")
    ])

    # Generate response
    outputs = model.generate(
        tokenized_chat,
        max_new_tokens=max_new_tokens,
        do_sample=False,
        temperature=0.7,
        use_cache=True,
        stopping_criteria=stopping_criteria
    )

    # Decode and return only the new tokens
    input_length = tokenized_chat.shape[1]
    response = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)
    return response

# Interactive chat loop
messages = [
    {"role": "system", "content": "You are a helpful assistant."}
]

def main():
    print("Chat with Hymba (type 'exit' to quit):")
    while True:
        print("User:")
        prompt = input()
        if prompt.lower() == "exit":
            break

        messages.append({"role": "user", "content": prompt})
        response = chat_with_model(messages, model, tokenizer)
        messages.append({"role": "assistant", "content": response})

        print(f"Model: {response}")