TL;DR
HyperSAE는 LLM 활성화에서 계층적 개념을 추출하기 위해 전방 전달과 가중치 규제를 분리하는 아키텍처를 채택했습니다. 전방은 bfloat16 기반의 빠른 Euclidean 경로에서 희소 인코더·디코더로 동작하고, 가중치 공간은 Poincaré 사영과 비대칭 포함성 손실로 쌍곡선 규제를 적용해 구조적 정보를 보존합니다. Gemma-2-2B Layer 13 실험에서는 동일 희소성에서 재구성 MSE 9.8% 감소와 CE Recovery +3.4%를 보고해 유효성을 입증했습니다. PyPI 패키지와 TransformerLens 후크를 제공해 실험 재현과 개입 연구에 바로 사용할 수 있습니다.
섹션별 상세
- 동일 희소성 조건에서 HyperSAE는 FlatSAE 대비 재구성 MSE를 9.8% 낮추었다. — 벤치마크 표의 Pareto Reconstruction & Sparsity Performance 섹션에서 'Key Takeaway'로 MSE 9.8% 감소를 보고함.
- Cross-Entropy Loss Recovery에서 HyperSAE는 동일 희소성 기준으로 +3.4% 향상을 기록했다. — 동일 테이블의 CE Loss Recovery 열에서 HyperSAE가 FlatSAE보다 +3.4% 우수한 수치를 보임.
- Gemma-2-2B Layer 13의 GPQA Diamond 단일-토큰 치환 테스트에서 HyperSAE는 100% 보전 능력을 보고했다. — Empirical Benchmark Results 표의 GPQA 행에서 HyperSAE와 FlatSAE 비교 수치가 제시되어 있음.
import torch
from hypersae import HyperSAE, CoActivationQueue, TriPartiteLoss, HyperSAETrainer
device = "cuda" if torch.cuda.is_available() else "cpu"
# 1. Instantiate HyperSAE model, CoActivationQueue, and TriPartiteLoss
sae = HyperSAE(d_model=2304, dict_size=16384).to(device)
queue = CoActivationQueue(dict_size=16384).to(device)
loss_fn = TriPartiteLoss(l1_coeff=0.005, entail_coeff=0.01)
# 2. Instantiate HyperSAETrainer
trainer = HyperSAETrainer(model=sae, queue=queue, loss_fn=loss_fn, lr=1e-3)
# 3. Train step on residual stream activation batch
x = torch.randn(64, 2304, device=device)
metrics = trainer.train_step(x)
print(f"Total Loss: {metrics['loss_total']:.4f}")
print(f"Reconstruction MSE: {metrics['loss_recon']:.4f}")
print(f"Poincaré Entailment Penalty: {metrics['loss_entail']:.4f}")이 코드 예시는 PyPI로 설치한 HyperSAE를 GPU에서 바로 초기화해 잔차 스트림 활성화 한 배치에 대해 한 스텝 학습을 수행하는 흐름을 보여줍니다. 모델과 CoActivationQueue를 생성한 뒤 TriPartiteLoss로 MSE·L1·포앙카레 포함성 항을 결합하고 HyperSAETrainer로 업데이트를 실행합니다. 출력은 전체 손실과 재구성 MSE, 포앙카레 포함성 패널티를 로그해 학습 진척을 확인하는 용도로 설계되어 있습니다.
용어 해설
- 쌍곡선 기하학(Hyperbolic Geometry)
- — 쌍곡선 기하학은 유클리드와 달리 음의 곡률을 가지는 기하학으로, 계층적 구조를 작은 반경에 압축해 표현하는 성질이 있습니다. 자연어의 개념 계층이나 온톨로지처럼 트리형 구조를 고밀도로 임베딩할 때 거리와 포함 관계를 더 잘 보존하는 점이 장점입니다. 기계 학습에서는 Poincaré ball 같은 표현을 통해 파라미터 공간에서 계층적 관계를 강제하는 용도로 주로 사용됩니다.
- 포앙카레 볼(Poincaré Ball)
- — 포앙카레 볼은 쌍곡선 공간을 단위구 내부에 매핑한 모델로, 각 점의 반경과 방향을 통해 계층적 포함관계를 표현할 수 있습니다. 가중치나 사전(dictionary) 좌표를 이 구면 내부로 사영하면 거리 기반의 포함성(entailment) 손실을 설계할 수 있습니다. HyperSAE는 이 사영을 가중치 정규화 경로에서 사용해 계층 관계를 학습합니다.
- 희소 오토인코더(Sparse Autoencoder)
- — 희소 오토인코더는 인코더·디코더 구조에서 출력 활성화의 희소성을 규제해 입력의 주요 특성만 복원하도록 유도하는 모델입니다. L1 규제나 L0 기준을 통해 활성화 수를 제어하며, 고차원 잔차 스트림에서 의미 있는 피쳐를 분리해 내는 용도로 쓰입니다. HyperSAE는 이 구조에 쌍곡선 기반의 가중치 규제를 결합해 계층적 개념을 사전(dictionary) 파라미터에 내장합니다.
- 잔차 스트림 활성화(Residual Stream Activations)
- — 잔차 스트림 활성화는 Transformer 계열 LLM의 레이어 내부에서 전달되는 중간 표현으로, 각 토큰의 문맥 정보를 담고 있습니다. 메커니즘 해석에서는 이 활성화를 대상으로 피처 딕셔너리나 인코더를 학습해 개념 단위를 추출합니다. HyperSAE 실험은 Gemma-2-2B Layer 13의 잔차 활성화(d=2304)를 벤치마크 입력으로 사용했습니다.
- 공동활성화 큐(CoActivation Queue)
- — 공동활성화 큐는 대규모 딕셔너리에서 특징들 간 동시 등장 통계를 메모리 효율적으로 추적하는 구조입니다. HyperSAE 구현은 O(M^2) 메모리 증식을 피하면서 GPU 비동기 큐로 피처 공출현을 집계해 Poincaré 기반 포함성 손실을 계산합니다. 이로 인해 대형 딕셔너리(M=16384)에서도 확장성이 확보됩니다.
기술
- PyTorch
- TransformerLens
- Poincaré manifold
- TriPartiteLoss
- CoActivationQueue
- AdamW
- RiemannianAdam
- Gemma-2-2B
- NVIDIA L4
활용 사례
- LLM 내부 활성화에서 계층적 개념 온톨로지를 추출해 메커니즘 해석 연구에 활용할 수 있습니다. HyperSAE는 대규모 딕셔너리와 비동기 공출현 큐로 확장성을 확보했기 때문에 레이어 차원의 개념 발견 실험에 적합합니다. 추출한 개념을 후크로 연결하면 모델 거동을 국소적으로 스티어링하거나 원인분석에 활용할 수 있습니다.
- 희소 표현을 유지하면서 재구성 품질을 개선하는 목적의 압축·디버깅 파이프라인에 적용할 수 있습니다. 재구성 MSE 개선과 CE Loss 회복 증가는 특정 downstream 유지능력 보전 측면에서 이득을 제공하며, 이를 통해 토큰 교체나 국부적 수정에 대한 민감도를 평가할 수 있습니다. 대형 딕셔너리 운용을 하는 연구용 워크로드에 적합합니다.
- 가중치 공간에서 포함성 관계를 학습해 개념 간 계층적 연결을 분석하는 응용에 유용합니다. 포앙카레 기반 엔테일먼트 손실은 개념의 포함성·상위-하위 관계를 파라미터 수준에서 측정 가능한 형태로 만듭니다. 이는 개념 계층을 시각화하거나 자동화된 온톨로지 구성 실험으로 확장될 수 있습니다.
AI 요약 · 북마크 · 개인 피드 설정 — 무료
출처 · 인용 안내
인용 시 "요약 출처: AI Trends (aitrends.kr)"를 표기하고, 사실 확인은 원문 보기 기준으로 진행해 주세요. 자세한 기준은 운영 정책을 참고해 주세요.
