开源 LLM Attention 可视化库

Inspectus 是一款用途广泛的机器学习可视化工具。它提供了易于使用的 Python API，可以在 Jupyter Notebook 中流畅运行。

Attention 可视化：预览、组件、用法、教程、Huggingface 模型、自定义 attention map

教程：Huggingface 模型、自定义 attention map

Token 可视化：预览、用法

分布图：预览、用法、示例使用场景

开发环境设置

pip install inspectus

Attention 可视化

Inspectus 提供了一套用于深度学习模型 attention 机制的可视化工具。它提供多种完整的视图，让开发者更容易理解这些模型的工作原理。

单击某个 token 可以选中它并取消选中其他 token。再次单击，则会重新选中所有 token。如果只想改变某一个 token 的状态，请使用 Shift + 单击。

Attention Matrix：将 token 之间的 attention 分数可视化，突出展示在处理过程中每个 token 如何关注其他 token。

Query Token Heatmap：展示每个 query 与所选 key token 之间 attention 分数的总和。

Key Token Heatmap：展示每个 key 与所选 query token 之间 attention 分数的总和。

Dimension Heatmap：展示各个维度（Layer 和 Head）中每一项 attention 分数的总和，并在对应维度上进行归一化。

import inspectus

# attn: Attention map; a 2-4D tensor or attention maps from Huggingface transformers
inspectus.attention(attn, tokens)

对于 query token 和 key token 不同的情况：

inspectus.attention(attns, query_tokens, key_tokens)

有关详细的 API 文档，请参阅官方文档。

from transformers import AutoTokenizer, GPT2LMHeadModel, AutoConfig
import torch
import inspectus

# Initialize the tokenizer and model
context_length = 128
tokenizer = AutoTokenizer.from_pretrained("huggingface-course/code-search-net-tokenizer")

config = AutoConfig.from_pretrained(
    "gpt2",
    vocab_size=len(tokenizer),
    n_ctx=context_length,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
)

model = GPT2LMHeadModel(config)

# Tokenize the input text
text= 'The quick brown fox jumps over the lazy dog'
tokenized = tokenizer(
    text,
    return_tensors='pt',
    return_offsets_mapping=True
)
input_ids = tokenized['input_ids']

tokens = [text[s: e] for s, e in tokenized['offset_mapping'][0]]

with torch.no_grad():
    res = model(input_ids=input_ids.to(model.device), output_attentions=True)

# Visualize the attention maps using the Inspectus library
inspectus.attention(res['attentions'], tokens)

可在此处查看 Notebook：Huggingface 教程

import numpy as np
import inspectus

# 2D attention representing attention values between Query and Key tokens
attn = np.random.rand(3, 3)

# Visualize the attention values using the Inspectus library
# The first argument is the attention matrix
# The second argument is the list of query tokens
# The third argument is the list of key tokens
inspectus.attention(arr, ['a', 'b', 'c'], ['d', 'e', 'f'])

可在此处查看 Notebook：自定义 attention map 教程

Token 可视化

这个工具用于可视化与 token 相关的指标。它支持多种指标，可以通过下拉菜单选择用于可视化的指标。除了指标之外，还可以为 token 添加任意附加信息。

import inspectus
inspectus.tokens(["Hello", " World", "!"], {
    'loss': [-7.375, -8.312, -7.5625],
    'entropy': [3.323, 1.0666, 0.432]
}, token_info=["info 1", "info 2", "info 3"])

可在此处查看 Notebook：GPT2 示例

分布图

分布图用于展示一系列数据的分布情况。在每一步中，工具都会计算数据分布，并根据 9 个基准点最多绘制 5 个区间带。（0、6.68、15.87、30.85、50.00、69.15、84.13、93.32、100.00）

import inspectus

inspectus.distribution({'x': [x for x in range(0, 100)]})

可以使用 minimap 聚焦并放大图表中的特定部分。如需选择单个图表，请使用右上角的图例。

有关完整的使用指南，请查看此 Notebook：分布图教程。

示例使用场景

该图表可用于识别数据中是否存在离群值。以下 Notebook 演示了如何使用分布图识别 MNIST 训练 loss 中的离群值。

开发环境设置

如果你在学术研究中使用 Inspectus，请使用以下 BibTeX 条目引用该库。

@misc{inspectus,
 author = {Varuna Jayasiri, Lakshith Nishshanke},
 title = {inspectus: A visualization and analytics tool for large language models},
 year = {2024},
 url = {https://github.com/labmlai/inspectus},
}