대규모 모델을 위한 벡터 양자화 기반 잠재 개념 탐색

초록

본 논문은 VQ‑VAE 구조를 활용해 토큰 표현을 이산 코드북으로 매핑하는 VQLC(Vector Quantized Latent Concepts) 방법을 제안한다. 기존 계층적 군집은 메모리·시간 복잡도가 높고 K‑Means는 의미가 얕은 군집을 만들기 쉬운 문제를 해결하고자, 학습 가능한 코드북을 통해 잠재 개념을 자동으로 형성한다. 실험에서는 세 가지 텍스트 분류 데이터셋과 두 종류의 사전학습 모델(BERT, RoBERTa) 및 LLM(Llama‑2, Qwen2.5)에서 VQLC가 계층적 군집과 비슷한 설명 품질을 유지하면서 메모리 사용량을 크게 절감함을 보였다.

상세 분석

VQLC는 VQ‑VAE의 세 가지 핵심 모듈(인코더, 벡터 양자화기, 디코더)을 텍스트 토큰 표현에 맞게 재구성한다. 인코더는 원본 컨텍스트 표현 hℓ(wᵢ)를 선형 변환 후 레이어 정규화하고, 학습 가능한 α(0 ~ 0.5)와 잔차 연결을 통해 원본 정보를 보존하면서 양자화에 적합한 공간으로 조정한다. 이는 기존 VQ‑VAE에서 발생할 수 있는 표현 왜곡을 최소화하고, 코드북 벡터와의 거리 계산이 안정적이도록 만든다.

벡터 양자화 단계에서는 K = 400개의 코드북을 EMA(Exponential Moving Average) 방식으로 업데이트한다. 초기화는 대규모 토큰 집합에 대해 K‑Means를 수행해 고밀도 영역을 잡아내며, 이는 무작위 초기화보다 빠른 수렴과 더 균등한 코드 사용을 보장한다. 양자화는 코사인 거리 기반으로 수행되며, 학습 시 top‑k(=5) 후보를 소프트맥스 온도 τ = 1.0으로 샘플링해 코드북 붕괴를 방지한다. 추론 단계에서는 가장 가까운 코드북을 직접 할당해 ‘잠재 개념 벡터’를 얻는다.

디코더는 양자화된 코드북 벡터를 다시 원본 차원(d = 768)으로 복원한다. 저차원(d′) 선형 다운프로젝션 → 4‑layer Transformer 인코더 → 선형 업프로젝션 순서로 구성돼, 복원 손실 L_rec과 커밋 손실 L_commit을 동시에 최소화한다. L_commit은 stop‑gradient를 이용해 인코더 출력이 코드북에 과도히 멀어지는 것을 억제한다. β = 0.25라는 가중치는 실험을 통해 최적화되었으며, 민감도 분석 결과 β 변화가 설명 품질에 큰 영향을 주지 않음을 확인했다.

실험 설계는 세 가지 텍스트 분류 태스크(감성, 독성, 뉴스 주제)와 두 종류의 인코더‑디코더 모델(BERT‑base, RoBERTa‑base) 및 두 LLM(Llama‑2‑7B‑chat, Qwen2.5‑3B)에서 수행됐다. 비교 대상은 기존 LA‑COAT(계층적 군집)과 단순 K‑Means이다. 스케일러빌리티 평가는 토큰 수 10k300k, 차원 2048에 대해 피크 메모리 사용량을 측정했으며, VQLC는 GPU 메모리 23배, CPU 메모리 4~5배 절감했다. 신뢰성(faithfulness) 평가는 가장 중요한 잠재 개념을 제거했을 때 분류 정확도가 얼마나 감소하는지를 측정했으며, VQLC는 LA‑COAT에 근접하거나 더 높은 감소율을 보였다. LLM 기반 평가에서는 GPT‑4가 생성한 설명을 인간 평가자와 비교했을 때, VQLC가 제공한 개념이 더 직관적이고 일관적이라는 결과가 나왔다.

핵심 기여는 다음과 같다. 1) VQ‑VAE를 텍스트 토큰에 적용해 학습 가능한 이산 코드북을 구축, 이를 잠재 개념으로 활용. 2) EMA 기반 코드북 업데이트와 온도 샘플링을 도입해 코드북 붕괴와 편향을 완화. 3) 기존 군집 기반 방법 대비 메모리·시간 효율성을 크게 향상시키면서도 설명 품질을 유지. 4) 다양한 모델·데이터셋에 대한 포괄적 실험을 통해 방법론의 일반성을 검증.