SAE 튜토리얼 진행해보기 - training SAE

SAELens을 사용한 기본 SAE 훈련 이 튜토리얼은 작은 크기의 Sparse Autoencoder(SAE)를 훈련하는 방법을 설명하며, 구체적으로 tiny-stories-1L-21M 모델을 대상으로 합니다.

SAELens 라이브러리는 활발히 개발 중이므로, 이 튜토리얼이 오래되었을 경우 여기를 통해 이슈를 제기해 주세요.

try:
    #import google.colab # type: ignore
    #from google.colab import output
    %pip install sae-lens transformer-lens circuitsvis
except:
    from IPython import get_ipython # type: ignore
    ipython = get_ipython(); assert ipython is not None
    ipython.run_line_magic("load_ext", "autoreload")
    ipython.run_line_magic("autoreload", "2")

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "5"
import torch

from sae_lens import LanguageModelSAERunnerConfig, SAETrainingRunner

if torch.cuda.is_available():
    device = "cuda"
elif torch.backends.mps.is_available():
    device = "mps"
else:
    device = "cpu"

print("Using device:", device)
os.environ["TOKENIZERS_PARALLELISM"] = "false"

모델 선택 및 평가 (건너뛰어도 됩니다)
우리는 TinyStories 모델에서 SAE를 훈련시키기 위해 runner를 사용할 것입니다. 이 모델은 매우 작기 때문에 비교적 빠르게 SAE를 훈련할 수 있습니다. 시작하기 전에 transformer_lens를 사용해 모델을 불러오고, 그 성능을 확인해 보겠습니다.

TransformerLens는 여기에서 유용한 두 가지 기능을 제공합니다(그리고 circuits viz는 세 번째 기능을 추가로 제공합니다):

transformer_lens.utils.test_prompt는 모델이 하나의 토큰을 추론할 수 있는지를 확인하는 데 도움이 됩니다.
HookedTransformer.generate는 우리가 모델에서 샘플링할 때 무슨 일이 발생하는지를 보여줍니다.
circuitsvis.logits.token_log_probs는 프롬프트 내의 여러 위치에서 토큰의 로그 확률을 시각화하는 데 도움이 됩니다.

from transformer_lens import HookedTransformer

model = HookedTransformer.from_pretrained(
    "tiny-stories-1L-21M"
)  # This will wrap huggingface models and has lots of nice utilities.

# here we use generate to get 10 completeions with temperature 1. Feel free to play with the prompt to make it more interesting.
for i in range(5):
    display(
        model.generate(
            "do you know BTS? ",
            stop_at_eos=False,  # avoids a bug on MPS
            temperature=1,
            verbose=False,
            max_new_tokens=20,
        )
    )

BTS는 확실하게 모르는 듯 하다.

우리가 알아차린 한 가지는 모델이 주인공의 이름을 매우 일관되게 반복할 수 있다는 점입니다. 대명사를 출력할 수도 있지만, 일부 이야기에서는 주인공의 이름을 반복합니다. 이는 SAE로 분석하기에 흥미로운 능력처럼 보입니다. 모델이 주인공의 이름을 기억하는 능력을 더 잘 이해하기 위해, 다음 문자가 결정된 프롬프트를 추출하고 TransformerLens의 "test_prompt" 유틸리티를 사용하여 해당 이름의 토큰 순위를 확인해 보겠습니다.

transformer_lens.utils.test_prompt를 사용한 모델 능력 점검

from transformer_lens.utils import test_prompt

# Test the model with a prompt
test_prompt(
    "Once upon a time, there was a little girl named Lily. She lived in a big, happy little girl. On her big adventure,",
    " Lily",
    model,
    prepend_space_to_answer=False,
)

위 출력에서 우리는 모델이 다음 토큰으로 "she"에 약 70%의 확률을 할당하고, "Lily"가 다음 토큰이 될 확률은 13%임을 볼 수 있습니다. Lucy나 Anna 같은 다른 이름들은 높은 순위를 차지하지 않았습니다.

로그 확률로 모델 능력 탐색
하나의 프롬프트에 대한 토큰 순위를 보는 것도 흥미롭지만, 모델을 더 잘 이해하는 방법은 텍스트 내 모든 토큰의 로그 확률을 보는 것입니다. 우리는 circuits_vis 패키지를 사용하여 토큰화를 확인하고, 마우스를 올리면 로그 확률에 따른 상위 5개의 토큰을 볼 수 있는 시각화를 만들 수 있습니다. 더 어두운 토큰은 모델이 실제 다음 토큰에 더 높은 확률을 할당한 경우입니다.

import circuitsvis as cv  # optional dep, install with pip install circuitsvis

Let's make a longer prompt and see the log probabilities of the tokens

example_prompt = """Hi, how are you doing this? I'm really enjoying your posts"""
logits, cache = model.run_with_cache(example_prompt)
cv.logits.token_log_probs(
    model.to_tokens(example_prompt),
    model(example_prompt)[0].log_softmax(dim=-1),
    model.to_string,
)

hover on the output to see the result.

model.generate와 토큰 로그 확률 시각화를 결합하여 모델이 생성한 텍스트에서 로그 확률을 확인해 봅시다. 여기서 온도를 조절하면 모델이 덜 가능성 있는 경로를 샘플링하게 할 수 있다는 점을 유의하세요. 저는 전체 이야기를 얻기 위해 최대 토큰 수를 늘렸습니다.

탐구할 몇 가지 사항:

• 모델이 높은 확률을 할당하는 토큰은 무엇인가요? 모델이 다음에 어떤 단어가 나와야 하는지 어떻게 알 수 있는지 확인해보세요.
• 온도를 높이거나 낮추면 무슨 일이 발생하나요?
• 토큰의 순위가 당신에게 합리적으로 보이나요? 모델이 다음에 나온 토큰에 높은 확률을 할당하지 않은 경우는 어떠신가요?

example_prompt = model.generate(
    "Once upon a time",
    stop_at_eos=False,  # avoids a bug on MPS
    temperature=1,
    verbose=True,
    max_new_tokens=30,
)
logits, cache = model.run_with_cache(example_prompt)
cv.logits.token_log_probs(
    model.to_tokens(example_prompt),
    model(example_prompt)[0].log_softmax(dim=-1),
    model.to_string,
)

LLM(대형 언어 모델)이 선택하는 단어는 항상 가장 높은 확률을 고르는 것이 아니라, 샘플링 전략에 따라 달라집니다. 특히, 당신이 사용하는 코드에서 설정된 매개변수인 temperature와 같은 요소들이 이러한 선택에 영향을 줍니다.

아래에서 이 현상의 주요 원인들을 설명하겠습니다:

1. Temperature (온도)

• temperature=1은 확률 분포를 그 자체로 사용하겠다는 의미입니다. 즉, 가장 높은 확률의 토큰만 선택하는 것이 아니라, 다른 확률이 낮은 토큰들도 적절한 확률에 따라 선택될 수 있습니다.
• temperature 값을 낮추면 모델은 높은 확률을 가진 토큰에 더 집중하고, 값을 높이면 확률이 더 고르게 분포되어 낮은 확률의 토큰도 선택될 가능성이 높아집니다.
  • 예: temperature=0이면 모델은 항상 가장 높은 확률의 토큰을 선택합니다(즉, 결정적).
  • temperature>1이면 낮은 확률을 가진 토큰도 더 자주 선택될 수 있습니다.

2. 샘플링 (Sampling) 전략

• 이 코드에서는 특정 샘플링 전략을 명시적으로 지정하지 않았지만, temperature 값에 따라 확률 기반 샘플링을 사용하는 것으로 보입니다.
• 확률 기반 샘플링은 매 토큰을 고를 때, 확률 분포에서 무작위로 선택합니다. 이 경우, 확률이 낮더라도 완전히 불가능한 것이 아니기 때문에 확률이 낮은 토큰이 선택될 수 있습니다.
• 이를 통해 모델이 항상 가장 높은 확률만을 선택하지 않고, 다양성을 더해 더 창의적이거나 자연스러운 출력을 만들 수 있습니다.

3. Top-k 또는 Top-p 샘플링

• 코드에서 언급된 부분이 없지만, 일반적으로 Top-k 샘플링이나 Top-p(또는 nucleus) 샘플링과 같은 다른 샘플링 기법들이 자주 사용됩니다.
  • Top-k 샘플링: 확률이 높은 상위 k개의 토큰만을 선택 대상으로 삼고 그 안에서 샘플링을 합니다.
  • Top-p 샘플링: 확률의 누적합이 일정 임계값(p)에 도달할 때까지 토큰을 선택 대상으로 삼고 그 안에서 샘플링을 합니다.

하지만 당신의 경우, 이 코드에서는 temperature=1로 설정되었기 때문에 Top-k나 Top-p가 아닌 일반적인 확률 분포에 따른 샘플링이 이루어지고 있으며, 이로 인해 46번째로 높은 확률의 토큰이 선택될 가능성도 생깁니다.

정리하자면, 가장 높은 확률의 토큰만을 고르지 않고, 확률 분포에 따라 다양성을 유지하기 위한 샘플링 과정에서 비교적 낮은 확률의 토큰이 선택된 것입니다.

SAE 훈련

이제 SAE를 훈련할 준비가 되었습니다. 러너 설정을 만들고, 러너를 인스턴스화하면 나머지는 자동으로 처리됩니다!

훈련 중에는 weights and biases를 사용하여 우리가 관심 있는 변수를 얼마나 잘 최적화하고 있는지 나타내는 주요 지표를 확인할 수 있습니다.

어떤 변수를 중점적으로 살펴봐야 하는지 더 잘 이해하려면, 제가 작성한 포스트를 읽어보시고, 특히 weights and biases 리포트는 여기에서 확인해보세요.

몇 가지 팁:

• wandb 대시보드를 자유롭게 재구성하여 L0, CE_Loss_score, 설명된 분산 및 기타 주요 지표를 상단의 한 섹션에 배치하세요.
• 하이퍼파라미터를 튜닝할 때는 run comparer를 만드세요.
• 훈련 결과로 나온 sparse autoencoder / 희소성 추정치를 wandb에서 다운로드하고, 다른 사람들과 공유하고 싶다면 huggingface에 업로드할 수 있습니다.
  • cfg.json (훈련 설정)
  • sae_weight.safetensors (모델 가중치)
  • sparsity.safetensors (희소성 추정치)

MLP 결과

아래의 하이퍼파라미터를 조정하여 86% CE 손실 회복률과 약 85의 L0를 달성하는 적절한 SAE를 얻었으며, M3 Max에서 약 2시간 안에 실행됩니다. L0와 CE 손실만 고려한다면 더 빠르게 성능이 좋은 SAE를 얻을 수 있지만, 더 밀집된 특징과 더 많은 죽은 특징이 있을 가능성이 큽니다. 두 가지 다른 L1으로 두 번 실행한 결과는 여기에서 확인할 수 있습니다.

total\_training\_steps = 30\_000 # probably we should do more  
batch\_size = 4096  
total\_training\_tokens = total\_training\_steps \* batch\_size

lr\_warm\_up\_steps = 0  
lr\_decay\_steps = total\_training\_steps // 5 # 20% of training  
l1\_warm\_up\_steps = total\_training\_steps // 20 # 5% of training

cfg = LanguageModelSAERunnerConfig(  
\# Data Generating Function (Model + Training Distibuion)  
model\_name="tiny-stories-1L-21M", # our model (more options here: [https://neelnanda-io.github.io/TransformerLens/generated/model\_properties\_table.html](https://neelnanda-io.github.io/TransformerLens/generated/model_properties_table.html))  
hook\_name="blocks.0.hook\_mlp\_out", # A valid hook point (see more details here: [https://neelnanda-io.github.io/TransformerLens/generated/demos/Main\_Demo.html#Hook-Points](https://neelnanda-io.github.io/TransformerLens/generated/demos/Main_Demo.html#Hook-Points))  
hook\_layer=0, # Only one layer in the model.  
d\_in=1024, # the width of the mlp output.  
dataset\_path="apollo-research/roneneldan-TinyStories-tokenizer-gpt2", # this is a tokenized language dataset on Huggingface for the Tiny Stories corpus.  
is\_dataset\_tokenized=True,  
streaming=True, # we could pre-download the token dataset if it was small.  
\# SAE Parameters  
mse\_loss\_normalization=None, # We won't normalize the mse loss,  
expansion\_factor=16, # the width of the SAE. Larger will result in better stats but slower training.  
b\_dec\_init\_method="zeros", # The geometric median can be used to initialize the decoder weights.  
apply\_b\_dec\_to\_input=False, # We won't apply the decoder weights to the input.  
normalize\_sae\_decoder=False,  
scale\_sparsity\_penalty\_by\_decoder\_norm=True,  
decoder\_heuristic\_init=True,  
init\_encoder\_as\_decoder\_transpose=True,  
normalize\_activations="expected\_average\_only\_in",  
\# Training Parameters  
lr=5e-5, # lower the better, we'll go fairly high to speed up the tutorial.  
adam\_beta1=0.9, # adam params (default, but once upon a time we experimented with these.)  
adam\_beta2=0.999,  
lr\_scheduler\_name="constant", # constant learning rate with warmup. Could be better schedules out there.  
lr\_warm\_up\_steps=lr\_warm\_up\_steps, # this can help avoid too many dead features initially.  
lr\_decay\_steps=lr\_decay\_steps, # this will help us avoid overfitting.  
l1\_coefficient=5, # will control how sparse the feature activations are  
l1\_warm\_up\_steps=l1\_warm\_up\_steps, # this can help avoid too many dead features initially.  
lp\_norm=1.0, # the L1 penalty (and not a Lp for p < 1)  
train\_batch\_size\_tokens=batch\_size,  
context\_size=512, # will control the lenght of the prompts we feed to the model. Larger is better but slower. so for the tutorial we'll use a short one.  
\# Activation Store Parameters  
n\_batches\_in\_buffer=64, # controls how many activations we store / shuffle.  
training\_tokens=total\_training\_tokens, # 100 million tokens is quite a few, but we want to see good stats. Get a coffee, come back.  
store\_batch\_size\_prompts=16,  
\# Resampling protocol  
use\_ghost\_grads=False, # we don't use ghost grads anymore.  
feature\_sampling\_window=1000, # this controls our reporting of feature sparsity stats  
dead\_feature\_window=1000, # would effect resampling or ghost grads if we were using it.  
dead\_feature\_threshold=1e-4, # would effect resampling or ghost grads if we were using it.  
\# WANDB  
log\_to\_wandb=True, # always use wandb unless you are just testing code.  
wandb\_project="sae\_lens\_tutorial",  
wandb\_log\_frequency=30,  
eval\_every\_n\_wandb\_logs=20,  
\# Misc  
device=device,  
seed=42,  
n\_checkpoints=0,  
checkpoint\_path="checkpoints",  
dtype="float32"  
)

# look at the next cell to see some instruction for what to do while this is running.

sparse\_autoencoder = SAETrainingRunner(cfg).run()

주석을 한글로 번역한 코드

total_training_steps = 30_000  # 아마도 더 많이 해야 할 것 같습니다
batch_size = 4096
total_training_tokens = total_training_steps * batch_size

lr_warm_up_steps = 0
lr_decay_steps = total_training_steps // 5  # 전체 훈련의 20%
l1_warm_up_steps = total_training_steps // 20  # 전체 훈련의 5%

cfg = LanguageModelSAERunnerConfig(
    # 데이터 생성 함수 (모델 + 훈련 분포)
    model_name="tiny-stories-1L-21M",  # 우리 모델 (더 많은 옵션은 여기에 있음: https://neelnanda-io.github.io/TransformerLens/generated/model_properties_table.html)
    hook_name="blocks.0.hook_mlp_out",  # 유효한 후크 지점 (자세한 내용은 여기에 있음: https://neelnanda-io.github.io/TransformerLens/generated/demos/Main_Demo.html#Hook-Points)
    hook_layer=0,  # 모델에 하나의 레이어만 존재함
    d_in=1024,  # MLP 출력의 너비
    dataset_path="apollo-research/roneneldan-TinyStories-tokenizer-gpt2",  # Huggingface에서 토큰화된 Tiny Stories 코퍼스 언어 데이터셋
    is_dataset_tokenized=True,
    streaming=True,  # 데이터셋이 작다면 토큰 데이터를 미리 다운로드할 수 있음
    # SAE 매개변수
    mse_loss_normalization=None,  # mse 손실을 정규화하지 않음
    expansion_factor=16,  # SAE의 너비. 값이 클수록 더 나은 통계를 얻지만 훈련 속도가 느려짐
    b_dec_init_method="zeros",  # 디코더 가중치를 초기화할 때 기하학적 중앙값을 사용할 수 있음
    apply_b_dec_to_input=False,  # 디코더 가중치를 입력에 적용하지 않음
    normalize_sae_decoder=False,
    scale_sparsity_penalty_by_decoder_norm=True,
    decoder_heuristic_init=True,
    init_encoder_as_decoder_transpose=True,
    normalize_activations="expected_average_only_in",
    # 훈련 매개변수
    lr=5e-5,  # 낮을수록 좋음, 튜토리얼 속도를 높이기 위해 꽤 높은 값을 사용함
    adam_beta1=0.9,  # adam 최적화 매개변수 (기본값이지만, 한때 이 값을 실험했었음)
    adam_beta2=0.999,
    lr_scheduler_name="constant",  # 고정 학습률과 워밍업. 더 나은 스케줄이 있을 수 있음
    lr_warm_up_steps=lr_warm_up_steps,  # 초기에는 너무 많은 죽은 특성을 방지하는 데 도움이 됨
    lr_decay_steps=lr_decay_steps,  # 과적합을 방지하는 데 도움이 됨
    l1_coefficient=5,  # 특성 활성화의 희소성을 제어함
    l1_warm_up_steps=l1_warm_up_steps,  # 초기에는 너무 많은 죽은 특성을 방지하는 데 도움이 됨
    lp_norm=1.0,  # L1 페널티 (p < 1에 대한 Lp 페널티가 아님)
    train_batch_size_tokens=batch_size,
    context_size=512,  # 모델에 제공하는 프롬프트의 길이를 제어함. 값이 클수록 좋지만 속도가 느려짐. 튜토리얼을 위해 짧은 값을 사용함
    # 활성화 저장 매개변수
    n_batches_in_buffer=64,  # 저장 / 셔플하는 활성화 수를 제어함
    training_tokens=total_training_tokens,  # 1억 개의 토큰이 꽤 많지만, 좋은 통계를 보기 위해 필요함. 커피 한잔하고 오세요.
    store_batch_size_prompts=16,
    # 재샘플링 프로토콜
    use_ghost_grads=False,  # 이제는 고스트 그라드를 사용하지 않음
    feature_sampling_window=1000,  # 특성 희소성 통계 보고를 제어함
    dead_feature_window=1000,  # 고스트 그라드 또는 재샘플링을 사용할 경우에만 영향을 미침
    dead_feature_threshold=1e-4,  # 고스트 그라드 또는 재샘플링을 사용할 경우에만 영향을 미침
    # WANDB
    log_to_wandb=True,  # 코드 테스트가 아니라면 항상 wandb를 사용하세요
    wandb_project="sae_lens_tutorial",
    wandb_log_frequency=30,
    eval_every_n_wandb_logs=20,
    # 기타
    device=device,
    seed=42,
    n_checkpoints=0,
    checkpoint_path="checkpoints",
    dtype="float32"
)
# 이 실행 중에 무엇을 해야 할지에 대한 지침은 다음 셀을 참조하세요.
sparse_autoencoder = SAETrainingRunner(cfg).run()

코드 설명

이 코드는 Sparse Autoencoder(SAE)를 훈련시키기 위한 설정을 정의하고, 훈련을 실행하는 역할을 합니다. 주요 기능을 설명하자면:

1. 훈련 설정: total_training_steps와 batch_size를 기반으로 훈련에서 사용할 총 토큰 수를 설정합니다. 그 외에도 학습률 스케줄링, L1 규제와 관련된 설정이 포함되어 있습니다.
2. SAE 매개변수: SAE의 구조적 매개변수를 정의하는 부분입니다. 확장 계수(expansion_factor), 디코더 초기화 방식, 희소성 페널티를 제어하는 옵션들이 정의됩니다.
3. 훈련 매개변수: 학습률, Adam 최적화기의 매개변수, 학습률 워밍업 및 감소 단계가 설정됩니다. 특히 L1 규제를 사용하여 희소성을 조절하며, 워밍업 단계가 포함되어 있어 초기에는 죽은 특성(dead features)을 방지하도록 설정됩니다.
4. 활성화 저장 및 재샘플링: 훈련 중 활성화 값들을 저장하고, 고스트 그라드를 사용하지 않으며, 재샘플링과 관련된 창 및 임계값 설정을 포함합니다.
5. WANDB 통합: 훈련 로그를 Weights and Biases(WANDB) 툴로 기록하여, 시각화 및 성능 추적을 쉽게 할 수 있도록 설정됩니다.
6. 기타 설정: device, seed, checkpoint 관련 설정을 통해 훈련 중 사용할 기기 및 체크포인트 옵션을 지정합니다.

마지막으로 SAETrainingRunner를 사용해 설정값을 적용한 후 SAE 훈련이 시작됩니다.

이 두 학습 과정의 차이를 설명드리겠습니다.

1. Estimating norm scaling factor:
   • 이 단계는 정규화 스케일링 계수를 추정하는 과정입니다. Sparse Autoencoder(SAE) 모델을 훈련하기 전에, 모델에서 사용하는 가중치나 활성화를 조정하기 위한 정규화 계수를 계산하는 단계입니다.
   • 이 단계는 짧은 시간(33초) 동안 이루어지며, 모델의 기본 구조를 준비하는 과정으로 볼 수 있습니다.
   • 보통 이 단계는 모델의 훈련과 직접적으로 관련되지 않으며, 이후 모델이 데이터에 맞춰 학습될 때 가중치나 활성화를 어떻게 스케일링할지를 결정하는 데 필요한 값을 설정하는 역할을 합니다.
2. 30000| MSE Loss 187.703 | L1 156.883:
   • 이 부분은 실제 훈련 과정입니다. 모델이 주어진 데이터셋을 바탕으로 MSE 손실과 L1 규제를 적용하여 학습이 진행되고 있는 상황을 나타냅니다.
   • MSE Loss 187.703는 현재 모델이 출력과 실제 값 사이에서 예측 오류가 어느 정도인지 보여주며, L1 156.883는 모델의 희소성(얼마나 많은 특성들이 활성화되고 있는지)을 제어하는 규제 항목을 의미합니다.
   • 1%| 1228800/122880000는 전체 훈련 과정 중 1%가 진행되었음을 나타내며, 이는 모델이 122,880,000개의 토큰 중 1,228,800개의 토큰에 대해 학습을 완료했음을 의미합니다.
   • 476.34it/s는 초당 처리되는 아이템 수를 나타냅니다. 이 속도는 훈련의 속도를 반영하며, 훈련이 완료되기까지 남은 시간이 약 70시간 정도임을 보여줍니다.

차이점:

• Estimating norm scaling factor는 훈련에 앞서 모델이 데이터를 효율적으로 학습할 수 있도록 사전 계산 단계를 수행하는 것이며, 이 단계는 짧고 단순한 계산입니다.
• 반면에 30000|MSE Loss...는 실제 훈련 과정으로, 모델이 데이터를 기반으로 손실을 줄이고 희소성을 제어하며 최적화되는 과정입니다. 이 단계는 시간이 오래 걸리고, 모델이 학습하면서 성능이 향상되는 것을 의미합니다.

따라서, 첫 번째는 훈련을 준비하는 단계이고, 두 번째는 본격적인 훈련 단계입니다.

해야 할 일: SAE를 사용하여 TinyStories-1L 이해하기

아직 이 섹션을 완성할 시간이 없었지만, 누군가 이 튜토리얼에서 훈련한 SAE를 사용하여 이 모델을 더 잘 이해하는 PR을 보면 좋겠습니다.

import pandas as pd

# Let's start by getting the top 10 logits for each feature
projection_onto_unembed = sparse_autoencoder.W_dec @ model.W_U


# get the top 10 logits.
vals, inds = torch.topk(projection_onto_unembed, 10, dim=1)

# get 10 random features
random_indices = torch.randint(0, projection_onto_unembed.shape[0], (10,))

# Show the top 10 logits promoted by those features
top_10_logits_df = pd.DataFrame(
    [model.to_str_tokens(i) for i in inds[random_indices]],
    index=random_indices.tolist(),
).T
top_10_logits_df