자동 세그멘테이션 및 블록 증류를 통한 블록 어텐션의 일반화 연구

입력을 서로 어텐션(Attention)할 수 없는 별도의 블록으로 처리하는 블록 어텐션(Block attention)은 검색 증강 생성 (RAG)과 같은 긴 문맥 (Long-context) 시나리오에서 KV 캐시 재사용을 개선할 수 있는 상당한 잠재력을 제공합니다. 그러나 입력 텍스트를 의미 있고 자기 완결적인 블록으로 분할하는 것의 어려움과, 성능 저하를 초래할 위험이 있는 기존 블록 미세 조정 (Fine-tuning) 방법의 비효율성이라는 두 가지 주요 과제로 인해 광범위한 적용이 저해되고 있습니다. 이를 해결하기 위해, 우리는 먼저 책, 코드, 웹 텍스트, 대화 등 16개 카테고리에 걸쳐 2k에서 32k 사이의 텍스트 길이를 포함하는 3만 개 이상의 인스턴스로 구성된 크고 다양한 시맨틱 세그멘테이션 (Semantic segmentation) 데이터셋인 SemanticSeg를 구축합니다. 이 데이터셋을 사용하여, 우리는 텍스트를 인간의 본능과 일치하며 제어 가능한 입도 (Granularity)를 가진 블록으로 자동 분할하는 경량 세그멘터 (Segmenter)를 학습시킵니다. 둘째, 우리는 블록 미세 조정보다 효율적인 학습 프레임워크인 블록 증류 (Block distillation)를 제안하며, 이는 동결된 풀 어텐션 (Full-attention) 교사 모델을 사용하여 블록 어텐션 학생 모델을 가이드합니다. 이 프레임워크는 세 가지 새로운 구성 요소를 통합합니다: 블록 경계에서의 정보 손실을 완화하기 위한 블록 싱크 토큰 (Block sink tokens), 모든 블록으로부터 학습 신호를 활용하기 위한 블록 드롭아웃 (Block dropout), 그리고 블록 어텐션에 민감한 토큰에 학습을 집중시키기 위한 토큰 수준 손실 가중치 (Token-level loss weighting)입니다. 여러 모델과 벤치마크에 걸친 실험을 통해 우리의 세그멘터가 휴리스틱 (Heuristic) 및 통계적 베이스라인보다 우수한 성능을 보임을 입증하였으며, 블록 증류가 블록 어텐션 환경에서 풀 어텐션에 근접한 성능을 달성함으로써 블록 어텐션 배포를 위한 실용적이고 확장 가능한 경로를 구축함을 보여줍니다.