샤오미가 먼저 MiMo-V3 신규 아키텍처를 공개했다: 백만 토큰 프리필 연산량 80% 감소

동찰 Beating AI 속보, 샤오미 MiMo 팀이 차세대 MiMo-V3에 적용될 HySparse2 아키텍처를 공개했다. 이는 주로 Agent가 실행할수록 비용이 커지는 문제를 해결한다. Agent는 매번 매우 짧은 명령만 내리지만, 웹페이지, 터미널, 도구로부터 대량의 내용을 돌려받을 수 있다. 모델은 먼저 이 새로운 내용을 모두 읽어야 하며, 컨텍스트가 길어질수록 이 단계에서 연산력 소모가 커지고 KV Cache도 VRAM을 더 많이 차지한다.

HySparse2는 먼저 모델을 전반부와 후반부로 나눈다. 전반부는 입력 처리를 담당하고, 후반부는 추론과 생성을 계속한다. 후반부에 필요한 컨텍스트 정보는 전반부에서 이미 계산된 결과로부터 직접 생성할 수 있으므로, prefill 시 후반부를 다시 완전히 실행할 필요가 없다. 이 아이디어는 마이크로소프트 리서치가 2024년에 제안한 YOCO에서 나왔으며, 이름은 「You Only Cache Once」로, 핵심은 후반부가 앞에서 계산된 정보를 공유하게 하여 캐시를 적게 저장하고 중복 계산도 줄이는 것이다.

또한 희소 어텐션도 다시 설계했다. 일반적인 전체 어텐션은 매번 전체 컨텍스트를 봐야 하지만, HySparse2는 소수 계층만 그렇게 한다. 이들은 먼저 긴 컨텍스트에서 가장 관련성 높은 1024개 token을 골라내고, 이후 계층은 이 결과를 그대로 이어받으며, 동시에 최근 128개 token을 고정적으로 유지한다. 이전 세대 HySparse는 64개 token마다 한 그룹으로 묶어 선택했기 때문에, 한 그룹에 중요한 정보가 있으면 그룹 전체를 남겨야 했다. 이제는 token을 하나씩 선택하는 방식으로 바뀌어, 같은 계산량으로 더 많은 실제 관련 내용에 집중할 수 있다. 이 변경만 단독으로 테스트한 결과, 두 가지 장문 검색 성적이 각각 6.57점과 8.14점 향상되었다.

논문은 총 파라미터 800억 개, 매번 약 30억 개 파라미터를 활성화하는 모델로 대조 실험을 했다. 49층 모델은 prefill 단계에서 앞 25층만 실행하면 된다. 입력이 100만 token에 도달했을 때, MiMo-V2 시리즈가 채택한 혼합 슬라이딩 윈도우 어텐션과 비교해 prefill 계산량은 약 1/5로 줄고, KV Cache는 12.09GB에서 2.69GB로 감소했다.