TopK SAE 어트리뷰션 그래프는 마진 재정렬로 정확해진다

무엇인가

이 논문은 SAE(희소 오토인코더) 피처 위에 만드는 어트리뷰션 그래프가 활성 피처만을 노드로 삼는 관행을 문제 삼는다. TopK SAE는 토큰마다 가장 큰 프리액티베이션 k개만 남기므로, 어떤 피처가 활성인지는 컷오프와의 거리인 마진이 결정한다. 기존 그래프는 클린 입력에서 활성인 피처를 고정하고 그 안에서 1차 변화를 예측하는데, 저자는 하나의 피처를 0으로 절제하는 것만으로도 다음 층 피처가 컷오프를 넘나들 수 있어 예측 대상 자체가 바뀐다고 지적한다. 그래서 절제가 다음 층 피처에 무엇을 하는지와 훈련 시드가 바뀌어도 그래프의 연산자가 얼마나 남는지를 묻는다.

어떻게 동작하나

제안하는 마진 그래프는 상류 피처 i의 코드를 α배로 쓰는 개입을 v=(α−1)u_i d_i로 두고, 하류 프리액티베이션 변화를 Δp̂=EJv로 계산한다. 고정 지원 그래프는 이 변화를 기존 활성 지원에 사영한 z+PΔp̂만 예측하지만, 마진 그래프는 p+Δp̂에 TopK 규칙을 다시 적용해 순위를 재정렬하고 어떤 피처가 들어오고 나가는지까지 예측한다. 비용은 상류 피처당 한 번의 야코비안-벡터 곱으로 같고, 기존 그래프가 활성 행만 쓰는 대신 마진 그래프는 하류 행 전체를 유지한다. 블록이 아핀인 구간에서는 마진 그래프가 정확한 코드를 준다는 정리 1을 제시하고, 그때 고정 지원 그래프의 오차는 들어오는 집합의 응답 에너지 비중의 제곱근 이상이라고 설명한다. 아핀에서 벗어난 구간에서는 각 쓰기 경로의 립시츠 상수를 세 점의 야코비안으로 추정해 1차 오차 밴드를 만들고 예측된 소속을 사후 감사한다.

무엇과 다른가

실험은 Pythia 모델의 공개 TopK SAE 다섯 쌍과 pythia-70m 3·4층에 직접 훈련한 testbed 셀 F(8런), 그리고 ReLU 쌍 하나를 쓴다. 휴지 상태에서 가장 가까운 비활성 마진의 중앙값은 -0.018~-0.008로 작았고, 단일 피처를 0으로 절제하면 하류 피처가 컷오프를 넘는 비율이 다섯 공개 쌍에서 0.74~0.89, 셀 F에서 0.92(런별 0.91~0.92)였다. 쓰기 크기가 0.9면 0.19~0.25로 줄고 4배면 0.95~0.99로 다시 늘었다. 교차 피처는 응답 에너지의 평균 0.73~0.84(셀 F 0.79)를 차지했고, 들어오는 피처만 0.33~0.39(셀 F 0.37)였다. 마진 그래프의 중앙 상대오차는 공개 다섯 쌍에서 0.013~0.033으로 고정 지원 그래프의 0.247~0.669를 크게 낮췄고, 셀 F에서는 0.017 대 0.786이었다. 재정렬만 제거하면 셀 F 오차가 0.558로 올라 재정렬이 핵심임을 보였다. 들어오는 집합의 정밀도는 0.80~0.90, 재현율은 0.82~0.92였고, 립시츠 밴드 감사는 668,393,327개 예측 소속 중 잘못된 소속을 0개 받아들였다. ReLU 쌍은 0.69에서 교차했지만 교차 좌표가 응답 에너지의 0.05만 차지했고 코드 변화가 프리액티베이션 변화의 0.43배에 그쳐 연속적인 규칙의 차이를 보였다.

어떻게 쓰나

문헌의 1차 보정 장치도 비교한다. 셀 F에서 적분 그래디언트는 고정 지원 위에서 5·10개 곱 모두 0.777, 어텐션 소프트맥스 재계산은 0.783이었고, 재정렬하면 소프트맥스 장치가 0.013으로 마진 그래프의 0.017보다 낮았다. 공개 다섯 쌍에서도 재정렬한 소프트맥스 장치는 0.011~0.026으로 마진 그래프보다 좋았지만 추가 경로 비용이 든다. 한 번 만든 그래프를 여러 피처 부분집합 절제에 재사용하면 셀 F의 2개 부분집합에서 오차 0.287(고정 지원 0.770), 16개에서는 0.458(0.551)이었으나, 공개 릴리스의 16개 부분집합에서는 0.670으로 고정 지원의 0.631보다 나빠지는 경우도 있었다. 적분 그래디언트를 재정렬하는 경로 예측기는 n=5에서 셀 F 0.0043 이하, 릴리스 0.124 이하였다.

전제와 한계

디코더가 코드 변화를 활성화 변화로 옮길 때 교차 피처의 몫이 크다는 결과도 중요하다. 교차 개입에서 들어오고 나가는 피처의 디코딩 변화는 모델 자체 활성화 변화 노름의 중앙값 1.24~1.72배(셀 F 1.56)였고, 전체 디코딩 변화는 1.41~1.82배였지만 코사인은 0.27~0.50에 그쳤다. 저자는 교차 단일 피처 쓰기 345,178개에서 측정값이 정리 4의 하한보다 낮지 않음을 확인했다. pythia-70m 3·4·5층 사슬 실험에서는 고정 지원, 마진, 정확한 중간 코드를 다음 블록에 통과시키는 사슬이 3층 코드에서 각각 0.884, 0.999, 0.996의 오차를 냈고, 정확한 활성화 변화 사슬은 0.029, 잔차 스트림 사슬은 0.053이었다. 즉 중간 코드를 디코딩해 넘기는 사슬은 실패하고 잔차 스트림을 통해 한 번 더 곱하는 쪽이 안정적이었다.

시드 간 비교에서는 같은 모델의 두 시드 사전 A,B에 대해 디코더 열을 D_A≈D_B M, 인코더 행을 E_A≈N E_B인 희소 비음수 맵으로 표현한다. 정확하면 W_A=N W_B M이 성립하고, 마스킹된 항등식에서 마스크가 바꾸는 부분은 B의 마진이 양이 아닌 원자에 실린 운반 항이며 크기가 아니라 마진의 부호에만 의존한다는 정리 3을 제시한다. 셀 F에서 비음수 직교 매칭 추적으로 상대 잔차 0.436, 원자 최대 32개로 맵을 적합했고, 6개 독립 시드와 재훈련 쌍 양방향의 30개 순서쌍을 평가했다. 마스크 없는 연산자에서 희소 맵은 시드 간 에너지의 0.57을 설명해 1대1 매칭의 0.22보다 0.35 높았고, 마스킹된 코드 야코비안에서는 0.47, 대상 사전 전체를 허용한 비교에서는 0.85를 설명했다. 손실의 대부분은 컷오프에 가장 가까운 비활성 원자에 있었고, 비활성 원자의 10%를 컷오프 근처에서 허용하면 시드 간 격차의 0.73, 재훈련 쌍의 0.76을 회복했다.

개발자 관점에서 이 논문은 SAE 어트리뷰션 그래프를 회로 해석이나 피처 절제 예측에 쓸 때 활성 노드만 고정하는 기본 가정을 재검토하라고 말한다. 예측하려는 개입이 컷오프 근처 피처의 소속을 바꿀 수 있다면, 그래프를 만들 때 하류 행 전체를 보관하고 개입 후 프리액티베이션에 TopK 재정렬을 적용해야 하며 추가 야코비안 곱은 들지 않는다. 또한 시드 간 그래프 비교에서는 피처별 코사인 유사도만 보지 말고 두 층 사이 연산자를 희소 맵으로 옮겨 마스크 손실, 특히 컷오프 근처 비활성 원자의 기여를 확인해야 한다. 다만 마진 그래프도 큰 부분집합 절제나 비아핀 구간에서는 오차가 커질 수 있어, 립시츠 밴드 감사와 쓰기 노름 영역을 함께 봐야 한다.

저자가 밝힌 한계는 분명하다. JumpReLU 사전은 연구하지 못했고, 인접 잔차층이 있는 공개 릴리스는 접근 계약 때문에 사용할 수 없었으며, 트랜스코더도 다루지 않았다. 가장 큰 모델은 pythia-160m이고, 시드 간 운반은 두 번째 시드 공개 릴리스가 없어 셀 F 한 곳과 재훈련 쌍에서만 측정했다. 적분 그래디언트는 셀 F에서만 평가했고, 런타임과 메모리는 측정하지 않고 야코비안-벡터 곱 수로 비용을 셈했다. 따라서 수치의 일반화 범위는 TopK SAE와 Pythia 규모, 셀 F 중심이라는 전제를 염두에 둬야 한다.