고정 상태 캐시가 Block Diffusion의 장거리 추론을 여는 방식

들어가며

확산 언어 모델은 여러 토큰을 병렬로 디코딩할 수 있지만, 양방향 디노이징을 사용하기 때문에 자동회귀 모델의 핵심 최적화인 키-값(KV) 캐시를 그대로 적용하기 어렵다. Block diffusion은 생성 과정을 블록 단위로 진행해 캐시를 되살리지만, 기존 방식은 대부분 어텐션에 묶여 있었다. 그 결과 컨텍스트가 길어질수록 캐시 메모리가 선형적으로 증가하고, 학습이 끝난 뒤 캐시를 덧붙이면 원래 모델의 계산을 근사하는 데 그칠 수 있다.

arXiv 논문 「Fixed State, Long Reach」는 상태를 고정 크기로 요약하는 접근을 제안한다. 이미 확정된 block을 시퀀스 믹서가 재사용 가능한 상태로 압축하고, 이에 맞춘 block-causal 학습 목표를 적용해 캐시를 사후 보정이 아니라 모델이 직접 학습한 계산으로 만든다.

핵심 결과

• 동일한 조건의 비교. 연구진은 어텐션, Mamba, 하이브리드 구조의 3B block-diffusion 디노이저를 3000억 토큰으로 사전학습했다. 세 모델은 같은 single-frontier 목표와 하나의 캐시 인터페이스로 평가됐다.
• 일정한 크기의 Mamba 캐시. Mamba는 완료된 block을 고정 크기 상태로 요약한다. 따라서 캐시 메모리와 단계별 지연 시간은 시퀀스 길이에 따라 증가하지 않는 반면, 어텐션 캐시는 O(L)로 커진다.
• 256k 토큰에서 뚜렷한 차이. 이 길이에서 어텐션 캐시는 82GB를 사용하고 단계당 29ms가 필요했다. Mamba는 지연 시간을 4.3배 낮추고 메모리를 11배 줄였으며, 단일 스트림 처리량은 2.6배 높였다.
• 배치 확장성. 상태 크기가 일정하므로 Mamba는 배치 크기를 늘리기에도 유리하다. 보고된 집계 처리량은 어텐션의 14배였고, 어텐션은 단일 스트림을 넘어 실행할 수 없었다.
• 더 긴 범위의 검색. Mamba와 하이브리드 모델은 학습 길이의 8~16배에서도 정보 검색을 유지했다. 어텐션은 2배 길이에서 검색 능력이 무너졌다. 초록에 따르면 측정된 품질 저하는 없었다.

의미와 한계

이 연구의 핵심은 캐시를 단순히 축소한 데 있지 않다. 모델 구조와 학습 목표, 추론 인터페이스를 함께 설계했다는 점이 중요하다. block-causal 목표는 요약 상태를 모델의 본래 계산에 포함시켜, 양방향 디노이저에 캐시를 나중에 붙일 때 생길 수 있는 계산 불일치를 줄인다.

서비스 관점에서 고정 크기 상태는 장문 컨텍스트와 다중 스트림 추론에 특히 유리하다. 입력이 길어져도 전체 이력에 비례해 메모리가 늘지 않으며, 더 큰 배치도 현실적으로 처리할 수 있다. 또한 상태공간 모델의 귀납적 편향이 학습 범위를 넘어선 정보 활용에 도움을 줄 가능성도 제시한다.

다만 이번 결과는 3B 모델 세 개, 특정 학습 방식, 하나의 캐시 인터페이스에 기반한다. 더 큰 모델이나 다른 작업, block 설정에서도 같은 장점이 유지되는지는 추가 연구가 필요하다. 그럼에도 이 논문은 block diffusion을 장문 추론 시스템으로 확장하기 위한 일관된 설계 방향을 보여준다.

arXiv