기술 요약: LAMP (대규모 언어 모델의 미리보기 혼합 정밀도 추론)

문제 제기

Transformer 기반 대규모 언어 모델 (LLM) 에 대한 현재 혼합 정밀도 추론 전략은 주로 입력의 균일 양자화와 출력의 고정밀도 누적에 의존합니다. 행렬 곱셈에는 효과적이지만, 이러한 접근 방식은 중간 계산의 출력을 균일하게 취급하여 심층 신경망의 구성적 특성을 간과합니다.

수치 해석학적 관점에서 부동소수점 (FP) 평가는 모든 단계에서 반올림 오차를 도입합니다. 심층 합성 $f(g(x))$에서 국소 반올림 오차는 전파되어 후속 연산자에 의해 증폭될 수 있습니다. DNN 에 대한 기존 최악의 경우 오차 상한선은 깊이에 따라 기하급수적으로 증가하지만, 표준 혼합 정밀도 방법은 중간 결과의 어떤 특정 구성 요소가 후속 층에 의해 증폭되기 가장 쉬운지 고려하지 않습니다. 이 논문은 양자화를 통한 메모리 풋프린트 감소뿐만 아니라 계산의 정확성을 제어하는 것을 목표로 하는, 구성적 효과를 명시적으로 인식하는 '혼합 정밀도 누적' 전략의 공백을 규명합니다.

방법론: 미리보기 혼합 정밀도 (LAMP)

핵심 기여는 LAMP 프레임워크로, 이는 중간 함수 $g(x)$ 내의 희소 부분 집합을 더 높은 정확도로 재계산하도록 선택하고 나머지는 낮은 정밀도로 계산하는 적응형 전략입니다. 이 선택은 $g(x)$의 반올림 오차가 후속 연산자 $f$를 통해 어떻게 전파될지에 대한 '미리보기' 분석에 기반합니다.

이론적 기반

이 방법은 합성 $f(g(x))$에 대한 정교화된 반올림 오차 분석에 기반합니다.

1. 오차 상한 유도: 저자는 합성의 상대 오차에 대한 상한을 유도합니다. 이 오차는 외부 함수 $f$의 조건수와 내부 함수 $g$에 의해 도입된 반올림 오차에 의존합니다.
2. 최적화 공식화: 목표는 전파된 오차가 임계값 $\tau$ 이하로 유지되도록 보장하면서 고정밀도 재계산의 수 (희소성) 를 최소화하는 것입니다. 이는 다음과 같은 최적화 문제로 공식화됩니다:
   $$\min \|q\|_0 \quad \text{s.t.} \quad \kappa(q) \leq \tau$$
   여기서 $q \in \{0, 1\}^n$은 $g(x)$의 어떤 구성 요소가 더 높은 정밀도를 필요로 하는지 나타내는 이진 벡터이며, $\kappa(q)$는 합성의 구성 요소별 또는 노름별 조건수를 나타냅니다.
3. Transformer 를 위한 폐쇄형 해: 중요한 이론적 통찰은 특정 Transformer 비선형성의 경우, 일반 경우의 조합적 어려움을 극복하는 폐쇄형 해가 최적화 문제에 허용된다는 점입니다:
   • 활성화 함수: 야코비안 (Jacobian) 이 대각 행렬이므로 민감한 구성 요소를 즉시 식별할 수 있습니다.
   • 레이어 정규화: 저자는 입력 크기를 정렬함으로써 거의 가장 희소한 해를 탐욕적으로 찾을 수 있음을 증명합니다.
   • Softmax (Attention): 어텐션 메커니즘의 경우, 구성 요소별 접근 방식은 매우 음수인 입력에 대해 불안정합니다. 저자는 총 확률 질량 이동을 상한하는 노름별 $\ell_1$ 목적 함수로 전환합니다. Key-Query(KQ) 내적 $y_j$를 재계산해야 하는 최적 조건은 다음과 같이 유도됩니다:
     $$2z_j(1 - z_j)|y_j| > \tau$$
     여기서 $z_j$는 결과 softmax 확률입니다. 이는 '혼란스러운'(동일한 확률) 결과나 특정 민감도 패턴을 초래하는 내적이 더 높은 정밀도를 필요로 함을 의미합니다.

구현 전략

LAMP 알고리즘은 두 단계로 작동합니다:

1. 기준선 평가: $g(x)$를 낮은 정밀도로 계산합니다.
2. 미리보기 선택: 기준선 출력을 기반으로 조건수 $\kappa$를 추정합니다. 최적화 문제를 풀어 재계산이 필요한 희소 인덱스 집합 $q$를 식별합니다.
3. 적응형 재계산: $g(x)$의 선택된 구성 요소만 더 높은 정밀도 (예: FP32) 로 재계산한 후 $f$에 전달합니다.

주요 기여

1. 새로운 계산 프레임워크: 이 논문은 구성적 오차 전파에 기반하여 행렬 곱셈의 누적을 대상으로 하는 수학적으로 원리 있는 혼합 정밀도 추론 접근법인 LAMP 를 소개하며, 기존 피연산자 양자화 전략을 보완합니다.
2. Transformer 에 대한 이론적 실현 가능성: 저자는 재계산을 위한 구성 요소 선택이라는 복잡한 조합적 문제가 표준 Transformer 빌딩 블록 (활성화, 레이어 노름, softmax) 에 대해 폐쇄형 해로 단순화됨을 입증합니다.
3. 개념 증명: 이 방법은 GPT-2 모델에서 수치적으로 검증되었으며, 특히 어텐션 메커니즘의 KQ 내적을 대상으로 합니다.

실험 결과

저자는 OpenWebText 데이터셋을 사용하여 GPT-2 XL 에서 LAMP 를 평가했으며, 다양한 가수부 비트 너비 ($\mu$) 를 가진 사용자 정의 부동소수점 형식 (PS($\mu$)) 을 시뮬레이션했습니다.

• 정확도 향상: 이 방법은 매우 낮은 재계산률로 정확도 (균일 FP32 기준에 대한 Kullback-Leibler 발산으로 측정) 에서 최대 두 자릿수만큼의 개선을 달성합니다.
  • 예를 들어, $\mu=4$(매우 낮은 정밀도) 의 경우, **0.3%**의 재계산률만으로도 KL 발산을 12 배 감소시켰습니다.
  • 1.6% 재계산 시 감소폭은 83 배였습니다.
  • 7.6% 재계산 시 감소폭은 385 배였습니다.
• 선택성: 결과는 적응형 선택이 결정적임을 확인시켜 주며, 동일한 수의 구성 요소를 무작위로 재계산하는 것은 유의미한 정확도 향상을 제공하지 않습니다.
• 완화된 LAMP: FlashAttention(한 번의 패스로 softmax 수행 필요) 과 같은 효율적인 추론 커널과의 호환성을 해결하기 위해, 저자는 '완화된 상대 임계값' LAMP 를 제안합니다. 이 변형은 정규화 상수에 대한 의존성을 제거하여 '엄격한' 최적 해에 비해 성능 손실이 미미함을 보여줍니다.

중요성과 주장

이 논문은 LAMP 를 완전히 배포된 생산 커널이 아닌 개념 증명이자 이론적 기반으로 위치시킵니다.

• 양자화와 보완적 관계: 저자는 LAMP 가 양자화를 대체하지 않는다고 명시적으로 밝힙니다 (메모리 풋프린트를 줄이거나 피연산자를 수정하지 않음). 대신 계산 정확도를 제어합니다. 이는 양자화 기술과 함께 작동하도록 설계되었습니다.
• 하드웨어 함의: 현재 실험은 낮은 정밀도 누적을 시뮬레이션합니다. 저자는 현재 하드웨어가 네이티브로 패킹된 낮은 정밀도 행렬 곱셈 누적을 지원하지 않는다고 지적합니다. 따라서 LAMP 는 이러한 연산을 네이티브로 지원할 수 있는 차세대 AI 가속기를 위한 '청사진' 역할을 합니다.
• 한계점:
  • softmax 를 위한 '엄격한' LAMP 알고리즘은 softmax 확률의 전체 벡터를 구체화해야 하므로 산업 표준인 FlashAttention 과 호환되지 않습니다.
  • 이 방법은 현재 GPT-2 모델에서 검증되었으나, 저자는 수치적 효과가 모델 깊이와 확률의 집중도에 따라 확장된다고 주장하며, 더 큰 모델에 대한 검증은 향후 과제로 남아 있습니다.
  • 초점은 엄격하게 추론에 맞춰져 있으며, LAMP 를 역방향 통과 (학습) 에 적용하는 것은 향후 작업으로 식별되었습니다.

요약하자면, 이 논문은 Transformer 합성을 통한 반올림 오차 전파를 수학적으로 분석함으로써, 소수의 적응형 선택된 부분 집합을 더 높은 정밀도로 재계산하는 조건 하에, 대부분의 계산에 대해 훨씬 낮은 정밀도로도 거의 FP32 수준의 정확도를 달성할 수 있다고 주장합니다.