AlignSAE: Concept-Aligned Sparse Autoencoders (ArXiv 2026)

https://www.dropbox.com/scl/fi/85qugzzkromfqz828fq38/arxiv25_The_AlignSAE_Control_Paradigm.pdf?rlkey=wfntedkpu0862gb92jb31lmyc&dl=0

논문 개요

이 논문은 기존 Sparse Autoencoder가 학습한 feature를 사후에 해석하는 데 그치지 않고, 사전에 정의한 ontology의 개념을 특정 SAE latent slot에 직접 정렬하여 이를 읽고 조작할 수 있는 인터페이스로 만드는 방법입니다.

핵심 발상은 다음과 같습니다.

일반 SAE 학습을 LLM의 pre-training에 대응시키고, 개념 supervision을 이용한 추가 학습을 LLM의 post-training에 대응시킨다.

먼저 일반 SAE처럼 비지도 재구성 학습을 수행한 뒤, 각 관계 개념을 고정된 latent slot에 연결하는 supervised post-training을 수행합니다. 이를 통해 하나의 개념이 여러 feature로 흩어지는 fragmentation과, 하나의 feature가 여러 개념을 나타내는 polysemanticity를 줄이고자 합니다. 최종 목표는 특정 slot을 관찰하면 모델이 어떤 관계를 사용하고 있는지 알 수 있고, 그 slot의 decoder direction을 주입하면 모델 출력을 해당 관계로 바꿀 수 있게 하는 것입니다.


1. 기존 SAE의 문제점

일반적인 SAE는 LLM activation h를 고차원의 sparse latent z로 분해합니다.

z=E(h),h^=D(z)z=E(h),\qquad \hat h=D(z)

이론적으로는 각각의 SAE feature가 하나의 인간 해석 가능한 개념을 나타내기를 기대하지만, 실제 비지도 SAE에는 다음을 보장하는 목적함수가 없습니다.

  • BIRTH_CITY가 정확히 한 feature에만 위치하는 것
  • 그 feature가 다른 개념을 함께 표현하지 않는 것
  • 동일한 개념이 여러 feature로 분산되지 않는 것
  • 인간이 원하는 개념을 어느 feature에서 찾을 수 있는지 미리 아는 것

따라서 기존 SAE steering은 보통 contrast pair, top-activating examples, feature search 등의 사후 탐색을 필요로 합니다. ALIGNSAE는 이 문제를 feature discovery 문제가 아니라 supervised binding 문제로 바꿉니다. 즉, “어느 feature가 BIRTH_CITY인가?”를 찾는 대신 처음부터 “slot j는 BIRTH_CITY”가 되도록 학습합니다.


2. 전체 구조

2.1 Activation 추출

기반 모델은 12-layer, hidden dimension 768인 GPT-2 124M입니다. 먼저 GPT-2를 synthetic relation-completion 데이터로 fine-tuning하고, 질문 prompt의 마지막 token, 즉 answer generation 직전의 residual-stream activation을 추출합니다.

모든 12개 transformer layer에서 activation을 별도로 수집하여, 어느 layer가 개념 정렬과 intervention에 가장 적합한지도 분석합니다.

논문 본문에서는 mean pooling도 가능한 예시로 언급하지만, 부록에 명시된 실제 실험은 질문의 마지막 token activation을 사용합니다.


2.2 Concept slot과 free-feature bank

추출한 activation을 hd,d=768h\in \mathbb{R}^{d},\qquad d=768 라고 하면 encoder는 다음 sparse representation을 만듭니다.

z=ReLU(Weh+be)z=\operatorname{ReLU}(W_eh+b_e) , z=[zconcept;zrest]z=[z_{\text{concept}};z_{\text{rest}}]

여기서 zconcept|R|,zrestK.z_{\text{concept}}\in\mathbb{R}^{|R|},\qquad z_{\text{rest}}\in\mathbb{R}^{K}.

  • zconceptz_{\text{concept}}: ontology relation마다 하나씩 할당된 supervised concept slot
  • zrestz_{\text{rest}}: 나머지 언어적·개체적 정보를 보존하는 unsupervised free features

Decoder는 h^=Wdz+bd\hat h=W_dz+b_d 로 원래 activation을 복원합니다.

사용된 크기는 다음과 같습니다.

| 과제 | 관계 수 |R| | Free features K |
|———————|———-:|—————-:|
| 1-hop biography QA | 6 | 10,000 |
| 2-hop reasoning | 20 | 100,000 |

즉, 해석해야 할 소수의 slot만 ontology에 정렬하고, 나머지 대규모 capacity는 일반적인 SAE feature로 유지합니다.

이 분리가 중요한 이유는 concept slot에 모든 정보를 억지로 병목시키면 reconstruction과 언어 능력이 크게 손상될 수 있기 때문입니다. ALIGNSAE는 “작은 해석 가능 영역 + 큰 자유 영역”을 둔다는 점에서 일반적인 concept bottleneck보다 덜 제한적입니다.


3. Pre-train → Post-train 학습

3.1 Stage 1: 비지도 SAE 초기화

먼저 개념 label을 사용하지 않고 SAE가 activation을 안정적으로 재구성하도록 학습합니다.

부록의 실제 구현 명세에서는 다음과 같이 되어 있습니다.

  • Stage 1: reconstruction-only, 100 epochs
  • Stage 2: full supervision, 500 epochs

즉 Stage 1에서는 encoder와 decoder가 GPT-2 activation의 기본 dictionary를 형성합니다. 그 후 개념 slot을 특정 관계에 정렬합니다.

다만 본문에서는 Stage 1을 reconstruction과 sparsity로 학습한다고 설명하는 반면, 부록에서는 reconstruction-only라고 명시합니다. 재현 시에는 이 차이를 확인할 필요가 있습니다.


3.2 Stage 2: Concept-aligned post-training

Stage 2에서는 다음 목적함수를 공동 최적화합니다.

total=λreconrecon+λsparsesparse+λalignalign+λorthoortho+λvaluevalue.\mathcal L_{\text{total}} = \lambda_{\text{recon}}\mathcal L_{\text{recon}} +\lambda_{\text{sparse}}\mathcal L_{\text{sparse}} +\lambda_{\text{align}}\mathcal L_{\text{align}} +\lambda_{\text{ortho}}\mathcal L_{\text{ortho}} +\lambda_{\text{value}}\mathcal L_{\text{value}}.

각 loss의 역할이 서로 다릅니다.


4. 목적함수 상세

4.1 Reconstruction loss

recon=1di=1d(h^ihi)2.\mathcal L_{\text{recon}} = \frac{1}{d}\sum_{i=1}^{d} (\hat h_i-h_i)^2.

SAE를 통과한 표현이 원래 residual-stream activation을 보존하도록 합니다. Concept alignment만 강하게 적용하면 concept slot은 잘 분류하지만 전체 activation 구조를 잃을 수 있으므로 이를 방지합니다.


4.2 Sparsity loss

sparse=1BKb=1Bj=1K|zb,j|.\mathcal L_{\text{sparse}} = \frac{1}{B K} \sum_{b=1}^{B}\sum_{j=1}^{K} |z_{b,j}|.

L1 penalty는 주로 대규모 zrestz_{\text{rest}} free-feature bank에 적용됩니다. 각 example에서 소수 feature만 활성화되게 하여 일반 SAE의 sparse dictionary 성질을 유지합니다.


4.3 Alignment 또는 binding loss

각 relation r에 대응하는 concept slot을 정해 놓고,

p(r|x)=softmax(zconcept)p(r\mid x)=\operatorname{softmax}(z_{\text{concept}})

로 관계를 예측합니다.

align=CE(softmax(zconcept),yrel).\mathcal L_{\text{align}} = \operatorname{CE} \left( \operatorname{softmax}(z_{\text{concept}}),y_{\text{rel}} \right).

입력이 BIRTH_CITY 질문이면 BIRTH_CITY slot activation이 다른 모든 concept slot보다 커지도록 학습합니다.

이 loss의 역할은 다음 두 가지입니다.

  1. 하나의 관계가 여러 slot으로 흩어지는 것을 방지한다.
  2. 각 slot의 semantic identity를 사람이 미리 알 수 있게 한다.

즉 일반 SAE에서 필요한 feature identification 절차를 제거합니다.


4.4 Orthogonality/ decorrelation loss

이름은 orthogonality loss이지만, 실제로는 decoder weight vector의 기하학적 직교성을 직접 강제하는 것이 아니라 concept-slot activation과 free-feature activation 사이의 cross-covariance를 최소화합니다.

ortho=r=1|R|j=1K[1Bb=1B(zb,rzr)(zb,jzj)]2.\mathcal L_{\text{ortho}} = \sum_{r=1}^{|R|} \sum_{j=1}^{K} \left[ \frac{1}{B} \sum_{b=1}^{B} (z_{b,r}-\bar z_r) (z_{b,j}-\bar z_j) \right]^2.

목적은 관계 정보가 concept slot과 free feature에 중복 저장되는 것을 막는 것입니다.

예를 들어 UNIVERSITY 정보가 UNIVERSITY slot에도 존재하지만 free features에도 강하게 복제되어 있다면, UNIVERSITY slot을 바꾸더라도 downstream computation은 free features에서 원래 관계 정보를 복구할 수 있습니다. 이 경우 slot classification은 잘되지만 causal steering은 실패합니다.

따라서 이 loss는 readability보다 controllability를 위해 특히 중요합니다.


4.5 Value/sufficiency loss

Alignment loss만 사용하면 concept slot은 단순히 “이 질문은 UNIVERSITY 질문이다”라는 relation label만 표현할 수 있습니다. 그러나 steering을 위해서는 이 slot이 실제 answer generation과 연결된 task-relevant information을 가져야 합니다.

이를 위해 gold relation slot의 scalar activation을 relation-specific two-layer MLP에 넣어 정답의 첫 token을 예측합니다.

value=1Bb=1BCE(Vrb(zb,rb),tb).\mathcal L_{\text{value}} = \frac{1}{B} \sum_{b=1}^{B} \operatorname{CE} \left( V_{r_b}(z_{b,r_b}),t_b \right).

  • rbr_b: gold relation
  • zb,rbz_{b,r_b}: 해당 relation slot의 scalar activation
  • VrV_r: 관계별 auxiliary MLP
  • tbt_b: 정답의 첫 token

Value head는 학습 시에만 사용되며 실제 swap intervention에는 사용되지 않습니다.

이 목적함수는 단순한 concept detector와 ALIGNSAE의 차이를 보여 줍니다. Concept slot은 관계 identity만 나타내는 것이 아니라, answer prediction에 실질적으로 유용한 표현이 되도록 학습됩니다.


4.6 Loss weights와 SAE 최적화 설정

사용한 loss weight는 다음과 같습니다.

λrecon=1.0,λsparse=103,\lambda_{\text{recon}}=1.0,\qquad \lambda_{\text{sparse}}=10^{-3}, λalign=1.0,λortho=102,λvalue=0.5.\lambda_{\text{align}}=1.0,\qquad \lambda_{\text{ortho}}=10^{-2},\qquad \lambda_{\text{value}}=0.5.

SAE 학습 설정은 다음과 같습니다.

  • AdamW
  • learning rate 10310^{-3}, constant
  • weight decay 0
  • batch size 64
  • β=(0.9,0.999)\beta=(0.9,0.999)
  • gradient clipping 없음
  • Stage 1: 100 epochs
  • Stage 2: 500 epochs

저자들은 extensive hyperparameter search를 수행하지 않았으며, collapse와 dead feature를 피하기 위한 최소한의 sanity check만 했다고 밝힙니다.


5. Causal controllability: Concept swap

ALIGNSAE가 단순한 relation classifier가 아니라 causal interface인지 평가하기 위해 swap test를 사용합니다.

질문의 원래 관계가 rr^\star이고, 이를 다른 관계 j로 바꾸려 한다고 합시다. Concept slot j의 decoder direction은 vj=Wdecejv_j=W_{\text{dec}}e_j 입니다. 이 방향을 residual stream에 주입합니다.

h,t=h,t+αvj=h,t+Wdec(αej).h’_{\ell,t} = h_{\ell,t}+\alpha v_j = h_{\ell,t}+W_{\text{dec}}(\alpha e_j).

예를 들어 다음 질문이 있다고 하겠습니다.

When was Alice born?

원래 답은 Alice의 생년월일입니다. 여기서 UNIVERSITY slot direction을 주입했을 때 출력이 Alice의 대학으로 바뀌면 swap 성공으로 판정합니다.

successy^=g(Alice,UNIVERSITY).\text{success} \iff \hat y=g(\text{Alice},\text{UNIVERSITY}).

즉 단순히 “대학 종류의 아무 답”을 내는 것이 아니라 동일 인물의 정확한 대학 값을 생성해야 strict success입니다.

직관적으로 원래 residual stream과 free features에는 인물 identity가 남아 있고, concept direction을 추가하면 “어떤 속성을 검색할 것인가”가 바뀌는 구조입니다.


6. 실험 1: 1-hop biography factual recall

6.1 데이터셋

Synthetic biography 데이터는 1,000명의 가상 인물로 구성됩니다. 각 인물은 다음 6개 속성을 가집니다.

  1. Birth date
  2. Birth city
  3. University
  4. Major
  5. Employer
  6. Work city

각 인물마다 5개의 biography variant를 생성합니다.

질문 template은 관계별로 네 개이며,

  • T0–T1: 학습
  • T2–T3: unseen-template 평가

로 분리합니다.

예를 들어 UNIVERSITY 관계에서:

  • 학습: “Where did p go to college?”
  • 평가: “What is p’s alma mater?”

와 같이 lexical trigger가 달라집니다. 다만 평가 인물 자체는 새로운 인물이 아니라 학습에 등장한 동일 인물이며, OOD는 entity OOD가 아니라 template OOD입니다.

Base GPT-2는 biography 전체를 예측하는 memorization 학습과, question 부분을 masking하고 answer token에만 loss를 주는 QA 학습을 함께 수행합니다.


6.2 Layer-wise 결과

가장 중요한 결과는 Layer 6이 sweet spot이라는 것입니다.

MetricLayer 0Layer 6
Diagonal binding accuracy0.2381.000
Max swap success0.0400.847
Train slot accuracy0.2321.000
Unseen-template slot accuracy0.1650.912
Reconstruction MSE6.53×1056.53\times10^{-5}7.42×1027.42\times10^{-2}

Layer 6에서는 관계–slot confusion matrix가 거의 완벽한 diagonal 형태가 됩니다. 반면 Layer 0에서는 relation mass가 여러 slot으로 분산됩니다.

해석

  • 초기 layer: lexical·embedding 정보 중심이어서 추상적 relation type이 안정적으로 분리되지 않음
  • 중간 layer: 질문의 semantic relation이 충분히 형성되어 concept binding이 쉬움
  • 후기 layer: output 생성에 필요한 정보가 압축되고 특수화되어 안정적인 slot interface를 만들기 어려움

다만 Layer 6의 reconstruction MSE가 Layer 0보다 약 1,100배 크다는 점은 상당한 trade-off입니다. 저자들은 reconstruction fidelity보다 semantic controllability가 중간 layer에서 더 우수하다는 결과로 해석합니다.


6.3 Fragmentation과 concentration

저자들은 개념 c의 feature별 평균 activation을 정규화하여 distribution Bc,kB_{c,k}를 만들고 두 metric을 정의합니다.

Effective number of features

EffFeat(c)=exp(kBc,klogBc,k).\operatorname{EffFeat}(c) = \exp \left( -\sum_k B_{c,k}\log B_{c,k} \right).

  • 작을수록 한 개념이 적은 수의 feature에 집중
  • 1에 가까우면 사실상 한 feature가 개념을 담당

Top-1 concentration

Top1C(c)=maxkBc,k.\operatorname{Top1C}(c)=\max_k B_{c,k}.

  • 클수록 하나의 feature가 해당 개념의 activation mass를 지배

결과는 다음과 같습니다.

  • 기존 unsupervised SAE:
    • 개념당 수백∼수천 feature에 분산
    • Top1C가 거의 0
  • ALIGNSAE:
    • 중간 layer 이후 EffFeat가 약 1
    • Top1C가 거의 1

즉 supervised post-training이 many-to-many representation을 거의 one-to-one representation으로 바꿉니다.

이 결과는 강력하지만, 어느 정도는 alignment CE loss가 직접 최적화한 결과이기도 합니다. 따라서 핵심 검증은 단순 diagonalization이 아니라 이후의 causal swap 성능입니다.


6.4 Swap intervention 결과

Layer 5–8에서 moderate intervention이 가장 잘 동작합니다.

Layer 6의 결과는 다음과 같습니다.

α\alphaSwap success
1.00.54
2.00.85
10.00.23

즉 intervention을 강하게 할수록 항상 좋아지는 것이 아니라, α2\alpha\approx2에서 최적입니다.

예를 들어 BIRTH_DATE 질문에서 원래 출력이 날짜인 상태에서 UNIVERSITY slot을 증폭하면, 동일 인물의 대학 이름으로 출력이 바뀝니다. 이는 concept slot이 단순한 diagnostic readout이 아니라 downstream output에 인과적으로 연결되어 있음을 보여 줍니다.

왜 너무 큰 α\alpha에서 성능이 떨어지는가?

α\alpha가 지나치게 크면 target relation category는 강하게 주입되지만, 원래 residual stream에 있는 person-specific 정보와 language-generation 구조를 훼손할 수 있습니다. 따라서 “무슨 종류의 답을 해야 하는가”는 바뀌지만 “정확히 어느 entity를 답해야 하는가”는 틀릴 수 있습니다.


6.5 실패 사례 분석

Strict exact-match swap에 실패한 경우에도 target semantic category를 유지하는지 조사합니다.

Strength실패 중 target category 유지
α=2\alpha=275.3%
α=10\alpha=1083.0%

예를 들어 UNIVERSITY에서 MAJOR로 swap했을 때:

  • 정답: Physical Therapy
  • 생성: Geography

라면 exact entity는 틀렸지만, university 이름 대신 academic major를 생성했으므로 category swap은 성공한 셈입니다.

이 결과는 concept slot이 크게 두 수준의 제어를 수행한다는 것을 시사합니다.

  1. Relation/type control: 무엇에 관한 답을 생성할 것인가
  2. Entity/value retrieval: 그 사람에게 대응하는 정확한 값이 무엇인가

α\alpha는 첫 번째는 강화하지만 두 번째를 손상시킵니다.


7. Ablation study

7.1 Post-training loss ablation

설정Train slotOOD slotUnseen answerBest swap
Full joint1.0000.9120.8090.847
No alignment0.1670.1670.8090.053
No orthogonality0.7110.6890.8090.046
No value0.1770.1720.8090.720

Swap success는 각 설정에 가장 유리한 별도 α\alpha에서 측정했습니다.

  • Full: α=2\alpha=2
  • No alignment: α=0.1\alpha=0.1
  • No orthogonality: α=10\alpha=10
  • No value: α=100\alpha=100

Alignment loss 제거

Chance가 1/60.1671/6\approx0.167 인데, OOD slot accuracy가 정확히 0.167로 떨어집니다. Swap도 0.053이므로 개념별 고정 주소가 형성되지 않습니다.

Orthogonality loss 제거

OOD binding은 0.689로 어느 정도 유지되지만 swap은 0.046으로 붕괴합니다.

이 결과가 매우 중요합니다.

“어떤 slot이 relation을 분류하는가”와 “그 slot만 조작하여 모델 행동을 바꿀 수 있는가”는 서로 다른 조건이다.

Orthogonality가 없으면 free features에 원래 relation 정보가 누출되어 concept slot을 조작해도 전체 모델이 이를 따르지 않습니다.

Value loss 제거

Slot binding은 chance 수준이지만 α=100\alpha=100 이라는 매우 큰 intervention에서 0.720 swap success를 얻습니다. 즉 decoder direction 자체에 유용한 방향은 남아 있지만, 낮은 intervention strength에서 안정적으로 작동하는 정교한 interface가 되지는 못합니다.

Full model은 α=2\alpha=2에서 0.847이므로 효율적이고 안정적인 causal direction을 형성합니다.

Base answer accuracy가 동일한 이유

모든 ablation에서 train/unseen answer accuracy가 각각 1.000/0.809로 같습니다. 이는 base GPT-2가 고정되어 있고 SAE loss가 기본 QA 모델 자체를 fine-tuning하지 않기 때문입니다. Ablation은 base capability가 아니라 SAE interface의 binding과 intervention만 변화시킵니다.


7.2 Stage-1 pretraining ablation

Stage 1 없이 처음부터 full supervised objective를 최적화한 결과는 다음과 같습니다.

MetricNo Stage 1With Stage 1
Total loss3.34381.8594
Reconstruction loss0.04120.0256
Sparsity loss0.26820.0209
Orthogonality loss0.53300.3760
Train binding0.97001.0000
Test binding0.96990.9538
Free-feature independence score157.804.76

Stage 1의 핵심 효과는 binding 자체가 아닙니다. Test binding은 오히려 거의 동일하거나 약간 낮습니다. 대신 다음이 크게 개선됩니다.

  • reconstruction
  • sparsity
  • concept/free separation
  • free-feature redundancy

특히 free-feature off-diagonal covariance 기반 independence score가 157.80에서 4.76으로 감소합니다. 즉 pre-training은 개념 label을 맞추기 위한 단계라기보다 안정적인 sparse dictionary를 먼저 형성하기 위한 단계입니다.


8. 실험 2: 2-hop compositional reasoning

8.1 문제 설정

2-hop task에서는

e1r1e2r2e3e_1\xrightarrow{r_1}e_2\xrightarrow{r_2}e_3

를 계산해야 합니다.

질문 예시는 다음과 같습니다.

Who is the friend of the report of Avery?

여기서

  • r1=𝚛𝚎𝚙𝚘𝚛𝚝𝚜_𝚝𝚘r_1=\texttt{reports\_to}
  • r2=𝚏𝚛𝚒𝚎𝚗𝚍_𝚘𝚏r_2=\texttt{friend\_of}

이고 모델은 중간 entity와 최종 entity를 모두 e2; e3e_2;\ e_3 형태로 생성합니다.

Step-wise concept supervision은 다음과 같습니다.

  • e2e_2 생성 token: r1r_1 slot 활성화
  • e3e_3 생성 token: r2r_2 slot 활성화

따라서 최종 정답만 평가하는 것이 아니라 각 reasoning step에서 모델이 어떤 relation을 사용하는지를 직접 읽을 수 있습니다.

데이터 설정은 다음과 같습니다.

  • 20개 relation
  • 60개 entity
  • 총 8,000개 QA pair
  • 4,000 train / 4,000 validation
  • Layer 6 사용
  • K=100,000K=100{,}000 free features

8.2 Step-wise binding 결과

학습 후 두 단계 모두 거의 완벽한 diagonal confusion matrix를 보입니다.

  • Token 1, e2e_2 생성: r1r_1 slot과 100% binding
  • Token 2, e3e_3 생성: r2r_2 slot과 100% binding

즉 동일한 sequence 안에서도 현재 reasoning step에 따라 활성 concept slot이 바뀝니다. 이는 단순히 전체 질문에 포함된 두 관계를 정적으로 탐지하는 것이 아니라, 적어도 supervised readout 수준에서는 step-specific relation usage를 구분할 수 있음을 의미합니다.


8.3 2-hop swap 결과

비교 대상은 다음 세 가지입니다.

  1. Post-trained ALIGNSAE
  2. Unsupervised traditional SAE
  3. Supervised linear probe

Linear probe는 relation classification accuracy가 거의 100%이지만, classifier weight를 steering direction으로 사용했을 때의 causal 제어는 제한적입니다.

ALIGNSAE는 α50\alpha\approx50에서 traditional SAE보다 약 4배 높은 swap success를 보이고, linear probe보다도 높은 최고 성능을 냅니다.

다만 Figure 8의 세로축은 percentage 단위로 약 1–8% 범위입니다. 따라서 상대적으로 4배 향상되었지만 절대적인 2-hop exact-swap 성공률은 최고 약 7–8% 수준으로 보입니다. 이는 1-hop의 84.7%와 큰 차이가 있습니다.

따라서 2-hop 결과는 다음처럼 해석하는 것이 적절합니다.

  • Step-wise concept localization: 매우 강함
  • Exact causal replacement of the final entity: 아직 제한적
  • Traditional SAE 대비 상대적 개선: 분명함
  • 실용적인 multi-hop knowledge editing 수준: 아직 아님

9. Grokking-like dynamics 분석

저자들은 2-hop base model에서 다음 현상을 관찰합니다.

  • Training loss는 계속 감소
  • Validation accuracy는 오랫동안 낮게 유지
  • 이후 validation accuracy가 급격히 상승

이를 전형적인 grokking과 구분하여 grokking-like generalization transition이라고 부릅니다.

ALIGNSAE를 각 학습 단계의 representation에 적용하면 다음 변화를 관찰합니다.

시점Token 1 bindingToken 2 binding
Epoch 1038.4%69.5%
Epoch 40100.0%100.0%

Epoch 10에는 off-diagonal activation이 많아 relation evidence가 여러 slot에 얽혀 있습니다. Epoch 40에는 거의 완벽한 diagonal structure로 바뀝니다.

더 흥미로운 점은 Top-1 binding accuracy가 validation accuracy보다 먼저 상승한다는 것입니다.

저자들은 이를 다음과 같이 해석합니다.

  1. 모델 내부에서 먼저 relation별 표현 구조가 정돈된다.
  2. 각 reasoning step의 concept가 고정된 slot에 안정적으로 대응한다.
  3. 그 이후에야 이 구조를 OOD example에서 일관되게 사용하여 validation accuracy가 상승한다.

논문은 이를 다음 표현으로 요약합니다.

  • Knowing: 내부 concept structure가 먼저 형성됨
  • Showing: 이후 외부 generalization 성능으로 나타남

다만 이는 상관관계 기반 mechanistic account입니다. Concept binding의 선행이 validation jump를 인과적으로 발생시켰다고 직접 증명한 것은 아닙니다.


10. 논문의 주요 기여

10.1 SAE를 discovery 도구에서 addressable interface로 전환

일반 SAE: activationunknown sparse featuresfeature interpretation/search\text{activation}\rightarrow \text{unknown sparse features} \rightarrow \text{feature interpretation/search}

ALIGNSAE: activationknown concept slots+free features\text{activation}\rightarrow \text{known concept slots} + \text{free features}

따라서 ontology가 정해져 있는 relation extraction, safety taxonomy, style control 등의 환경에서는 어떤 slot을 조작해야 하는지 미리 알 수 있습니다.

10.2 Binding과 causal controllability를 구분

Linear probe나 alignment accuracy가 높다고 해당 direction이 causal control에 유용한 것은 아닙니다. Orthogonality ablation은 이를 명확하게 보여 줍니다.

  • no_ortho: OOD binding 68.9%
  • swap success: 4.6%

즉 probe accuracy만으로 feature의 causal usefulness를 판단해서는 안 된다는 결과입니다.

10.3 Step-wise concept interface

2-hop reasoning에서 각 decoding step마다 다른 relation slot이 활성화되는 구조를 제시합니다. 이는 향후 다음 방향으로 확장할 수 있습니다.

token-wise concept slotsreasoning trace verificationstep-specific intervention.\text{token-wise concept slots} \rightarrow \text{reasoning trace verification} \rightarrow \text{step-specific intervention}.

10.4 Grokking 내부 구조에 대한 새로운 관찰

단순 activation 통계가 아니라 human-defined concept와 정렬된 slot을 이용해, generalization transition 전에 relation representation이 정돈되는 현상을 관찰합니다.


11. 한계와 비판적 해석

11.1 Concept discovery가 아니라 supervised concept allocation

ALIGNSAE는 unknown concept를 발견하는 방법이 아닙니다. 다음 정보가 필요합니다.

  • 사전에 정의된 ontology
  • 각 sample의 relation label
  • concept별 dedicated slot
  • answer supervision

따라서 “SAE가 스스로 monosemantic feature를 발견했다”기보다는 supervised concept bottleneck을 SAE 내부에 삽입했다고 보는 것이 더 정확합니다.


11.2 실험 환경이 매우 통제됨

실험은 다음 범위에 한정됩니다.

  • GPT-2 124M
  • synthetic biography
  • 6개 relation
  • synthetic 2-hop reasoning
  • 최대 20개 relation
  • template-based language

저자들도 실제 대규모 LLM, 자연 발생 factual knowledge, 대규모 ontology로의 확장은 검증되지 않았다고 인정합니다.


11.3 1-hop OOD의 범위가 제한적

Unseen evaluation은 동일한 인물과 동일한 relation에 대해 질문 template만 바꿉니다.

따라서 91.2%는 다음을 의미합니다.

relation slot이 새로운 표면 표현에서도 동일한 relation을 인식한다.

그러나 다음을 입증하지는 않습니다.

  • unseen entity generalization
  • unseen relation generalization
  • ontology expansion
  • domain transfer
  • 자연스러운 자유 형식 질문

실제로 Claude가 생성한 더 간접적인 질문은 GPT-2가 잘 처리하지 못해 본 실험에서 제외되었습니다.


11.4 2-hop swap의 절대 성능

논문은 traditional SAE보다 “4배 높다”고 강조하지만 Figure 8상 절대 성공률은 약 7–8%입니다. 따라서 step-wise binding 100%와 exact output control 사이에는 큰 간극이 있습니다.

이는 다음을 의미합니다.

concept decodability⇏reliable causal compositional control.\text{concept decodability} \not\Rightarrow \text{reliable causal compositional control}.


11.5 “Lightweight”라는 표현의 한계

2-hop에서는 K=100,000K=100{,}000, d=768입니다. Encoder와 decoder가 모두 dense matrix라면 두 행렬만으로도 약 1.5억 개 이상의 parameter가 필요합니다. 이는 GPT-2 124M 자체보다 클 수 있습니다.

Base LLM을 고정한다는 점에서는 학습 대상이 분리되어 있지만, parameter 수와 메모리 관점에서 2-hop ALIGNSAE를 단순히 lightweight라고 보기는 어렵습니다.


11.6 Reconstruction과 controllability의 trade-off

Layer 6의 reconstruction MSE가 Layer 0보다 약 1,100배 높으면서 controllability는 크게 개선됩니다.

이는 ALIGNSAE가 다음 두 목표를 동시에 완전히 달성한 것은 아니라는 뜻입니다.

high-fidelity activation decompositionvs.clean concept-addressability.\text{high-fidelity activation decomposition} \quad\text{vs.}\quad \text{clean concept-addressability}.

현재 intervention은 원래 residual에 direction을 추가하므로 reconstruction error가 직접 base prediction을 항상 손상시키지는 않지만, SAE reconstruction을 원래 activation의 완전한 대체물로 사용하려면 문제가 될 수 있습니다.


11.7 일부 구현 설명의 불일치

재현 관점에서 다음 차이가 있습니다.

  • 본문: Stage 1은 reconstruction + sparsity
    부록: reconstruction-only
  • 본문 swap: 원래 h에 target decoder direction을 추가
    부록 metric 설명: original slot을 0으로 만들고 target slot을 α\alpha로 설정
  • Ablation 본문: full Test-Unseen binding 0.891
    Table 6: 0.912
  • 부록은 “six components”라고 표현하지만 weighted total equation에는 다섯 loss가 있으며, free-slot independence loss는 optional diagnostic으로 제시됨

핵심 결론을 바꿀 정도는 아니지만, 정확한 재현에는 code 확인이 필요합니다.


12. 종합 평가

이 논문의 가장 중요한 메시지는 다음입니다.

SAE feature가 자동으로 인간 개념과 일치하기를 기다리지 말고, 일부 latent capacity를 ontology에 명시적으로 정렬하여 검증 가능하고 조작 가능한 인터페이스로 만들자.

방법론적으로는 다음 세 요소의 결합입니다.

Unsupervised SAE initialization+supervised concept binding+concept/free decorrelation\boxed{ \text{Unsupervised SAE initialization} + \text{supervised concept binding} + \text{concept/free decorrelation} }

실험적으로 가장 설득력 있는 결과는 1-hop Layer 6에서의 다음 조합입니다.

  • OOD slot binding: 91.2%
  • one-to-one diagonal binding: 100%
  • exact swap success: 84.7%
  • moderate intervention: α=2\alpha=2

반면 2-hop에서는 concept binding은 거의 완벽하지만 exact causal swap은 낮습니다. 따라서 ALIGNSAE는 현재 개념 localization과 단순 causal control에는 강하지만, 복잡한 compositional reasoning을 안정적으로 다시 쓰는 수준에는 아직 도달하지 못한 방법으로 평가할 수 있습니다.

댓글

답글 남기기

이메일 주소는 공개되지 않습니다. 필수 필드는 *로 표시됩니다