큰 중간 행렬을 피하기
질의와 키의 모든 점수 행렬을 느린 GPU 메모리에 저장한 뒤 다시 읽는 대신 작은 블록을 빠른 메모리에서 처리한다. 소프트맥스의 누적 통계와 출력을 갱신해 전체 주의 결과를 얻는다. 역전파에서는 일부 중간값을 재계산한다.
블록 정규화의 독립 예제
한 질의의 점수가 (0, ln 3), 값이 (2, 6)이면 출력은 (1×2+3×6)/(1+3)=5다. 두 원소를 따로 정규화한 결과 2와 6을 단순 평균하면 4가 되어 틀린다. 블록을 나눠도 전체 정규화 정보를 올바르게 합쳐야 한다.
정확성과 수치 오차
정확한 주의라는 말은 희소화나 저랭크 근사로 연산 자체를 바꾸지 않는다는 뜻이다. 부동소수점 연산 순서가 달라져 구현 사이에 작은 수치 차이가 생길 수는 있다. 마스크와 정밀도 등 비교 조건도 같아야 한다.
복잡도의 구분
메모리 사용과 읽기·쓰기 감소가 모든 토큰 쌍의 조밀한 주의 계산을 선형 연산량으로 바꾸는 것은 아니다. 실제 속도는 길이와 헤드 크기, 장치에 따라 달라진다. 근사 주의인 블록 희소 확장과 기본 정확 연산도 구분해야 한다.
관련 문서
FlashAttention-2 / 게이트 선형 주의 / 자기 주의 / 소프트맥스 / GPU / 트랜스포머