상수 크기 캐시가 대규모 확산 모델에 제공하는 이점
Fixed State, Long Reach: What a Constant-Size Cache Buys Block Diffusion at Scale
arXiv확산 언어 모델(Diffusion LLM)의 추론 효율성을 높이기 위해, 컨텍스트 길이에 관계없이 메모리 사용량이 일정하게 유지되는 상수 크기 캐싱 기술이 개발되었습니다.
- 기존 어텐션 방식은 컨텍스트 길이($L$)에 비례하는 메모리 사용량 문제를 안고 있었으나, 새로운 블록 캐시는 이를 상수 시간($O(1)$)으로 해결했습니다.
- 이는 지연 시간을 획기적으로 줄이고 배치 크기 확장성까지 확보하여, 모델이 수십만 토큰에 달하는 초장문맥을 효율적으로 처리할 수 있게 합니다.
- 다만 이 캐시를 구현하려면 블록 요약 기능과 특정 학습 목적 함수가 필수적이며, 기존 어텐션 구조와는 근본적인 아키텍처 변화가 요구됩니다.
확산 언어 모델(Diffusion )은 을 병렬로 디코딩하지만, 양방향 디노이저 특성상 일반적인 키-값(KV) 캐시를 이용한 빠른 자가 회귀 추론이 어렵습니다. 블록 확산 방식은 이를 블록 단위로 디코딩하여 캐싱 기능을 복원합니다. 기존의 블록 캐시는 어텐션에 종속되어 메모리 사용량이 $O(L)$이었으나, 시퀀스 믹서가 최종화된 블록을 재사용 가능한 상태로 요약하는 방식을 통해 이 문제를 해결하고 정확한 블록 캐시를 구현할 수 있습니다.
연구진은 세 가지 블록 확산 디노이저(어텐션, mamba, 하이브리드)를 300B 토큰으로 사전 학습시키고 단일 캐시 인터페이스로 테스트했습니다. 상태 공간 캐시는 시퀀스 길이에 대해 $O(1)$을 유지하여 메모리와 단계별 지연 시간이 일정하게 유지되는 장점이 있습니다. 특히 256k 토큰 환경에서 어텐션 방식이 82GB의 메모리와 29 ms/step에 도달한 반면, Mamba 캐시는 4.3배 낮은 지연 시간, 11배 적은 메모리, 그리고 2.6배 높은 단일 스트림 처리량을 제공했습니다.
상수 상태 공간 캐시를 사용하면 컨텍스트 길이에 관계없이 메모리와 단계별 지연 시간이 일정하게 유지됩니다. 이 덕분에 Mamba와 하이브리드 백본은 학습 길이의 8~16배까지 정보를 검색할 수 있는 반면, 어텐션 방식은 측정된 품질 저하 없이 정보 검색이 2배까지만 가능했습니다.
용어 풀이
- LLM
- 방대한 글을 학습해 문장을 이해하고 만들어 내는 AI 모델. ChatGPT, Claude, Gemini가 여기에 속해요.
- 토큰
- AI 모델이 글을 처리하는 단위. 단어보다 작은 조각이며 사용량과 요금을 셀 때 기준이 돼요.