Convert a tensor of an imperative framework (i.e. numpy/cupy/torch/jax/etc.) to numpy.ndarray

Parse a tensor shape to dictionary mapping axes names to their lengths.

# Use underscore to skip the dimension in parsing. 
>>> x = np.zeros([2, 3, 5, 7]) >>> parse_shape(x, 'batch _ h w') {'batch': 2, 'h': 5, 'w': 7} 
# `parse_shape` output can be used to specify axes_lengths for other operations: 
>>> y = np.zeros([700]) 
>>> rearrange(y, '(b c h w) -> b c h w', **parse_shape(x, 'b _ h w')).shape 
(2, 10, 5, 7)

也就是把维度的维数映射到对应的命名。与数据无关，只看得到维度。

该工作主要梳理了LLM-based Agent 中的规划（planning）能力，原文链接：

Understanding the planning of LLM agents: A survey

文章中，作者将planning能力进一步细分为了五个维度：

gather

参数：

• input (Tensor) – the source tensor
• dim (int) – the axis along which to index
• index (LongTensor) – the indices of elements to gather
• out (Tensor_,__optional_) – the destination tensor
• sparse_grad (bool,optional) – If True, gradient w.r.t. input will be a sparse tensor. > gather操作是scatter操作的逆操作，如果说scatter是根据index和src求self(input)，那么gather操作是根据self(input)和index求src。具体来说gather操作是根据index指出的索引，沿dim指定的轴收集input的值。

out[i][j][k] = input[index[i][j][k]][j][k]  # if dim == 0
out[i][j][k] = input[i][index[i][j][k]][k]  # if dim == 1
out[i][j][k] = input[i][j][index[i][j][k]]  # if dim == 2

对于gather操作来说，有三个约束需要满足：

LLaMA介绍

LLaMA 是目前为止，效果最好的开源 LLM 之一。

论文的核心思想：相比于GPT，更小的模型+更多的训练数据**也可以获得可比的效果

基于更多 tokens 的训练集，在各种推理预算下，训练出性能最佳的一系列语言模型，称为 LLaMA，参数范围从 7B 到 65B 不等，与现有最佳 LLM 相比，其性能是有竞争力的。比如，LLaMA-13B 在大多数基准测试中优于 GPT-3，尽管其尺寸只有 GPT-3 的十分之一。作者相信，LLaMA 将有助于使 LLM 的使用和研究平民化，因为它可以在单个 GPU 上运行！在规模较大的情况下，LLaMA-65B 也具有与最佳大型语言模型（如 Chinchilla 或 PaLM-540B）相竞争的能力。

LLM解码时采用的自回归采样，其过程如下：

1. 小模型使用前缀作为输入，将输出结果处理+归一化成概率分布后，采样生成下一个token。
2. 将生成的token和前缀拼接成新的前缀，重复执行1，直到生成EOS或者达到最大token数目。

将模型输出logits的转换成概率，有几种常用的采样方法，包括argmax、top-k和top-n等 # 贪心搜索 直接选择概率最高的单词。这种方法简单高效，但是可能会导致生成的文本过于单调和重复 # 随机采样 按照概率分布随机选择一个单词。这种方法可以增加生成的多样性，但是可能会导致生成的文本不连贯和无意义。 # beam search 维护一个大小为 k 的候选序列集合，每一步从每个候选序列的概率分布中选择概率最高的 k 个单词，然后保留总概率最高的 k 个候选序列。这种方法可以平衡生成的质量和多样性，但是可能会导致生成的文本过于保守和不自然。 # top-k

线性Transformer

\[V_i'=\frac{\sum_{j=1}^N sim(Q_i,K_j)V_j}{\sum_{j=1}^N sim(Q_i,K_j)}\] 注意下标i。 其中

\[sim(Q_{i},K_{j})=\phi(Q_{i},K_{j})\]

此时有：

\[V_{i}^{\prime}=\frac{\phi(Q_{i})\sum_{j=1}^{i}\phi(K_{j})^{T}V_{j}}{\phi(Q_{i})\sum_{j=1}^{i}\phi(K_{j})^{T}}\]

注意可以将\(\phi(Q_{i})\)提出来。

原始Transformer的计算复杂度随序列长N呈二次方增长，这是因为attention的计算包含两层for循环，外层是对于每一个Query，我们需要计算它对应token的新表征；内层for循环是为了计算每一个Query对应的新表征，需要让该Query与每一个Key进行计算。 所以外层是 for q in Queries，内层是 for k in Keys。Queries数量和Keys数量都是N，所以复杂度是 O(N^2) 。而Linear Transformer，它只有外层for q in Queries这个循环了。因为求和项的计算与i无关，所以所有的 Qi 可以共享求和项的值。换言之，求和项的值可以只计算一次，然后存在内存中供所有 Qi 去使用。所以Linear Transformer的计算复杂度是O(N) 。

证明

核心思想就是找到一个转换，可以通过点积操作将位置信息注入，即： \[<f_q\left(x_m,m\right),f_k\left(x_n,n\right)>=g\left(x_m,x_n,m-n\right)\] 而通过复数的一些性质，找到了满足上述操作的转换：

\[\begin{aligned} &f_{q}\left(\boldsymbol{x}_{m},m\right)=\left(\boldsymbol{W}_{q}\boldsymbol{x}_{m}\right)e^{im\theta} \\ &f_{k}\left(\boldsymbol{x}_{n},n\right)=\left(\boldsymbol{W}_{k}\boldsymbol{x}_{n}\right)e^{in\theta} \\ &g\left(\boldsymbol{x}_{m},\boldsymbol{x}_{n},m-n\right)=\mathrm{Re}\left[\left(\boldsymbol{W}_{q}\boldsymbol{x}_{m}\right)\left(\boldsymbol{W}_{k}\boldsymbol{x}_{n}\right)^{*}e^{i(m-n)\theta}\right] \end{aligned}\] 可以发现g函数中存在相对位置信息。 欧拉公式：\(e^{ix}=\cos x+i\sin x\)

Data Engineering for Scaling Language Models to 128K Context

💡 Meta Data

📜 研究背景 & 基础 & 目的

论文主要研究了如何通过数据工程的方法，将语言模型的上下文长度扩展到128K个token。这项研究的重点在于数据工程，作者们提出了一个假设：长上下文建模的能力，特别是利用任意输入位置信息的能力，主要是通过大规模预训练获得的，并且这种能力可以通过轻量级的持续预训练在适当的数据混合上扩展到训练期间未见过的更长上下文（例如，从4K扩展到128K）。

KV cache

LLM推理过程分为Prefill和Decode两个阶段，其中Prefill阶段会对Prompt中所有的token做并行计算，得到Prompt中所有Tokens的KV Cache以及计算得到首Token。Prompt阶段Token计算得到的KV Cache会保存下来，留给Decode阶段复用，Decode阶段是一个自回归过程，每decode一个新的Token，都需要用到所有之前计算得到的KV Cache来计算当前query token的Attention。因此，当输出长度越来越大或者context很长时，KV Cache将会占用大量的显存。如何优化KV Cache的显存占用，一直都是LLM推理的核心主题之一。