《Long Context Compression with Activation Beacon》笔记

Activation Beacon出自智源与人大在2024年1月放在arxiv上的论文《Long Context Compression with Activation Beacon》(v1版的题目：Soaring from 4K to 400K: Extending LLM’s Context with Activation Beacon）。它引入了Beacon token将上下文信息蒸馏到其激活(activations)；在压缩时将文本切分成固定大小的块(chunk)，并根据压缩比$\alpha$进一步将chunk分成更小的单元，beacon token插入在每个单元后面；LLM每次编码一个chunk，在自注意力机制执行过程中将chunk的信息蒸馏到beacon token的激活信息(activation)中，逐步地对整个长文本完成压缩过程，论文实验结果表明此方法可以有效加速推理过程并节省KV cache内存占用。

实现思路

如论文图1所示意，对输入文本$X=[x_1,\dots ,x_n]$，将其划分为相同尺寸w(如1024)的chunk：
$$[x_1,\dots ,x_n]\stackrel{\text{Partition}}{\to }[X_1,\dots X_{\lceil n/w\rceil }],X_i=[x_{(i-1)w+1},\dots ,x_{iw}]=[x_1^i,\dots ,x_w^i]$$
对每一个chunk $X_i$，使用一个压缩比${\alpha }_i$（w可由${\alpha }_i$整除），即将chunk划分到大小为$\alpha$的更细粒度单元，一组共$k_i=w/{\alpha }_i$个beacon token：$B_i=[⟨\mathbf{b}{⟩}_1^i,\dots ,⟨\mathbf{b}{⟩}_{k_i}^i]$被交替地插入到这些单元后。
$$X_i\stackrel{\text{Interleave}{B}_i}{\to }{X}_i^{\prime }=[x_1^i,\dots ,x_{{\alpha }_i}^i,⟨\mathbf{b}{⟩}_1^i,\dots ,x_{w-{\alpha }_i+1}^i,\dots ,x_w^i,⟨\mathbf{b}{⟩}_{k_i}^i]$$
LLM逐一地编码这些chunk，在自注意力机制过程中将每个chunk的信息压缩到beacon token的激活(activations)中，在编码了$X_i^{\prime }$后，将$X_i$的所有原始token(raw tokens)的激活信息给丢弃，但一直保留并累积beacon token $B_i$的激活信息；在编码下一个chunk $X_{i+1}^{\prime }$时，LLM将累积的beacon激活作为原始上下文$X_{\le i}$的代理。

如论文图2所示，Activation Beacon与一般的LLM相比只做少许修改。对于第i个chunk $X_i^{\prime }$，编码过程可以写作：
$$\mathrm{LLM}(\underset{\text{beacon activations accumulated from }{X}_{<i}^{\prime }}{\underbrace{⟨\mathbf{b}{⟩}_1^i,\dots ,⟨\mathbf{b}{⟩}_{k_{i-1}}^{i-1}}},\underset{\text{the current chunk }{X}_i^{\prime }}{\underbrace{x_1^i,\dots ,x_{{\alpha }_i}^i,⟨\mathbf{b}{⟩}_1^i,\dots ,x_{w-{\alpha }_i+1}^i,\dots ,x_w^i,⟨\mathbf{b}{⟩}_{k_i}^i}}),$$
也就是LLM的输入是前面chunk的激活累积和当前chunk需要被编码的token的混合物。设D表示LLM的隐藏层尺寸，$\mathbf{H}\in {\mathbb{R}}^{(w+k_i)\times D}$表示LLM任意层的self attention的输入隐藏状态。我们会区分raw token和beacon token：
$${\mathbb{I}}^r=\left\{j\mid x_j^i\ne ⟨\mathbf{b}⟩\right\},{\mathbb{I}}^b=\left\{j\mid x_j^i=⟨\mathbf{b}⟩\right\};{\mathbf{H}}^r=\mathbf{H}\left[{\mathbb{I}}^r\right],{\mathbf{H}}^b=\mathbf{H}\left[{\mathbb{I}}^b\right].$$
将隐状态变成query, key, value：
$$\begin{array}{lll}{\mathbf{Q}}^r={\mathbf{W}}_Q^r{\mathbf{H}}^r,& {\mathbf{K}}^r={\mathbf{W}}_K^r{\mathbf{H}}^r,& {\mathbf{V}}^r={\mathbf{W}}_V^r{\mathbf{H}}^r,\\ {\mathbf{Q}}^b={\mathbf{W}}_Q^b{\mathbf{H}}^b,& {\mathbf{K}}^b={\mathbf{W}}_K^b{\mathbf{H}}^b,& {\mathbf{V}}^b={\mathbf{W}}_V^b{\mathbf{H}}^b,\end{array}$$
上式中${\mathbf{W}}_{\ast }^r$是LLM原来的投影矩阵， ${\mathbf{W}}_{\ast }^b$是新引入的只处理beacon token的投影矩阵。再将raw token和beacon token的query/key/value状态来得到$\mathbf{Q},\mathbf{K},\mathbf{V}\in {\mathbb{R}}^{(w+k_i)\times D}$
$$\mathbf{Q}\left[{\mathbb{I}}^r\right]={\mathbf{Q}}^r,\mathbf{Q}\left[{\mathbb{I}}^b\right]={\mathbf{Q}}^b,\mathbf{K}\left[{\mathbb{I}}^r\right]={\mathbf{K}}^r,\mathbf{K}\left[{\mathbb{I}}^b\right]={\mathbf{K}}^b,\mathbf{V}\left[{\mathbb{I}}^r\right]={\mathbf{V}}^r,\mathbf{V}\left[{\mathbb{I}}^b\right]={\mathbf{V}}^b$$
最后，用标准方法计算self attention:
A = softmax ( mask ( Q { K a c ; K } T D ) ) , V = A { V a c ; V } \boldsymbol{A} = \text{softmax}\left(\text{mask} \left( \frac{\boldsymbol{Q}\{\boldsymbol{K}^{ac}; \boldsymbol{K} \}^T }{\sqrt{D}} \right)\right), \quad \boldsymbol{V} = \boldsymbol{A}\{\boldsymbol{V}^{ac};\boldsymbol{V} \} A=softmax(mask(D ​Q{Kac;K}T​)),V=A{Vac;V}
上式中的 { ⋅ ; ⋅ } \{ \cdot ; \cdot\} {⋅;⋅}表示矩阵连接，${\mathbf{K}}^{ac},{\mathbf{V}}^{ac}\in {\mathbb{R}}^{m_{i-1}\times D}$是从之前的chunk累积得到的beacon token的激活参数， $m_{i-1}={\sum }_{j=1}^{i-1}{k}_j$， mask就是causal attention mask。在self attention过程中，所有的token与其他token进行交互，使得beacon tokens的key和value(${\mathbf{K}}^b,{\mathbf{V}}^b$)蒸馏了$X_i$的上下文信息，它们会增量累积：
K a c = { K a c ; K b } , V a c = { V a c ; V b } \boldsymbol{K}^{ac} = \{\boldsymbol{K}^{ac}; \boldsymbol{K}^{b}\}, \boldsymbol{V}^{ac} = \{\boldsymbol{V}^{ac};\boldsymbol{V}^{b} \} Kac={Kac;Kb},Vac={Vac;Vb}

### 下面代码是activation beacon在实现时，interleave插入beacon token的代码，位于model_beacon.py的Memory类的_step函数
						input_len = input_ids.shape[1]
            if beacon_size > 0:
                # insert beacon tokens in between raw tokens，对应论文中的式(2)
                input_ids_with_beacons = input_ids.new_full((input_ids.shape[0], input_len + beacon_size), self.beacon_token.item())
                raw_token_indices = torch.arange(input_ids_with_beacons.shape[1], device=input_ids.device)
                interleave_start_idx = compression_ratio - self._interleave_remainder
                raw_token_indices = raw_token_indices[raw_token_indices % (compression_ratio + 1) != interleave_start_idx].unsqueeze(0).expand_as(input_ids)
                input_ids_with_beacons = input_ids_with_beacons.scatter(dim=1, index=raw_token_indices, src=input_ids)
                input_ids = input_ids_with_beacons
                # attention mask
                ## beacon token是参与attention的，所以默认值为1
                attention_mask_with_beacons = attention_mask.new_full((attention_mask.shape[0], attention_mask.shape[1] + beacon_size), 1)
                attention_mask_with_beacons = attention_mask_with_beacons.scatter(dim=1, index=raw_token_indices, src=attention_mask)
                attention_mask = attention_mask_with_beacons
                # labels
                if labels is not None:
                    ## beacon token不参与loss的计算，所以标签为-100
                    labels_with_beacons = labels.new_full((labels.shape[0], labels.shape[1] + beacon_size), -100)
                    labels_with_beacons = labels_with_beacons.scatter(dim=1, index=raw_token_indices, src=labels)
                    labels = labels_with_beacons

训练过程

Activation Beacon的学习目标是在当前chunk上下文和之前压缩信息的条件下提高生成质量，损失函数如下：
$$\underset{{\mathbf{\Theta }}^b}{\min }.\sum _{i=2}^{\lceil N/w\rceil }\sum _{j=1}^w\mathrm{Pr}\left(x_j^i\mid ⟨\mathbf{b}{⟩}_1^1,\dots ,⟨\mathbf{b}{⟩}_{k_{i-1}}^{i-1},x_1^i,\dots x_{j-1}^i;\mathbf{\Theta },{\mathbf{\Theta }}^b\right).$$
上式中$\mathbf{\Theta }$是LLM的参数，在训练过程中被冻结，${\mathbf{\Theta }}^{\mathbf{b}}$是每一层中beacon token对应的投影矩阵${\mathbf{W}}_{\ast }^b$和beacon token 的embedding ${\mathbf{e}}_{⟨b⟩}$（所有beacon token使用共享embedding )，训练时beacon token不参与损失计算（标签被设置为-100）因为它们仅用作压缩。

训练时第i个chunk的压缩比${\alpha }_i$是随机地从{2, 4, 8, 16, 32}中选取的，意在让模型灵活地支持不同的压缩粒度。而在推理时可以根据下游任务选择一个压缩比并应用到所有chunk。

训练过程分为预训练和微调，消融实验表明两个阶段对模型效果都有提升。

注意，activation beacon默认的方式是将其交替地插入在原始上下文中(代码中的interleave)，论文做消融实验时尝试将beacon token全部放在chunk的最后时效果是会下降的(代码中的append)。

注：Activation Beacon与MemoRAG是同一个团队出的，理解这篇思路之后，就能更好地理解MemoRAG的记忆模型了。(对比这两篇论文对应的记忆模型的代码，几乎是一样的，有点奇怪为什么memorag没有引用这篇文章，也没有对代码做说明。因为不理解memorag的记忆模型的代码，通过搜索关键字beacon搜到了这篇论文)。