로봇 정책 π0.5에서 목표 단어 특징 교체가 다른 작업을 유도한다
Editing the Goal: Per-Token SAE Steering and Its Depth Limits in $π_{0.5}$
무엇인가
이 논문은 로봇 파운데이션 모델(RFM) π0.5 내부에서 지시문의 특정 단어 표현을 조작했을 때 로봇이 실제로 다른 목표를 수행하게 만들 수 있는지, 그리고 그 조작이 모델의 몇 번째 층까지 유효한지를 다룬다. 배경에는 두 가지 선행 결과가 있다. Swann 등은 π0.5에서 해석 가능하고 인과적인 SAE 특징을 찾았지만 대부분 타임스텝 단위로 풀링한 활성값에서 나왔고, 개별 지시문 토큰 단위의 SAE 특징은 해석성이 떨어졌다. Grant 등은 π0.5의 action expert에서 토큰 단위 SAE 특징을 스티어링하면 정책이 다른 과제로 전환되기보다 오히려 성능이 나빠진다고 보고했다. 저자들은 "한 단어의 SAE 특징 교체가 목표를 바꿀 수 있는가", "지시문은 어느 깊이까지 편집 가능한가"가 미해결로 남아 있다고 본다. 물리 세계에서 위험한 궤적은 되돌릴 수 없기 때문에, 모델이 어떤 목표를 추구하는지 투명하게 확인하고 변조 여부를 검출하는 수단이 필요하다는 것이 동기다.
어떻게 동작하나
실험 기반은 π0.5-LIBERO의 PaliGemma 잔차 스트림이다. openpi 체크포인트를 사용해 층 0, 5, 11, 17에서 두 종류의 SAE를 학습했다. 하나는 Swann等人的 설정을 따른 BatchTopK SAE로 AuxK 손실과 행별 TopK, 1배 확장(2,048개 latent)을 쓴다. 다른 하나는 64·64·128·256·512·1024개 latent 그룹으로 구성한 Matryoshka SAE로, 특징 분할이 적고 계층 구조가 있다는 점이 RFM에도 전이되리라는 가정에서 출발했다. SAE 입력 행은 카메라 한 대의 이미지 토큰 256개의 평균이거나 개별 텍스트 토큰 하나다. 학습된 SAE는 모두 분산의 98% 이상을 재구성했다. 장면 선택은 엄격하다. LIBERO 정책은 지시문을 무시하는 경우가 많아 LIBERO-Spatial에서 프롬프트 없이도 61/100 에피소드가 성공했기 때문에, 빈 프롬프트 성공률이 10회 중 3회 이하인 장면만 사용했다. LIBERO-Spatial에서는 검은 그릇 두 개 중 올바른 것을 골라 접시 위에 놓는 과제에서 7번("on the stove")과 9번("on the wooden cabinet")이라는 거울 쌍을 쓰고, LIBERO-Object에서는 지시문이 담을 물체만 달라지는 장면 중 지시문 교체만으로 5/5 롤아웃이 전환되는 4개 장면으로 교차 검증했다.
무엇과 다른가
개입 절차는 다음과 같다. 프롬프트 A의 목표 단어 토큰들에 대해 잔차 스트림 h를 SAE로 인코딩해 코드 z를 얻고, 선택한 latent 집합 S의 각 z_j를 목표값 c_j로 바꾼 뒤 디코더 변화량을 원래 활성값에 더한다(x' = x + Dec(z') − Dec(z)). 재구성 오차는 그대로 유지하며, SAE가 각 행을 단위 노름으로 재스케일하기 때문에 8번 재인코딩해 편집된 코드를 목표값의 0.18 이내로 수렴시킨다. 특징 선택은 두 가지다. 단어 선택 방식은 해당 단어를 포함한 모든 지시문에서의 평균 활성 대 다른 모든 텍스트 토큰에서의 평균 활성 비율이 가장 높은 16개(최대 64개) latent를 고르고, 전체 행의 절반 이상에서 켜져 있는 attention-sink류 latent는 제외한다. A의 latent는 0으로, B의 latent는 B에서의 평균 활성으로 설정한다. 대조 선택 방식은 같은 시연 프레임에서 프롬프트 A와 B 사이에 코드가 가장 크게 변하는 16개(또는 64개) latent를 고르고 B에서의 평균값을 목표로 쓴다. 비교를 위해 차이 벡터 방식(과제 A의 256개 시연 프레임에 걸쳐 평균한 목표 단어 상태 차이를 α배로 더함, α=1이 원래 크기)과 활성 패칭(목표 단어 상태를 프롬프트 B의 상태로 교체하고 이후 블록을 재계산)도 함께 썼다. 모든 조건에는 같은 수의 무작위 latent를 같은 목표값으로 바꾸는 대조 편집, 같은 노름의 무작위 벡터, 프롬프트 B의 임의 비목표 단어 패칭이라는 통제가 붙는다.
어떻게 쓰나
핵심 결과는 층 0에서 BatchTopK 특징 집합 하나로 프롬프트 A를 받은 로봇이 B의 과제를 20/20 롤아웃에서 완수했다는 것이다. 같은 장면에서 무작위 특징 교체는 0/20이었다. A 단어의 latent를 0으로만 만들고 B의 latent를 쓰지 않으면, 층 0의 16-latent 교체가 20회 중 16회 이상 성공한 5개 장면에서 최대 1/20만 전환됐고, B 과제를 놓친 롤아웃의 89%는 A 과제를 완수했다. 즉 목표를 지우는 것과 다른 목표를 채워 넣는 것은 다르다. 성공률은 두 LIBERO 스위트의 6개 장면과 층에 따라 달랐다. Matryoshka는 단어를 덜 쪼갠다. 층 5에서 서로의 발화 토큰 절반 이상에서 같이 켜지는 latent 쌍을 세면 Matryoshka 1,323쌍, BatchTopK 3,853쌍(live latent당 0.65 대 1.88)이었다. 그러나 이 성질은 스티어링에 도움이 되지 않아, LIBERO-Object에서 Matryoshka의 단어 선택 교체는 무작위 교체와 비슷한 성적을 냈다. Matryoshka의 높은 무작위 성공률(최대 17/20)은 무작위 층 0 편집의 크기가 실제 단어 변화량의 중앙값 16~54배였기 때문으로 분석된다. 단일 특징 제거 실험도 인과 효과가 작음을 보여준다. 층 0에서 SAE·장면별 상위 8개 목표 단어 특징과 무작위 8개, 총 40건을 제거했을 때 39건은 원래 과제 성공률이 80% 이상으로 유지됐다. 예외는 카메라 행에서도 단어만큼 활성화되는 BatchTopK의 "stove" 관련 특징으로, 장면 7을 20% 미만으로 떨어뜨렸지만 장면 9는 90%로 남겼고, 카메라 토큰에서만 제거하면 6/10, 단어에서만 제거하면 9/10이었다.
전제와 한계
SAE가 필수 도구는 아니라는 점도 분명히 한다. 목표 단어의 상태를 프롬프트 B에서의 상태로 패칭하면 두 LIBERO-Spatial 장면 모두에서 층 0과 5에서 10/10으로 B 과제를 완수했다. 테스트 시점에 프롬프트 B 순전파가 필요 없는 고정 차이 벡터도 1배에서 5~9/10, 2배에서 8~10/10을 기록했고, 같은 노름의 무작위 벡터는 1배·2배 모두 0/10이었다. 저자들은 SAE의 부가가치를 "프롬프트 A의 프레임에서 프롬프트 B를 한 번도 실행하지 않는 단어 단위 교체"로 규정한다. 패칭은 매 정책 호출마다 B를 실행해야 하고, 차이 벡터는 시연 단계에서 한 번은 B를 실행해 두어야 한다.
편집 가능 깊이의 지도는 활성 패칭으로 그렸다. 목표 단어는 물체 위치 과제에서는 층 6까지 편집이 유효했지만 물체 정체성 과제에서는 층 2까지만 통했다. 패치를 해당 층에서 17층까지 유지하거나 지시문 전체로 넓혀도 성능 하락 양상은 비슷했다. 카메라 토큰까지 함께 패칭하면 두 스위트 모두 층 13까지 과제가 전이됐다. 층 17의 편집은 키·값을 통해 하위 블록으로 전달될 경로가 없어 작동할 수 없고, 층 L에서 17까지 유지한 패치는 L+1부터의 KV 캐시 패치와 동등하다. SAE 교체는 층 0에서 작동하고 층 5에서는 부분적으로(LIBERO-Object에서 최대 4/20), 층 11에서는 작동하지 않았다.
개발자 관점에서 이 결과는 두 가지를 시사한다. 첫째, 지시문 한 단어의 잔차 스트림 표현만 바꿔도 정책의 목표가 바뀔 수 있으므로, VLA를 배포하거나 파인튜닝할 때 프롬프트·활성값 무결성 검사를 프롬프트 문자열 검사와 별개로 두어야 한다. 둘째, SAE 특징을 뽑는 것만으로는 부족하고 단일 특징 제거가 대부분 성공률을 유지했다는 사실처럼 특징 단위 인과 효과를 반드시 개별 검증해야 하며, 간단한 차이 벡터나 패칭 베이스라인을 함께 돌려 SAE가 실제로 이득인지 확인해야 한다. 저자들도 언어 모델에서 단순 베이스라인이 SAE 기반 스티어링과 대등하거나 더 나은 사례가 보고됐음을 인용한다.
한계는 저자들이 명시한다. 이 결과는 시뮬레이션 환경의 단일 파인튜닝 VLA(π0.5-LIBERO)에서의 존재 증명이며, SAE는 4개 층만 덮어 깊이 연구는 원시 활성 패칭에 의존했다. 성공률은 장면과 층에 따라 크게 변하고, Matryoshka에 대한 사전 가설은 빗나갔다. 또한 LIBERO 정책이 지시문을 자주 무시한다는 전제 때문에 빈 프롬프트 성공률이 낮은 장면만 골라 실험했으므로, 일반적인 지시 준수 정책에 그대로 일반화된다고 주장하지 않는다. 결론적으로 이 논문은 목표 단어 편집이 자연스러운 크기에서 어느 층까지 잠재적 목표 하이재킹 위험을 갖는지에 대한 초기 지도를 제공한다.