在解釋機器學習模型方面，稀疏自編碼器（SAE）是一種越來越常用的工具（雖然 SAE 在 1997 年左右就已經(jīng)問世了）。

機器學習模型和 LLM 正變得越來越強大、越來越有用，但它們?nèi)耘f是黑箱，我們并不理解它們完成任務的方式。理解它們的工作方式應當大有助益。

SAE 可幫助我們將模型的計算分解成可以理解的組件。近日，LLM 可解釋性研究者 Adam Karvonen 發(fā)布了一篇博客文章，直觀地解釋了 SAE 的工作方式。

可解釋性的難題

神經(jīng)網(wǎng)絡最自然的組件是各個神經(jīng)元。不幸的是，單個神經(jīng)元并不能便捷地與單個概念相對應，比如學術引用、英語對話、HTTP 請求和韓語文本。在神經(jīng)網(wǎng)絡中，概念是通過神經(jīng)元的組合表示的，這被稱為疊加（superposition）。

之所以會這樣，是因為世界上很多變量天然就是稀疏的。

舉個例子，某位名人的出生地可能出現(xiàn)在不到十億分之一的訓練 token 中，但現(xiàn)代 LLM 依然能學到這一事實以及有關這個世界的大量其它知識。訓練數(shù)據(jù)中單個事實和概念的數(shù)量多于模型中神經(jīng)元的數(shù)量，這可能就是疊加出現(xiàn)的原因。

近段時間，稀疏自編碼器（SAE）技術越來越常被用于將神經(jīng)網(wǎng)絡分解成可理解的組件。SAE 的設計靈感來自神經(jīng)科學領域的稀疏編碼假設?，F(xiàn)在，SAE 已成為解讀人工神經(jīng)網(wǎng)絡方面最有潛力的工具之一。SAE 與標準自編碼器類似。

常規(guī)自編碼器是一種用于壓縮并重建輸入數(shù)據(jù)的神經(jīng)網(wǎng)絡。

舉個例子，如果輸入是一個 100 維的向量（包含 100 個數(shù)值的列表）；自編碼器首先會讓該輸入通過一個編碼器層，讓其被壓縮成一個 50 維的向量，然后將這個壓縮后的編碼表示饋送給解碼器，得到 100 維的輸出向量。其重建過程通常并不完美，因為壓縮過程會讓重建任務變得非常困難。

一個標準自編碼器的示意圖，其有 1x4 的輸入向量、1x2 的中間狀態(tài)向量和 1x4 的輸出向量。單元格的顏色表示激活值。輸出是輸入的不完美重建結果。

解釋稀疏自編碼器

稀疏自編碼器的工作方式

稀疏自編碼器會將輸入向量轉(zhuǎn)換成中間向量，該中間向量的維度可能高于、等于或低于輸入的維度。在用于 LLM 時，中間向量的維度通常高于輸入。在這種情況下，如果不加額外的約束條件，那么該任務就很簡單，SAE 可以使用單位矩陣來完美地重建出輸入，不會出現(xiàn)任何意料之外的東西。但我們會添加約束條件，其中之一是為訓練損失添加稀疏度懲罰，這會促使 SAE 創(chuàng)建稀疏的中間向量。

舉個例子，我們可以將 100 維的輸入擴展成 200 維的已編碼表征向量，并且我們可以訓練 SAE 使其在已編碼表征中僅有大約 20 個非零元素。

稀疏自編碼器示意圖。請注意，中間激活是稀疏的，僅有 2 個非零值。

我們將 SAE 用于神經(jīng)網(wǎng)絡內(nèi)的中間激活，而神經(jīng)網(wǎng)絡可能包含許多層。在前向通過過程中，每一層中和每一層之間都有中間激活。

舉個例子，GPT-3 有 96 層。在前向通過過程中，輸入中的每個 token 都有一個 12,288 維向量（一個包含 12,288 個數(shù)值的列表）。此向量會累積模型在每一層處理時用于預測下一 token 的所有信息，但它并不透明，讓人難以理解其中究竟包含什么信息。

我們可以使用 SAE 來理解這種中間激活。SAE 基本上就是「矩陣 → ReLU 激活 → 矩陣」。

舉個例子，如果 GPT-3 SAE 的擴展因子為 4，其輸入激活有 12,288 維，則其 SAE 編碼的表征有 49,512 維（12,288 x 4）。第一個矩陣是形狀為 (12,288, 49,512) 的編碼器矩陣，第二個矩陣是形狀為 (49,512, 12,288) 的解碼器矩陣。通過讓 GPT 的激活與編碼器相乘并使用 ReLU，可以得到 49,512 維的 SAE 編碼的稀疏表征，因為 SAE 的損失函數(shù)會促使實現(xiàn)稀疏性。

通常來說，我們的目標讓 SAE 的表征中非零值的數(shù)量少于 100 個。通過將 SAE 的表征與解碼器相乘，可得到一個 12,288 維的重建的模型激活。這個重建結果并不能與原始的 GPT 激活完美匹配，因為稀疏性約束條件會讓完美匹配難以實現(xiàn)。

一般來說，一個 SAE 僅用于模型中的一個位置舉個例子，我們可以在 26 和 27 層之間的中間激活上訓練一個 SAE。為了分析 GPT-3 的全部 96 層的輸出中包含的信息，可以訓練 96 個分立的 SAE—— 每層的輸出都有一個。如果我們也想分析每一層內(nèi)各種不同的中間激活，那就需要數(shù)百個 SAE。為了獲取這些 SAE 的訓練數(shù)據(jù)，需要向這個 GPT 模型輸入大量不同的文本，然后收集每個選定位置的中間激活。

下面提供了一個 SAE 的 PyTorch 參考實現(xiàn)。其中的變量帶有形狀注釋，這個點子來自 Noam Shazeer，參見：https://medium.com/@NoamShazeer/shape-suffixes-good-coding-style-f836e72e24fd 。請注意，為了盡可能地提升性能，不同的 SAE 實現(xiàn)往往會有不同的偏置項、歸一化方案或初始化方案。最常見的一種附加項是某種對解碼器向量范數(shù)的約束。更多細節(jié)請訪問以下實現(xiàn)：

• OpenAI：https://github.com/openai/sparse_autoencoder/blob/main/sparse_autoencoder/model.py#L16
• SAELens：https://github.com/jbloomAus/SAELens/blob/main/sae_lens/sae.py#L97
• dictionary_learning：https://github.com/saprmarks/dictionary_learning/blob/main/dictionary.py#L30

import torch
import torch.nn as nn

# D = d_model, F = dictionary_size
# e.g. if d_model = 12288 and dictionary_size = 49152
# then model_activations_D.shape = (12288,) and encoder_DF.weight.shape = (12288, 49152)

class SparseAutoEncoder (nn.Module):
    """
    A one-layer autoencoder.
    """
    def __init__(self, activation_dim: int, dict_size: int):
        super ().__init__()
        self.activation_dim = activation_dim
        self.dict_size = dict_size

        self.encoder_DF = nn.Linear (activation_dim, dict_size, bias=True)
        self.decoder_FD = nn.Linear (dict_size, activation_dim, bias=True)

    def encode (self, model_activations_D: torch.Tensor) -> torch.Tensor:
        return nn.ReLU ()(self.encoder_DF (model_activations_D))

    def decode (self, encoded_representation_F: torch.Tensor) -> torch.Tensor:
        return self.decoder_FD (encoded_representation_F)

    def forward_pass (self, model_activations_D: torch.Tensor) -> tuple [torch.Tensor, torch.Tensor]:
        encoded_representation_F = self.encode (model_activations_D)
        reconstructed_model_activations_D = self.decode (encoded_representation_F)
        return reconstructed_model_activations_D, encoded_representation_F

標準自編碼器的損失函數(shù)基于輸入重建結果的準確度。為了引入稀疏性，最直接的方法是向 SAE 的損失函數(shù)添加一個稀疏度懲罰項。對于這個懲罰項，最常見的計算方式是取這個 SAE 的已編碼表征（而非 SAE 權重）的 L1 損失并將其乘以一個 L1 系數(shù)。這個 L1 系數(shù)是 SAE 訓練中的一個關鍵超參數(shù)，因為它可確定實現(xiàn)稀疏度與維持重建準確度之間的權衡。

請注意，這里并沒有針對可解釋性進行優(yōu)化。相反，可解釋的 SAE 特征是優(yōu)化稀疏度和重建的一個附帶效果。下面是一個參考損失函數(shù)。

# B = batch size, D = d_model, F = dictionary_size
def calculate_loss (autoencoder: SparseAutoEncoder, model_activations_BD: torch.Tensor, l1_coeffient: float) -> torch.Tensor:
    reconstructed_model_activations_BD, encoded_representation_BF = autoencoder.forward_pass (model_activations_BD)
    reconstruction_error_BD = (reconstructed_model_activations_BD - model_activations_BD).pow (2)
    reconstruction_error_B = einops.reduce (reconstruction_error_BD, 'B D -> B', 'sum')
    l2_loss = reconstruction_error_B.mean ()


    l1_loss = l1_coefficient * encoded_representation_BF.sum ()
    loss = l2_loss + l1_loss
    return loss

稀疏自編碼器的前向通過示意圖。

這是稀疏自編碼器的單次前向通過過程。首先是 1x4 大小的模型向量。然后將其乘以一個 4x8 的編碼器矩陣，得到一個 1x8 的已編碼向量，然后應用 ReLU 將負值變成零。這個編碼后的向量就是稀疏的。之后，再讓其乘以一個 8x4 的解碼器矩陣，得到一個 1x4 的不完美重建的模型激活。

假想的 SAE 特征演示

理想情況下，SAE 表征中的每個有效數(shù)值都對應于某個可理解的組件。

這里假設一個案例進行說明。假設一個 12,288 維向量 [1.5, 0.2, -1.2, ...] 在 GPT-3 看來是表示「Golden Retriever」（金毛犬）。SAE 是一個形狀為 (49,512, 12,288) 的矩陣，但我們也可以將其看作是 49,512 個向量的集合，其中每個向量的形狀都是 (1, 12,288)。如果該 SAE 解碼器的 317 向量學習到了與 GPT-3 那一樣的「Golden Retriever」概念，那么該解碼器向量大致也等于 [1.5, 0.2, -1.2, ...]。

無論何時 SAE 的激活的 317 元素是非零的，那么對應于「Golden Retriever」的向量（并根據(jù) 317 元素的幅度）會被添加到重建激活中。用機械可解釋性的術語來說，這可以簡潔地描述為「解碼器向量對應于殘差流空間中特征的線性表征」。

也可以說有 49,512 維的已編碼表征的 SAE 有 49,512 個特征。特征由對應的編碼器和解碼器向量構成。編碼器向量的作用是檢測模型的內(nèi)部概念，同時最小化其它概念的干擾，盡管解碼器向量的作用是表示「真實的」特征方向。研究者的實驗發(fā)現(xiàn)，每個特征的編碼器和解碼器特征是不一樣的，并且余弦相似度的中位數(shù)為 0.5。在下圖中，三個紅框?qū)趩蝹€特征。

稀疏自編碼器示意圖，其中三個紅框?qū)?SAE 特征 1，綠框?qū)谔卣?4。每個特征都有一個 1x4 的編碼器向量、1x1 的特征激活和 1x4 的解碼器向量。重建的激活的構建僅使用了來自 SAE 特征 1 和 4 的解碼器向量。如果紅框表示「紅顏色」，綠框表示「球」，那么該模型可能表示「紅球」。

那么我們該如何得知假設的特征 317 表示什么呢？目前而言，人們的實踐方法是尋找能最大程度激活特征并對它們的可解釋性給出直覺反應的輸入。能讓每個特征激活的輸入通常是可解釋的。

舉個例子，Anthropic 在 Claude Sonnet 上訓練了 SAE，結果發(fā)現(xiàn)：與金門大橋、神經(jīng)科學和熱門旅游景點相關的文本和圖像會激活不同的 SAE 特征。其它一些特征會被并不顯而易見的概念激活，比如在 Pythia 上訓練的一個 SAE 的一個特征會被這樣的概念激活，即「用于修飾句子主語的關系從句或介詞短語的最終 token」。

由于 SAE 解碼器向量的形狀與 LLM 的中間激活一樣，因此可簡單地通過將解碼器向量加入到模型激活來執(zhí)行因果干預。通過讓該解碼器向量乘以一個擴展因子，可以調(diào)整這種干預的強度。當 Anthropic 研究者將「金門大橋」SAE 解碼器向量添加到 Claude 的激活時，Claude 會被迫在每個響應中都提及「金門大橋」。

下面是使用假設的特征 317 得到的因果干預的參考實現(xiàn)。類似于「金門大橋」Claude，這種非常簡單的干預會迫使 GPT-3 模型在每個響應中都提及「金毛犬」。

def perform_intervention (model_activations_D: torch.Tensor, decoder_FD: torch.Tensor, scale: float) -> torch.Tensor:
    intervention_vector_D = decoder_FD [317, :]
    scaled_intervention_vector_D = intervention_vector_D * scale
    modified_model_activations_D = model_activations_D + scaled_intervention_vector_D
    return modified_model_activations_D

稀疏自編碼器的評估難題

使用 SAE 的一大主要難題是評估。我們可以訓練稀疏自編碼器來解釋語言模型，但我們沒有自然語言表示的可度量的底層 ground truth。目前而言，評估都很主觀，基本也就是「我們研究一系列特征的激活輸入，然后憑直覺闡述這些特征的可解釋性?！惯@是可解釋性領域的主要限制。

研究者已經(jīng)發(fā)現(xiàn)了一些似乎與特征可解釋性相對應的常見代理指標。最常用的是 L0 和 Loss Recovered。L0 是 SAE 的已編碼中間表征中非零元素的平均數(shù)量。Loss Recovered 是使用重建的激活替換 GPT 的原始激活，并測量不完美重建結果的額外損失。這兩個指標通常需要權衡考慮，因為 SAE 可能會為了提升稀疏性而選擇一個會導致重建準確度下降的解。

在比較 SAE 時，一種常用方法是繪制這兩個變量的圖表，然后檢查它們之間的權衡。為了實現(xiàn)更好的權衡，許多新的 SAE 方法（如 DeepMind 的 Gated SAE 和 OpenAI 的 TopK SAE）對稀疏度懲罰做了修改。下圖來自 DeepMind 的 Gated SAE 論文。Gated SAE 由紅線表示，位于圖中左上方，這表明其在這種權衡上表現(xiàn)更好。

Gated SAE L0 與 Loss Recovered

SAE 的度量存在多個難度層級。L0 和 Loss Recovered 是兩個代理指標。但是，在訓練時我們并不會使用它們，因為 L0 不可微分，而在 SAE 訓練期間計算 Loss Recovered 的計算成本非常高。相反，我們的訓練損失由一個 L1 懲罰項和重建內(nèi)部激活的準確度決定，而非其對下游損失的影響。

訓練損失函數(shù)并不與代理指標直接對應，并且代理指標只是對特征可解釋性的主觀評估的代理。由于我們的真正目標是「了解模型的工作方式」，主觀可解釋性評估只是代理，因此還會有另一層不匹配。LLM 中的一些重要概念可能并不容易解釋，而且我們可能會在盲目優(yōu)化可解釋性時忽視這些概念。

總結

可解釋性領域還有很長的路要走，但 SAE 是真正的進步。SAE 能實現(xiàn)有趣的新應用，比如一種用于查找「金門大橋」導向向量（steering vector）這樣的導向向量的無監(jiān)督方法。SAE 也能幫助我們更輕松地查找語言模型中的回路，這或可用于移除模型內(nèi)部不必要的偏置。

SAE 能找到可解釋的特征（即便目標僅僅是識別激活中的模式），這一事實說明它們能夠揭示一些有意義的東西。還有證據(jù)表明 LLM 確實能學習到一些有意義的東西，而不僅僅是記憶表層的統(tǒng)計規(guī)律。

SAE 也能代表 Anthropic 等公司曾引以為目標的早期里程碑，即「用于機器學習模型的 MRI（磁共振成像）」。SAE 目前還不能提供完美的理解能力，但卻可用于檢測不良行為。SAE 和 SAE 評估的主要挑戰(zhàn)并非不可克服，并且現(xiàn)在已有很多研究者在攻堅這一課題。

有關稀疏自編碼器的進一步介紹，可參閱 Callum McDougal 的 Colab 筆記本：https://www.lesswrong.com/posts/LnHowHgmrMbWtpkxx/intro-to-superposition-and-sparse-autoencoders-colab