Jacobian Lens: 언어 모델 내부 표현을 어휘 토큰으로 읽어내는 해석 도구

9bow · 7월 22, 2026, 11:30오전

Jacobian Lens 소개

언어 모델은 하나의 답을 내놓기까지 내부에서 방대한 계산을 수행하지만, 그 중간 활성화(activation)는 고차원 벡터라서 사람이 곧바로 읽을 수 없습니다. 어떤 레이어의 어떤 표현이 무엇을 담고 있는지 확인하려면 별도의 해석 도구가 필요합니다. 기존에는 로짓 렌즈(logit lens) 처럼 중간 활성화를 곧장 모델의 출력 어휘로 사영해 보는 방법이 쓰였지만, 이 방식은 모든 레이어가 같은 좌표계를 쓴다고 가정하기 때문에 초반 레이어에서는 해석하기 어려운 결과를 냅니다.

Jacobian Lens(J-lens)는 Anthropic이 공개한 해석 도구로, 중간 레이어의 잔차 스트림(residual stream) 벡터를 최종 레이어의 기저로 선형 전송한 뒤 모델 자신의 언임베딩(unembedding)으로 디코딩 합니다. 이렇게 하면 특정 활성화가 모델이 앞으로 말하려는 어떤 토큰과 연결되어 있는지를 어휘 순위로 읽어낼 수 있습니다. 저자들은 이를 로짓 렌즈를 원칙적으로 다듬은 기법이라고 설명하며, 레이어마다 달라지는 표현 좌표계를 보정하기 때문에 로짓 렌즈가 실패하던 초반 레이어에서도 의미 있는 정보를 끌어낸다고 밝히고 있습니다.

이 저장소는 Anthropic의 "Verbalizable Representations Form a Global Workspace in Language Models" 연구의 동반 코드(companion code)입니다. 오픈 웨이트(open-weight) 디코더 트랜스포머에 렌즈를 학습(fit)하고 적용(apply)하는 최소 구현과, 레이어 × 위치(position) 격자를 인터랙티브하게 보여주는 시각화를 담고 있습니다. 예제는 Qwen 계열 모델을 사용하며, 다른 HuggingFace 디코더에도 그대로 맞춰 쓸 수 있습니다. 다만 저자들은 이 저장소가 최적화되지 않은 참조 구현(reference implementation) 이며 유지보수나 기여를 받지 않는다고 명시합니다.

Jacobian Lens의 동작 원리

J-lens의 핵심은 어떤 레이어 l 의 활성화 h 를 최종 레이어 기저로 옮기는 선형 변환입니다. 이 변환은 다음과 같이 정의됩니다.

\text{lens}_l(h) = \text{unembed}(J_l \, h), \quad J_l = \mathbb{E}\left[\frac{\partial h_\text{final}}{\partial h_l}\right]

여기서 J_l 은 레이어 l 의 활성화가 최종 레이어 활성화에 미치는 입력–출력 자코비안(Jacobian)의 평균 입니다. 기댓값은 일반 웹 텍스트 코퍼스의 프롬프트, 소스 위치, 그리고 현재와 미래의 모든 타깃 위치에 대해 취합니다. 저자들에 따르면 이 평균화가 핵심인데, 이는 어떤 문맥에서 우연히 말해진 표현과 기회가 되면 말할 준비가 된 표현을 구분해 줍니다. 즉 J-lens는 활성화가 어떤 토큰을 앞으로 발화(verbalize)할 잠재력을 담고 있는지를 측정합니다. 정확한 추정 방식은 jlens.fitting 모듈의 docstring에 정리되어 있습니다.

위 그림은 저장소가 제공하는 ASCII 얼굴 예제입니다. 모델에게 아스키 아트가 무엇을 그린 것인지 묻고, J-lens로 각 (위치, 레이어) 칸의 최상위 토큰을 읽어낸 결과입니다. 코를 나타내는 ^ 문자 위치를 선택하면 중간 레이어에서 렌즈가 "nose(코)" 를 읽어내는데, 정작 프롬프트에는 그 단어가 한 번도 등장하지 않습니다. 모델이 그림을 공간적으로 파싱해 각 부위 개념을 해당 문자 위치에 담고 있었다는 뜻입니다. 격자의 각 칸은 그 자리의 최상위 토큰과 전체 어휘에서의 순위(위첨자)를 함께 보여주며, 맨 아랫줄은 모델의 실제 출력에 해당합니다.

Jacobian Lens 설치와 사용법

설치는 저장소를 클론한 뒤 editable 모드로 진행합니다.

pip install -e .

이미 학습된 렌즈를 적용할 때는 HuggingFace 모델을 jlens.from_hf 로 감싼 뒤, 원하는 위치의 활성화를 판독합니다.

import transformers, jlens

hf = transformers.AutoModelForCausalLM.from_pretrained("org/model").cuda()
tok = transformers.AutoTokenizer.from_pretrained("org/model")
model = jlens.from_hf(hf, tok)

lens = jlens.JacobianLens.from_pretrained("org/lens-repo", filename="model/lens.pt")
lens_logits, model_logits, _ = lens.apply(
    model, "Fact: The currency used in the country shaped like a boot is",
    positions=[-2])
for layer, logits in sorted(lens_logits.items()):
    print(layer, [tok.decode([t]) for t in logits[0].topk(5).indices])

직접 렌즈를 학습하려면 프롬프트 집합과 체크포인트 경로를 넘겨 jlens.fit 을 호출합니다.

lens = jlens.fit(model, prompts=my_prompts, checkpoint_path="out/ckpt.pt")
lens.save("out/jacobian_lens.pt")

논문에서 쓴 렌즈는 사전학습에 준하는 코퍼스에서 128토큰짜리 시퀀스 1,000개를 사용했습니다. 저자들은 품질이 빠르게 포화하며 약 100개 프롬프트만으로도 쓸 만하다고 설명합니다(논문 §9.3). 학습 시간은 대부분 모델 자체의 역전파(backward pass)에 좌우되며, 프롬프트를 겹치지 않는 조각으로 나눠 fit() 을 병렬로 돌린 뒤 JacobianLens.merge() 로 합칠 수 있습니다.

전체 흐름은 저장소의 walkthrough.ipynb 노트북에서 확인할 수 있습니다. 모델과 렌즈를 불러오고, 몇 개 레이어에 적용한 뒤, 위 그림 같은 slice 페이지를 렌더링하는 과정을 담고 있습니다.

Jacobian Lens의 라이선스

Jacobian Lens의 코드는 Apache License 2.0으로 공개되어 개인 및 상업적 목적으로 자유롭게 사용할 수 있습니다. 저장소의 data/ 에 포함된 복제·평가용 프롬프트 집합도 Anthropic이 직접 만든 합성 데이터로, 코드와 동일한 Apache 2.0으로 배포됩니다. 다만 모델 가중치나 텍스트 코퍼스는 저장소에 포함되지 않으며, 실행 시점에 내려받는 모델·데이터셋은 각자의 라이선스를 따릅니다.

Jacobian Lens 논문

Jacobian Lens 프로젝트 GitHub 저장소

더 읽어보기

이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있습니다. 관심있는 내용이시라면 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다.

파이토치 한국 사용자 모임이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일로 보내드립니다! 텔레그램(Telegram)이나 Slack/Discord/Teams/Dooray/GoogleChat 등으로도 새 글 알림을 받으실 수 있습니다.

아래쪽에 좋아요를 눌러주시면 새로운 소식들을 정리하고 공유하는데 힘이 됩니다~