#

评估 k 最近邻语言模型

11from typing import Optional, List
12
13import faiss
14import numpy as np
15import torch
16
17from labml import monit, lab
18from labml.logger import inspect
19from labml_nn.transformers.knn.train_model import Configs

#

$k$ -NN 要得到 $p (w_{t}, c_{t})$

这里我们称 $f (c_{t})$ 之为查询、 $f (c_{i})$ 键 $w_{i}$ 和值。

22def knn(queries: torch.Tensor, index: faiss.IndexFlatL2, keys_store: np.ndarray, vals_store: np.ndarray, n_tokens: int):

#

保存查询形状以重塑结果

31    queries_shape = queries.shape

#

展平查询的batch 和sequence 维度

34    queries = queries.view(-1, queries_shape[-1])

#

找到 10 个最近的邻居 $f (c_{i})$ 。 $f (c_{t})$ distance 是 FAISS 给出的距离idx ， $i$ 是其中的索引keys_store 。

38    distance, idx = index.search(queries.numpy(), 10)

#

得到 $f (c_{i})$

41    keys_found = queries.new_tensor(keys_store[idx])

#

得到 $w_{i}$

43    vals_found = torch.tensor(vals_store[idx]).squeeze(-1)

#

我们将计算归一化向量之间的余弦相似度

#

规范化 $f (c_{i})$

48    keys_found_n = keys_found / torch.sqrt((keys_found ** 2).sum(-1, keepdims=True) + 1e-10)

#

规范化 $f (c_{t})$

50    queries_n = queries / torch.sqrt((queries ** 2).sum(-1, keepdims=True) + 1e-10)

#

获取点积或余弦相似度

53    dot_prod = (keys_found_n * queries_n.unsqueeze(1)).sum(-1)

#

令牌明智的 logits

56    logits_token = dot_prod.new_zeros(queries.shape[0], n_tokens)

#

根据最近的邻居分散和累积令牌日志

58    _ = logits_token.scatter_(dim=1, index=vals_found, src=dot_prod, reduce='add')

#

重塑 logits

61    logits_token = logits_token.reshape(queries_shape[0], queries_shape[1], -1)
62
63    return logits_token

#

计算验证损失

我们计算了 $k$ -NN 预测和变压器预测组合的验证损失。分配给 $k$ -NN 模型的权重由给出knn_weight 。这是一个权重列表，我们计算每个权重的验证损失。

66def validation_loss(knn_weights: List[float], last_n: Optional[int], conf: Configs, index: faiss.IndexFlatL2,
67                    keys_store: np.ndarray, vals_store: np.ndarray):

#

每种的损失清单knn_weights

77    losses = [[] for _ in knn_weights]

#

每批样品的数量

79    n_samples = []
80    with torch.no_grad():

#

遍历验证数据

82        for i, batch in monit.enum("Validation", conf.validator.data_loader, is_children_silent=True):

#

获取数据和目标标签

84            data, target = batch[0].to(conf.device), batch[1].to(conf.device)

#

运行模型并获得预测 $p (w_{t}, c_{t})$

86            res = conf.model(data)

#

获取 $k$ -NN 预测

88            res_knn = knn(conf.model.ff_input.cpu(), index, keys_store, vals_store, conf.n_tokens)
89            res_knn = res_knn.to(conf.device)

#

这是为了只计算last_n 代币的损失。这一点很重要，因为变压器模型的第一个预测（沿顺序）几乎没有过去的令牌可供考虑。

94            if last_n:
95                res = res[-last_n:]
96                res_knn = res_knn[-last_n:]
97                target = target[-last_n:]

#

样本数量

100            n_s = res.shape[0] * data.shape[1]
101            n_samples.append(n_s)

#

计算每项的分数knn_weights 。

104            for i, c in enumerate(knn_weights):

#

计算损失

106                loss = conf.loss_func(res_knn * c + (1 - c) * res, target)
107                losses[i].append(loss * n_s)
108
109    return losses, n_samples

#

加载索引

112def load_index(conf: Configs, n_probe: int = 8):

#

的尺寸 $f (c_{i})$

117    d_model = conf.transformer.d_model

#

训练数据加载器

119    data_loader = conf.trainer.data_loader

#

上下文的数量；即训练数据中的令牌数减一。 $(f (c_{i}), w_{i})$ 对于 $i \in [2, T]$

122    n_keys = data_loader.data.shape[0] * data_loader.data.shape[1] - 1

#

加载 FAISS 指数

125    with monit.section('Load index'):
126        index = faiss.read_index(str(lab.get_data_path() / 'faiss.index'))

#

设置要探测的细胞数量

128    index.nprobe = n_probe

#

加载内存映射的 numpy 数组

131    keys_store = np.memmap(str(lab.get_data_path() / 'keys.npy'), dtype=np.float32, mode='r', shape=(n_keys, d_model))
132    vals_store = np.memmap(str(lab.get_data_path() / 'vals.npy'), dtype=np.int, mode='r', shape=(n_keys, 1))
133
134    return index, keys_store, vals_store

#

137def main():
138    from labml_nn.transformers.knn.build_index import load_experiment

#

加载实验。将 run uuid 替换为你在训练模型时运行的 uuid。

141    conf = load_experiment('4984b85c20bf11eb877a69c1a03717cd')

#

将模型设置为评估模式

143    conf.model.eval()

#

负荷指数

146    index, keys_store, vals_store = load_index(conf)

#

为

$k$ -NN 预测赋予的权重列表。我们将评估每个权重的验证损失

149    knn_weights = [i / 20 for i in range(10)]

#

评估验证损失

151    losses, n_samples = validation_loss(knn_weights, None, conf, index, keys_store, vals_store)

#

输出每个的损失knn_weights 。

153    inspect({c: np.sum(losses[i]) / np.sum(n_samples) for i, c in enumerate(knn_weights)})
154
155
156if __name__ == '__main__':
157    main()

评估 k 最近邻语言模型

k-NN 要得到p(wt​,ct​)

计算验证损失

加载索引

$k$ -NN 要得到 $p (w_{t}, c_{t})$