logit lens是机器学习可解释性中的一种技术，用于检查基于Transformer (architecture)的大型语言模型的内部表示。其工作原理是取给定层的隐藏状态，应用模型的最终去嵌入矩阵（该矩阵将隐藏状态映射到词汇表上的logits），然后应用softmax以获得概率分布。该分布揭示了模型若在该层停止时会预测什么，从而提供了预测逐层演化的窗口。该方法由nostalgebraist于2020年提出，此后成为研究GPT-2和GPT-3等模型内部计算的标准工具。
与需要训练辅助模型的探针分类器不同，logit lens使用模型自身的输出头，因此是一种零参数且无需训练的方法。它特别适用于具有残差流的模型，因为每层的隐藏状态是加性的，可以直接投影。该技术已用于识别模型何时对特定预测做出承诺，检测“晚期”与“早期”解码，并分析事实回忆和算术推理等现象。然而，其适用性仅限于具有绑定或共享嵌入的模型，对于具有复杂层归一化或非线性变换的模型可能效果不佳。
机制与实现
logit lens作用于transformer的残差流。在典型的transformer中，每层将其输出添加到残差流中，因此第\( l \)层的隐藏状态是初始嵌入和到\( l \)为止所有层输出的总和。最终logits通过应用层归一化（如果存在）然后应用去嵌入矩阵\( W_U \)来计算。logit lens将相同的去嵌入应用于中间隐藏状态，通常在应用最终层归一化（或学习的仿射变换）以考虑分布偏移之后。
对于具有绑定嵌入（输入嵌入和输出去嵌入共享权重）的模型，投影是直接的。对于具有单独去嵌入矩阵的模型，只要去嵌入可用，该技术仍然有效。结果是每个层词汇表上的概率分布序列，可以可视化为热图或用于计算候选标记之间的“logit差异”等指标。
在可解释性中的应用
logit lens已应用于各种可解释性任务。一个常见用途是追踪预测在各层中的发展。例如，在要求模型完成“法国的首都是__”的模型中，logit lens可能显示早期层仅在若干层之后才赋予“巴黎”高概率，从而揭示模型何时检索相关事实。研究人员已利用这一点来识别“归纳头”和有助于特定行为的其他注意力模式。
另一个应用是检测“晚期”解码，即模型可能最初预测一个标记但后来改变主意。通过检查每层的logit lens，可以看到最终预测何时“锁定”以及何时仍然灵活。这对理解模型置信度和设计干预措施具有意义。
该技术还用于研究算术和推理。在GPT-3等模型中，logit lens可以显示中间层表示中间和或进位，为多步计算提供证据。这导致了机制可解释性的进一步工作，例如Anthropic和其他研究小组提出的“归纳头”和“电路”框架。
局限性与批评
尽管有其效用，logit lens有几个局限性。首先，它假设去嵌入矩阵对所有层都是有意义的投影，如果模型的隐藏状态经历了未考虑的显著变换（例如层归一化），这可能不成立。一些模型，尤其是具有预归一化架构的模型，需要将最终层归一化应用于中间状态，但这并不总是足够的。
其次，logit lens对于具有大词汇表或需要多标记预测的任务效果较差，因为它仅显示下一个标记的分布。它还无法捕获模型的完整状态，因为它忽略了注意力和前馈子层除其加性效应之外的贡献。
第三，如果模型使用与残差流不对齐的单独输出头，该技术可能产生误导性结果。在这种情况下，投影的logits可能嘈杂或无信息。研究人员提出了诸如“调谐透镜”之类的变体，学习每层仿射变换以改善对齐，但这些需要额外的训练数据。
与其他可解释性方法的关系
logit lens是更广泛的可解释性技术家族的一部分，包括探针分类器、激活修补和因果追踪。探针分类器在隐藏状态上训练线性或非线性模型以预测特定特征，但它们需要监督且可能不反映模型的实际计算。相比之下，logit lens使用模型自身的输出头，因此更直接。
激活修补涉及干预隐藏状态以观察它们如何影响输出，可用于验证logit lens的发现。因果追踪由David Kaplan等人推广，使用损坏的运行来识别哪些层负责特定事实，而logit lens可以通过显示每层的预测标记来补充这一点。
另一个相关方法是“残差流”视图，将transformer视为一系列加性更新。logit lens自然适合此视图，因为它投影每个点的残差流。这启发了诸如“Transformer调试器”和其他可视化库等工具。
扩展与变体
已提出logit lens的几种扩展。“调谐透镜”学习每层仿射变换以更好地对齐隐藏状态与去嵌入，从而在标准logit lens失败的模型上提高性能。另一个变体是“softmax透镜”，对logits应用温度以锐化或平坦化分布，使低概率预测更易见。
对于具有多个输出头或解码器的模型，logit lens可以分别应用于每个头。在编码器-解码器模型中，该技术可应用于解码器的隐藏状态，但如何应用于编码器尚不清楚。
最近的工作还探索了使用logit lens分析多模态模型，其中词汇表包括图像标记或其他模态。然而，这仍是一个活跃的研究领域。
计算考虑
logit lens计算效率高，因为它仅需要一次前向传递和每层的矩阵乘法。对于具有\( L \)层和词汇表大小\( V \)的模型，额外成本为\( O(L \cdot V \cdot d) \)，其中\( d \)是隐藏维度。与前向传递本身的成本相比，这可以忽略不计。
然而，存储所有层的logits可能内存密集，尤其是对于大词汇表。在实践中，可以计算每层的top-k标记以减少内存使用。该技术在多个开源库中实现，包括TransformerLens库和Hugging Face可解释性工具。
未来方向
随着深度学习模型规模不断增长，像logit lens这样的可解释性工具变得越来越重要。未来的工作可能侧重于将该技术适应新架构，如混合专家或状态空间模型，并将其集成到自动化可解释性管道中。logit lens也正用于安全研究，以检测模型何时“思考”有害输出，这可能为对齐策略提供信息。
总体而言，logit lens仍然是窥视神经网络黑箱的简单而强大的工具，随着机制可解释性领域的成熟，其应用可能会扩大。