기술 요약: 자동 분할과 블록 증류를 통한 블록 어텐션의 일반화

1. 문제 제기

대규모 언어 모델 (LLM) 은 표준 전체 어텐션 메커니즘의 2 차적 확장성으로 인해 긴 컨텍스트 입력을 처리할 때 상당한 계산 및 메모리 병목 현상에 직면합니다. 이러한 비효율성은 중복된 문서 집합이 서로 다른 쿼리에 대해 Key-Value(KV) 상태 재계산을 요구하는 검색 증강 생성 (RAG) 시나리오에서 특히 심각합니다.

블록 어텐션은 입력 시퀀스를 서로 어텐션할 수 없는 독립적인 블록으로 분할하여 사전 계산된 KV 캐시를 재사용할 수 있게 함으로써 유망한 해결책으로 부상했습니다. 그러나 블록 어텐션의 실용적 일반화는 두 가지 주요 장애물에 의해 방해받고 있습니다:

1. 분할의 어려움: 기존 방법들은 종종 의미적 일관성을 깨뜨리는 휴리스틱 규칙 (예: 개행 문자로 분할) 에 의존하여 성능 저하를 초래합니다. 다양한 텍스트 유형을 의미 있고 독립적인 블록으로 자동 분할할 수 있는 견고한 데이터 기반 방법이 부재합니다.
2. 훈련 비효율성: 블록 파인튜닝과 같이 모델을 블록 어텐션에 적응시키는 현재 방법들은 계산 비용이 많이 들고 다양한 도메인에 걸쳐 일반화하는 데 어려움을 겪습니다. 또한, 프롬프트 캐시나 슈퍼포지션 프롬프팅과 같은 훈련 없는 전략들은 전문적인 훈련 없이 일반 도메인에 직접 적용될 경우 심각한 성능 저하를 겪습니다.

2. 방법론

저자들은 이러한 과제를 해결하기 위해 학습된 분할기를 통한 자동 분할과 새로운 훈련 프레임워크인 블록 증류라는 양면적 접근법을 제안합니다.

2.1 자동 분할 (SemanticSeg)

휴리스틱의 한계를 극복하기 위해 저자들은 책, 코드, 웹 텍스트, 대화 등을 포함한 16 개 카테고리에서 2k~32k 토큰 길이의 30,000 개 이상의 인스턴스를 포함하는 대규모 데이터셋인 SemanticSeg를 구축했습니다.

• 데이터 구성: 원본 텍스트에 후보 분할 토큰을 삽입하고, 강력한 모델 (Gemini-2.5-Pro) 을 사용하여 최적의 의미적 경계를 결정합니다.
• 분할기: 사전 훈련된 언어 모델 백본 (Qwen3-4B-Instruct-2507) 과 분류 헤드를 사용하여 SemanticSeg 에서 경량 신경 분할기를 훈련합니다.
• 메커니즘: 분할기는 텍스트 세그먼트와 후보 분할 토큰을 처리하여 각 후보에 대한 이진 확률 분포를 출력합니다. 후속 텍스트의 컨텍스트를 포착하기 위해 다음 후보 토큰의 숨겨진 벡터를 사용하여 현재 후보의 분할을 결정합니다.
• 세분화 제어: 시스템은 조정 가능한 확률 임계값과 재귀적 적용 깊이를 통해 제어 가능한 블록 세분화를 허용합니다.

2.2 블록 증류

비효율적인 블록 파인튜닝을 대체하기 위해 저자들은 고정된 전체 어텐션 교사 모델을 사용하여 블록 어텐션 학생 모델을 안내하는 프레임워크인 블록 증류를 도입합니다. 이 프레임워크는 세 가지 새로운 구성 요소를 통합합니다:

1. 블록 싱크 토큰: "블록 헤드에서의 손실" 현상 (키 상태의 낮은 L2-노름으로 인해 블록 경계에서 어텐션 패턴이 비정상적인 경우) 을 해결하기 위해 각 블록의 시작 부분에 특수 토큰 <|block_start|> 을 네 번 중복합니다. 이는 블록 경계에서 어텐션 메커니즘을 안정화합니다.
2. 블록 드롭아웃: 신호 희소성 (표준 파인튜닝이 최종 블록의 신호만을 기반으로 최적화하는 경우) 을 완화하기 위해, 프레임워크는 컨텍스트 블록의 하위 집합을 무작위로 선택하여 개별적으로 인코딩합니다. 그런 다음 나머지 비손상 블록에 KL 발산 손실을 적용하여 모델이 입력 시퀀스의 더 큰 비율에서 검색 능력을 학습하도록 강제합니다.
3. 토큰 수준 손실 가중치: 표준 교차 엔트로피 손실은 모든 토큰을 동일하게 취급합니다. 저자들은 블록 어텐션과 전체 어텐션 순전파 간의 손실 차이를 기반으로 가중치를 도입합니다. 큰 성능 격차 (블록 어텐션에 민감한) 를 가진 토큰에는 더 높은 가중치를 부여하고, 민감하지 않은 토큰은 노이즈를 줄이기 위해 가중치를 낮춥니다.

최종 훈련 손실은 학생 모델의 가중 교차 엔트로피와 KL 발산 손실을 결합합니다.

3. 주요 기여

• SemanticSeg 데이터셋 및 분할기: 휴리스틱 및 통계적 기준선보다 인간 직관과 일치하는 블록 생성에서 더 우수한 성능을 보이는 다양하고 대규모의 의미 분할 데이터셋과 해당 신경 분할기의 생성.
• 블록 증류 프레임워크: 전체 어텐션 능력을 유지하면서 블록 어텐션 성능을 크게 향상시키는 블록 파인튜닝 대비 훈련 효율적인 대안.
• 새로운 훈련 구성 요소: 블록 어텐션의 고유한 실패 모드 (경계에서의 정보 손실 및 신호 희소성) 를 해결하기 위한 블록 싱크 토큰, 블록 드롭아웃, 토큰 수준 손실 가중치의 도입.
• 종합적 평가: 제안된 접근법의 확장성과 효과를 입증하는 여러 모델 계열 (Qwen, Llama) 과 벤치마크 (LongBench, LoCoMo) 에 대한 엄격한 테스트.

4. 실험 결과

4.1 분할 성능

LongBench 에 대한 실험은 제안된 분할기가 휴리스틱 기준선 (무작위, 평균, 구두점) 과 통계적 기준선 (손실, 엔트로피) 일관적으로 우세함을 보여줍니다.

• 훈련되지 않은 모델 (Qwen3-8B) 과 블록 파인튜닝된 모델 (Tulu3-Block-FT) 에 적용될 때, 분할기는 가장 높은 평균 성능을 보였으며, 의미 인식 분할이 하류 작업 성공에 중요함을 확인했습니다.

4.2 블록 어텐션의 일반화

• 블록 파인튜닝의 실패: 본 연구는 블록 파인튜닝 (이전 연구 [Ma et al., 2025] 에서 사용됨) 이 일반 도메인으로 효과적으로 일반화되지 못하며, 분할된 데이터로 훈련되더라도 전체 어텐션 기준선에 비해 상당한 성능 격차를 보임을 확인합니다.
• 블록 증류의 성공: 블록 증류는 모든 테스트된 모델 (Qwen3-4B/8B/14B, Llama-3.1-8B) 에서 블록 어텐션 설정 하에 거의 전체 어텐션 성능을 달성합니다.
  • LongBench에서 Block-Dist 모델은 원래 전체 어텐션 모델과 비교 가능한 점수를 달성했습니다 (예: Qwen3-14B Block-Dist: 42.53 vs. Original Full: 42.72).
  • LoCoMo에서도 유사한 경향이 관찰되었으며, Block-Dist 모델은 멀티홉, 싱글홉, 적대적 작업에서 높은 성능을 유지했습니다.

4.3 효율성

• 훈련: 블록 증류는 블록 파인튜닝 대비 훈련 단계당 약 26% 더 빠릅니다 (25,859.9ms vs. 34,941.1ms).
• 추론: 블록 어텐션은 첫 번째 토큰 도달 시간 (TTFT) 을 크게 단축합니다. 이 감소는 시퀀스 길이에 따라 증가하여 8k 토큰에서 약 58ms 에서 64k 토큰에서 3,100ms 이상으로 증가하며, 이는 원본 전체 어텐션 대비 수치입니다.

4.4 제거 연구

• 블록 싱크 토큰: 이를 제거하면 블록 헤드에서 검색이 필요한 합성 작업과 퓨샷 (Few-shot) 작업에서 성능이 크게 저하되어 경계 어텐션 안정화에서의 역할을 검증했습니다.
• 블록 드롭아웃: 그 부재는 퓨샷 및 합성 작업에서 성능 저하를 초래하여 최종이 아닌 블록에서 학습하는 데 필수적임을 확인했습니다.
• 토큰 가중치: 이러한 가중치를 제거하면 단일 문서 QA 및 합성 작업에서 성능이 저하되어 민감한 토큰에 학습을 집중하는 것의 중요성을 강조했습니다.

5. 중요성 및 주장

본 논문은 긴 컨텍스트 애플리케이션에서 블록 어텐션을 배포하기 위한 실용적이고 확장 가능한 경로를 확립한다고 주장합니다. 데이터 기반 학습을 통한 분할 문제 해결과 효율적인 증류를 통한 훈련 문제 해결을 통해, 저자들은 블록 어텐션이 전체 어텐션 모델의 성능을 희생하지 않고 특정 RAG 시나리오를 넘어 일반 도메인으로 일반화될 수 있음을 입증합니다.

이 연구는 블록 증류를 단순한 훈련 트릭이 아닌, 이전 파인튜닝 접근법의 한계를 극복하기 위한 필수적인 진화로 위치시키며, 모델의 일반 능력을 유지하면서 다양한 실제 시나리오에서 KV 캐시를 재사용할 수 있게 합니다.