무한한 컨텍스트를 지원하는 획기적인 MoBA, 키미 출시!

전문가와 희박한 관심의 혼합으로 사실상 무제한의 컨텍스트가 가능합니다. 이를 통해 RAG AI 에이전트는 컨텍스트의 제한 없이 전체 코드베이스와 문서를 삼킬 수 있습니다.

📌 긴 컨텍스트 주의의 과제

시퀀스가 매우 커지면 트랜스포머는 여전히 과중한 계산 부담을 안게 됩니다. 기본 주의 모델은 각 토큰 다른 모든 토큰과 비교하면 계산 비용이 4배로 증가합니다. 이러한 오버헤드는 전체 코드 베이스, 여러 장으로 구성된 문서 또는 대량의 법률 텍스트를 읽을 때 문제가 됩니다.

📌 MoBA

MoBA(혼합 블록 주의)는 전문가 혼합 개념을 주의 메커니즘에 적용합니다. 이 모델은 입력 시퀀스를 여러 블록으로 나눈 다음, 학습 가능한 게이팅 함수가 각 쿼리 토큰과 각 블록 간의 상관관계 점수를 계산합니다. 주의도 계산에는 가장 높은 점수를 받은 블록만 사용되므로 전체 시퀀스에서 각 토큰에 주의를 기울이지 않아도 됩니다.

블록은 시퀀스를 동일한 범위로 분할하여 정의됩니다. 각 쿼리 토큰은 각 블록에 있는 키의 집계된 표현을 살펴본 다음(예: 평균 풀링 사용) 중요도를 평가하여 세부적인 관심도 계산을 위해 몇 개의 블록을 선택합니다. 쿼리가 포함된 블록은 항상 선택됩니다. 인과 관계 마스킹은 토큰이 미래 정보를 볼 수 없도록 하여 왼쪽에서 오른쪽으로 생성 순서를 유지합니다.

📌 집중과 집중 사이의 원활한 전환

MoBA는 표준 주의 메커니즘을 대체하지만 매개변수 수는 변경하지 않습니다. 표준과 유사합니다. 트랜스포머 인터페이스는 호환 가능하므로 서로 다른 레이어 또는 훈련 단계 간에 희소 주의와 전체 주의가 전환될 수 있습니다. 일부 계층은 특정 작업(예: 감독 미세 조정)을 위해 전체 주의를 예약할 수 있지만, 대부분의 계층은 계산 비용을 줄이기 위해 MoBA를 사용합니다.

📌 이는 표준 주의 호출을 대체하여 대규모 Transformer 스택에 적용됩니다. 게이팅 메커니즘은 각 쿼리가 소수의 블록에만 집중하도록 보장합니다. 인과관계는 미래의 블록을 필터링하고 쿼리의 현재 블록 내에 로컬 마스크를 적용하여 처리됩니다.

📌 아래 그림은 쿼리가 전체 시퀀스가 아닌 몇 개의 '전문가' 키/값 블록으로만 라우팅되는 것을 보여줍니다. 게이팅 메커니즘은 각 쿼리를 가장 관련성이 높은 블록에 할당하여 관심도 계산의 복잡성을 이차적에서 이차적 이하로 줄입니다.

📌 게이팅 메커니즘은 각 쿼리와 각 블록의 응집력 있는 표현 사이의 상관관계 점수를 계산합니다. 그런 다음 쿼리의 순서에 상관없이 각 쿼리에서 가장 높은 점수를 받은 상위 k개의 블록을 선택합니다.

각 쿼리는 몇 개의 블록만 처리하므로 계산은 여전히 이차적이지만, 게이팅 점수가 높은 상관관계를 보이면 모델은 현재 블록에서 멀리 떨어진 토큰으로 이동할 수 있습니다.

PyTorch 구현  

이 의사 코드는 키와 값을 블록으로 나누고 각 블록의 평균 풀링된 표현을 계산한 다음 쿼리(Q)에 풀링된 표현을 곱하여 게이팅 점수(S)를 계산합니다.

그런 다음 인과 관계 마스크를 적용하여 쿼리가 향후 블록에 집중할 수 없도록 하고, 상위-k 연산자를 사용하여 각 쿼리에 대해 가장 관련성이 높은 블록을 선택하고, 효율적인 관심도 계산을 위해 데이터를 정리합니다.

📌 플래시 어텐션 를 각각 자기 주의 블록(현재 위치)과 MoBA 선택 블록에 적용하고, 마지막으로 온라인 소프트맥스를 사용하여 출력을 병합했습니다.

그 결과 인과 구조를 보존하고 장거리 종속성을 포착하는 동시에 표준 주의의 전체 이차적 계산 비용을 피할 수 있는 희소 주의 메커니즘이 탄생했습니다.

이 코드는 각 쿼리가 몇 개의 블록에만 집중할 수 있도록 여러 전문가가 혼합된 로직과 드문드문 주의를 결합합니다.

게이팅 메커니즘은 각 블록과 쿼리에 점수를 매기고 상위 k명의 '전문가'를 선택해 키/값 비교 횟수를 줄입니다.

이렇게 하면 주의의 계산 오버헤드가 2차적 수준 이하로 유지되므로 계산이나 메모리 부담을 늘리지 않고도 매우 긴 입력을 처리할 수 있습니다.

동시에 게이팅 메커니즘은 필요한 경우 쿼리가 여전히 멀리 떨어진 토큰에 집중할 수 있도록 하여 트랜스포머의 글로벌 컨텍스트 처리 능력을 유지합니다.

이 블록 및 게이팅 기반 전략은 바로 MoBA가 LLM에서 거의 무한에 가까운 컨텍스트를 구현하는 방식입니다.

실험적 관찰  

MoBA를 사용하는 모델은 언어 모델링 손실과 다운스트림 작업 성능 측면에서 전체 주의와 거의 비슷합니다. 수십만 또는 수백만 개의 토큰이 포함된 컨텍스트 길이에서도 일관된 결과를 얻을 수 있습니다. '꼬리 토큰'으로 평가한 실험에서는 쿼리가 관련 청크를 식별할 때 중요한 장거리 종속성이 여전히 포착되는 것을 확인했습니다.

확장성 테스트 결과, 비용 곡선은 2차함수 이하인 것으로 나타났습니다. 연구원들은 100만 토큰에서 최대 6배까지 속도가 빨라지고 그 범위를 벗어나면 더 큰 이득을 얻었다고 보고했습니다.

MoBA는 전체 주의 행렬을 사용하지 않고 블록 기반 계산을 위해 표준 GPU 커널을 활용하여 메모리 친화성을 유지합니다.

최종 관찰  

MoBA는 쿼리가 어떤 블록이 중요한지 학습하고 나머지는 모두 무시하도록 하는 간단한 아이디어로 주의 오버헤드를 줄입니다.

표준 소프트맥스 기반 주의 인터페이스는 유지하면서 경직된 로컬 모델을 적용하지 않습니다. 많은 대규모 언어 모델에서 이 메커니즘을 플러그 앤 플레이 방식으로 통합할 수 있습니다.

따라서 MoBA는 사전 학습 가중치를 크게 변경하거나 많은 재학습 오버헤드를 소비하지 않고도 전체 코드베이스를 스캔하거나 방대한 문서를 요약하는 등 매우 긴 컨텍스트를 처리해야 하는 워크로드에 매우 유용합니다.