Vector Quantized Latent Concepts: A Scalable Alternative to Clustering-Based Concept Discovery

기술 요약: 벡터 양자화 잠재 개념 (Vector Quantized Latent Concepts, VQLC)

문제 정의

대규모 언어 모델(LLM)은 내부 은닉 상태(hidden states)에 풍부한 의미론적 정보를 인코딩하지만, 이러한 내부 표현을 해석하는 것은 여전히 어려운 과제입니다. 토큰 수준의 기여도 산출 방식(예: Integrated Gradients, SHAP)은 중요한 입력 토큰을 식별할 수는 있지만, 모델에 인코딩된 고차원적인 의미 구조에 대한 통찰을 제공하는 데는 한계가 있습니다. 잠재 개념 발견(Latent concept discovery)은 문맥적 표현들을 이산적인(discrete) 해석 가능한 개념들로 그룹화함으로써 이 문제를 해결하고자 합니다. 그러나 기존의 접근 방식들은 **확장성(scalability)**과 개념의 품질(concept quality) 사이에서 근본적인 트레이드오프에 직면해 있습니다:

• 계층적 클러스터링 (예: LACOAT): 일관된 개념을 생성하지만, 이차 메모리 비용($O(N^2)$) 문제로 인해 소규모 데이터셋에만 적용이 제한됩니다.
• K-Means 클러스터링: 효율적으로 확장되지만, 의미론적 일관성이 떨어지는 경우가 많으며 확장성과 충실도(faithfulness) 사이의 트레이드오프가 약합니다.

본 논문은 대규모 LLM을 대상으로 할 때, 중심점 기반 방식(centroid-based methods)의 확장성을 유지하면서도 계층적 클러스트링의 의미론적 일관성을 달取得할 수 있는 방법론이 필요하다고 주장합니다.

방법론: 벡터 양자화 잠재 개념 (VQLC)

저자들은 연속적인 은닉 상태를 유한한 집합의 학습된 이산 코드(discrete codes)로 매핑하기 위해 **벡터 양자화(Vector Quantization, VQ)**를 응용한 이산 개념 학습 프레임워크인 VQLC를 제안합니다. 이 프레임워크는 특정 레이어에서 추출된 고정된(frozen) 은한 상태를 기반으로 작동합니다.

아키텍처 구성 요소:

1. 인코더 (Encoder): 경량 잔차 다층 퍼셉트론(residual MLP)이 문맥적 토큰 표현($h^{(\ell)}_{w_i}$)을 코드 공간으로 매핑합니다. 이는 직접적인 선형 투영(linear projection)과 비선형 보정 브랜치를 결합하여, 원래의 모델 공간을 보존하는 동시에 안정적인 이산 할당을 가능하게 합니다.
2. 벡터 양자화기 (Vector Quantizer): 이 모듈은 학습 가능한 코드북 $E = \{e_1, \dots, e_K\}$를 유지합니다. 각 인코딩된 표현은 제곱 유클리드 거리를 기준으로 가장 가까운 코드북 벡터에 할당됩니다.
   • 초기화: 코드북은 표현 공간의 다양한 커버리지를 보장하기 위해 중복이 제거된 훈련 표현들에 대해 "가장 먼 첫 번째(farthest-first)" 순회 방식을 사용하여 초기화됩니다.
   • 학습: 코드북은 직접적인 그래디언트 단계 대신, 할당된 인코더 출력의 지수 이동 평균(EMA)을 사용하여 업데이트됨으로써 안정성을 확보합니다. "사멸 코드 복구(dead code recovery)" 메커니즘은 사용되지 않는 코드를 할당 오차가 가장 큰 인코더 출력을 사용하여 재초기화합니다.
3. 디코더 (Decoder): 잔차 디코더는 양자화된 벡터로부터 원래의 은닉 상태를 재구성합니다. 훈련 목적 함수는 재구성 손실($L_{rec}$)과 인코더 출력이 할당된 코드북 벡터와 가깝게 유지하도록 하는 커밋 손실($L_{commit}$)을 최소화합니다.

개념 구축 (Concept Construction):
학습 후 모델은 고정됩니다. 개념 구축 패스(pass)는 동일한 코드 인덱스 $k$에 할당된 모든 인코딩된 토큰 표현들을 집계하여 개념 벡터 $v_k$를 형성합니다. 잠재 개념 $c_k$는 쌍 $(v_k, \text{할당된 토큰})$으로 정의됩니다. 추론 시, 새로운 토큰 표현은 학습된 코드북에 할당되며, 그에 대응하는 개념들이 해당 표현에 인코딩된 의미 정보를 설명합니다.

주요 기여

• 확장 가능한 이산 학습: VQLC는 계층적 클러스터링의 이차 메모리 의존성을 피하는 벡터 양자화 프레임워크를 도입하여 잠재 개념 발견을 위한 확장성을 제공합니다.
• 최적화된 트레이드오프: 이 방법은 데이터셋 크기가 증가하더라도 거의 일정한 메모리 사용량을 유지하며, 충실도(개념 품질)와 확장성 사이의 유리한 균형을 달성합니다.
• 포괄적인 평가: 본 논문은 다음을 포함한 12개의 데이터셋-모델 설정에 걸쳐 다각적인 평가를 제공합니다:
  • 확장성: 실행 시간 및 메모리 분석.
  • 충실도: 직교 투영(orthogonal projection)을 통해 개념 방향을 제거(ablation)한 후 모델의 신뢰도 및 예측 레이블의 변화로 측정.
  • LLM-as-a-Judge: 여러 LLM(Claude, Gemini, DeepSeek)을 사용하여 개념이 모델 예측과 얼마나 정렬되어 있는지 순위를 매기는 평가.
  • 정성적 분석: 개념 워드 클라우드 시각화 및 사례 연구.
  • SAE와의 비교: 희소 오토인코더(Sparse Autoencoders, SAEs)와의 대비를 통해 특징 수준(feature-level)과 개념 수준(concept-level) 설명의 차이를 강조.

결과

• 확장성: Qwen 모델을 사용한 ERASER 데이터셋에서 VQLC는 토큰 수에 따라 선형적으로 확장됩니다. K-Means의 메모리 사용량이 1.8GB에서 17.5GB로 증가하고 LACOAT가 200k 토큰에서 실패(약 229GB 필요)한 반면, VQLC는 500k 토큰에서도 1.94GB에서 2.54GB로 증가하며 거의 일정하게 유지되었습니다.
• 충실도: VQLC는 12개 데이터셋-모델 설정 중 6개에서 가장 높은 평균 신뢰도 변화를 달성했습니다. 특히 디코더 전용 모델(Llama-2, Qwen)에서 가장 뚜렷한 성능 향상을 보였으며, 6개 디코더 설정 중 5개에서 베이스라인을 능가했습니다. 인코더 기반 모델(RoBERTa)에서는 K-Means 및 LACOAT와 경쟁력 있는 수준을 유지했습니다.
• LLM 평가: LLM-as-a-judge 순위에서 VQLC는 12개 설정 중 8개에서 가장 좋은 평균 순위를 기록했으며, 이는 VQLC의 개념이 K-Means나 LACOAT보다 작업 정렬(task-aligned)도가 높음을 시사합니다.
• SAE와의 비교: 낮은 활성화율을 가진 다수의 희소 특징을 생성하는 SAE와 달리, VCLC는 훨씬 더 큰 비율의 테스트 예시에서 활성화되는 더 작은 규모의 개념 집합(예: ~400개)을 생성합니다. Qwen 모델 설정에서 VCLC 개념이 예측을 위한 작업 관련 정보를 더 많이 인코딩하는 것으로 나타났습니다.
• 정성적 결과: 시각화를 통해 VCLC 개념이 의미론적으로 일관됨을 보여주었습니다 (예: 스포츠 분야에서 "야구 선수"와 "축구"의 뚜렷한 클러스터 구분, 또는 과학/기술 분야에서 "임상 약물"과 "네트워크 인프라"의 구분). 사례 연구에서 VCLC는 K-Means나 LACOAT의 광범위한 클러스터에 비해 모델 예측에 대해 더 정밀한 설명을 제공했습니다.

의의 및 주장

본 논문은 VCLC가 LLM을 위한 클러스터링 기반 개념 발견의 실용적이고 확장 가능한 대안을 제공한다고 주장합니다. VCLC는 개념의 의미론적 일관성을 희생하지 않으면서 계층적 클러스터링의 메모리 병목 현상을 성공적으로 해결했습니다. 저자들은 VCLC가 특히 디코더 전용 모델에 효과적이며, 기존 베이스라인보다 작업 관련 개념 방향을 더 효과적으로 식별한다고 강조합니다.

이 연구는 VCLC를 SAE와 같은 기계론적 해석 도구를 대체하는 것이 아니라, 모노세맨틱(monosemantic) 특징이 아닌 **작업 수준의 잠재 개념(task-level latent concepts)**으로 표현을 조직하는 보완적 접근 방식으로 포지셔닝합니다. 저자들은 프레임워크가 하이퍼파라미터(코드북 크기, 커밋 가중치)에 민감하다는 점과 현재 분류 작업에 집중되어 있다는 점을 한계로 언급하며, 이 접근 방식을 생성적인 토큰 단위 추론으로 확장하는 것이 향후 연구 방향임을 시사했습니다.