SoloQ, 확산 언어 모델을 위한 무교정 저비트 양자화 제안

들어가며

확산 언어 모델은 자기회귀 모델처럼 토큰을 왼쪽에서 오른쪽으로 하나씩 생성하는 대신, 마스크되거나 노이즈가 포함된 토큰을 양방향으로 갱신하며 문장을 완성합니다. 이 방식은 새로운 생성 구조를 제공하지만, 추론 효율 측면에서는 다른 부담을 만듭니다. 전체 시퀀스를 처리하는 확산 모델은 여러 디노이징 단계에서 비용이 큰 전방 패스를 반복 실행해야 합니다. 블록 확산 모델은 여기에 더해 이미 처리한 블록의 KV 캐시를 계속 저장하므로, 긴 생성에서는 메모리 사용량이 커질 수 있습니다.

저비트 가중치·활성값 양자화는 이러한 비용을 줄일 수 있는 방법입니다. 하지만 확산 모델의 활성값 분포는 마스킹 상태와 디노이징 단계에 따라 바뀝니다. 한정된 보정 데이터로 얻은 범위나 통계가 모든 생성 상태를 대표하지 못할 가능성이 있는 것입니다. SoloQ는 이 문제를 보정 데이터를 더 많이 모으는 방식이 아니라, 양자화하기 쉬운 표현을 먼저 만드는 방식으로 접근합니다.

SoloQ의 핵심 구성

• 정규화된 회전 기저: 가중치와 활성값을 정규화된 회전 기저로 변환하고, 회전 이후의 주변 분포가 예측 가능하도록 설계합니다. 이를 통해 실행 데이터에 의존하지 않는 양자화 규칙을 만들 수 있습니다.
• 구조화된 K-RPBH 회전: 구조화된 K-RPBH 회전에 가벼운 재스케일링 보정을 더합니다. 이 조합은 값의 분포를 정리하면서 변환 자체의 추가 부담을 제한하는 방향입니다.
• 코드북과 NVFP4 지원: 회전 후 분포를 예측할 수 있기 때문에 분포에 맞춘 코드북과 하드웨어 네이티브 NVFP4를 모두 활용할 수 있습니다. 배포 환경에 따라 수치 표현을 선택할 여지가 생깁니다.
• 커밋 시점 KV 캐시 양자화: 블록 확산 모델에서는 현재 디노이징 중인 블록을 바로 양자화하지 않습니다. 블록이 커밋된 뒤 지속적으로 보관되는 KV 상태를 압축해, 현재 연산에 대한 간섭을 피하면서 캐시 메모리를 줄입니다.

결과와 영향

제공된 자료에 따르면 SoloQ는 LLaDA와 Dream 같은 전체 시퀀스 확산 모델, Fast-dLLM v2와 Nemotron-Labs-Diffusion 같은 블록 확산 모델에서 평가되었습니다. 4비트 양자화에서도 정확도를 유지했으며, 지식과 추론 중심 벤치마크에서 비교된 보정 기반 방법보다 좋은 결과를 보였다고 설명됩니다. NVFP4를 사용했을 때는 피크 메모리를 최대 2.61배 줄이고 엔드투엔드 추론을 최대 2.24배 빠르게 했다는 결과가 제시됩니다.

이 연구의 의미는 확산 언어 모델의 양자화를 파라미터 압축만의 문제로 보지 않는 데 있습니다. dLLM에서는 반복적인 디노이징 과정에서 활성값이 변하고, 블록 구조에서는 지속 상태가 메모리를 차지합니다. SoloQ는 회전과 재스케일링을 통한 수치 변환, KV 캐시 관리, 하드웨어 형식을 하나의 배포 관점에서 결합합니다. 그 결과 모델이나 추론 상태마다 별도의 보정 데이터를 준비해야 하는 부담을 줄이는 방향을 제시합니다.

다만 제공된 초록만으로는 모델별 세부 설정과 모든 하드웨어 조건에서의 성능을 확인할 수 없습니다. 따라서 실제 메모리 절감과 속도 향상은 구현과 환경에 따라 달라질 수 있습니다. 그럼에도 SoloQ는 확산 언어 모델을 위한 저비트 추론 최적화가 가중치 압축을 넘어 활성값 분포와 캐시 수명까지 함께 설계해야 한다는 점을 보여주는 접근입니다.

출처: Hugging Face Daily Papers