SnapKV: LLM Knows What You are Looking for Before Generation（实现超长上下文的压缩方法无需训练）

2024-04-28 09:44
文章标签 实现 训练 方法 llm 压缩 无需 上下文 超长 generation looking knows snapkv

本文主要是介绍SnapKV: LLM Knows What You are Looking for Before Generation（实现超长上下文的压缩方法无需训练），希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

地址

https://arxiv.org/pdf/2404.14469

核心

这篇论文介绍了一种名为SnapKV的创新方法，旨在提高大型语言模型处理长上下文时的效率和内存利用率。主要贡献包括： 1. 设计实验探索在输出生成过程中注意力特征的模式，发现注意力分配具有一致性，可以提取重要信息。 2. 提出了SnapKV算法，利用观察窗口和投票机制选择每个注意力头的重要键值对，并使用池化进行细粒度聚类。 3. 在多个模型和数据集上评估SnapKV，结果显示其可以大幅压缩键值对缓存，提高解码速度，同时保持模型性能。 总之，SnapKV为长序列输入提供了一种高效压缩键值对缓存的方法，有助于降低内存和计算成本，同时保持了生成质量。

import torch
import time
import torch.nn.functional as F
import torch.nn as nn
import math# perform qk calculation and get indices
# this version will not update in inference mode# Copied from transformers.models.llama.modeling_llama.repeat_kv
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:"""This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)"""batch, num_key_value_heads, slen, head_dim = hidden_states.shapeif n_rep == 1:return hidden_stateshidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)class KVCluster():def __init__(self, window_size = 64, max_capacity_prompt = 256 + 64, kernel_size = 5, pooling = 'avgpool'):self.window_size = window_sizeself.max_capacity_prompt = max_capacity_promptassert self.max_capacity_prompt - self.window_size > 0self.kernel_size = kernel_sizeself.pooling = poolingdef reset(self, window_size = 64, max_capacity_prompt = 256 + 64, kernel_size = 5, pooling = 'avgpool'):self.window_size = window_sizeself.max_capacity_prompt = max_capacity_promptassert self.max_capacity_prompt - self.window_size > 0self.kernel_size = kernel_sizeself.pooling = poolingdef update_kv(self, key_states, query_states, value_states, attention_mask, num_key_value_groups):# check if prefix phaseassert key_states.shape[-2] == query_states.shape[-2]bsz, num_heads, q_len, head_dim = query_states.shapeif q_len < self.max_capacity_prompt:return key_states, value_stateselse:attn_weights = torch.matmul(query_states[..., -self.window_size:, :], key_states.transpose(2, 3)) / math.sqrt(head_dim)mask = torch.full((self.window_size, self.window_size), torch.finfo(attn_weights.dtype).min, device=attn_weights.device)mask_cond = torch.arange(mask.size(-1), device=attn_weights.device)mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)mask = mask.to(attn_weights.device)attention_mask = mask[None, None, :, :]attn_weights[:, :, -self.window_size:, -self.window_size:] += attention_maskattn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)attn_weights_sum = attn_weights[:, :, -self.window_size:, : -self.window_size].sum(dim = -2)if self.pooling == 'avgpool':attn_cache = F.avg_pool1d(attn_weights_sum, kernel_size = self.kernel_size, padding=self.kernel_size//2, stride=1)elif self.pooling == 'maxpool':attn_cache = F.max_pool1d(attn_weights_sum, kernel_size = self.kernel_size, padding=self.kernel_size//2, stride=1)else:raise ValueError('Pooling method not supported')indices = attn_cache.topk(self.max_capacity_prompt - self.window_size, dim=-1).indicesindices = indices.unsqueeze(-1).expand(-1, -1, -1, head_dim)k_past_compress = key_states[:, :, :-self.window_size, :].gather(dim = 2, index = indices)v_past_compress = value_states[:, :, :-self.window_size, :].gather(dim = 2, index = indices)k_cur = key_states[:, :, -self.window_size:, :]v_cur = value_states[:, :, -self.window_size:, :]key_states = torch.cat([k_past_compress, k_cur], dim = 2)value_states = torch.cat([v_past_compress, v_cur], dim = 2)return key_states, value_states

这段代码实现了一个KVCluster类，用于更新键值对（key-value pairs）。该类具有以下方法：

  • __init__(self, window_size = 64, max_capacity_prompt = 256 + 64, kernel_size = 5, pooling = 'avgpool')：初始化KVCluster对象，可以设置窗口大小、最大容量、卷积核大小和池化方法。

  • reset(self, window_size = 64, max_capacity_prompt = 256 + 64, kernel_size = 5, pooling = 'avgpool')：重置KVCluster对象的参数。

  • update_kv(self, key_states, query_states, value_states, attention_mask, num_key_value_groups)：根据输入的键、查询和值的状态以及注意力掩码更新键值对。在查询状态长度小于最大容量时，直接返回原始键值对。否则，根据注意力权重计算窗口内的聚合权重，并根据聚合权重选择top-k的索引。然后将过去的键值对和当前的键值对进行拼接，返回更新后的键值对。

该类主要用于处理键值对的更新，其中关键的部分是计算注意力权重、选择top-k索引和拼接过去和当前的键值对。

根据文档内容，使用SnapKV方法可以按照以下步骤进行：

  1. 确定观测窗口大小：选择输入序列末尾的一部分作为观测窗口，以捕获重要的注意力特征。通常选择窗口大小为32。
  2. 计算注意力权重：对观测窗口的查询和输入序列的前缀进行注意力计算，得到注意力权重矩阵。
  3. 进行投票：对每个注意力头，将观测窗口的注意力权重相加，选出权重最大的前缀位置作为重要特征。
  4. 聚类：对选出的重要特征进行聚类，以保留相邻特征。可以通过1D最大池化实现聚类。
  5. 更新KV缓存：将聚类后的特征与前缀特征拼接，形成新的Key-Value对，并更新KV缓存。这可以将KV缓存的大小压缩到指定值。
  6. 生成：使用更新后的KV缓存进行解码生成。由于KV缓存大小不再随输入序列增长，因此可以显著提高解码速度和内存效率。
  7. 调整参数：根据需要调整观测窗口大小、聚类核大小、KV缓存压缩目标值等参数，以平衡性能和效率。 总的来说，SnapKV通过自动识别输入序列中重要的注意力特征，并仅保留这些特征来压缩KV缓存，实现高效的长序列生成。该方法无需训练，可直接应用于现有模型中。

是的，根据文档中对SnapKV方法的描述，该步骤是在softmax之前进行的。 文档中提到SnapKV包含两个阶段：

  1. 投票阶段：计算观测窗口内查询和前缀的注意力权重，并进行投票，以选择出重要的前缀特征。
  2. 更新和存储阶段：根据投票结果，选择重要特征进行聚类，并拼接这些特征与前缀特征，形成新的KV对，以更新KV缓存。 这一过程发生在softmax之前，也就是在计算注意力权重时进行的。文档中并未明确指出是在softmax之前，但从上下文来看，这一过程发生在注意力权重计算阶段，因此是在softmax之前进行的。
    总之，SnapKV是在计算注意力权重时，通过压缩模型中提示的KV缓存来提高生成效率的，因此是在softmax之前进行的。

这篇关于SnapKV: LLM Knows What You are Looking for Before Generation（实现超长上下文的压缩方法无需训练）的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！


http://www.chinasem.cn/article/942919。

相关文章

Java中流式并行操作parallelStream的原理和使用方法

Java中流式并行操作parallelStream的原理和使用方法

《Java中流式并行操作parallelStream的原理和使用方法》本文详细介绍了Java中的并行流（parallelStream）的原理、正确使用方法以及在实际业务中的应用案例,并指出在使用并行流... 目录Java中流式并行操作parallelStream0. 问题的产生1. 什么是parallelS
阅读更多...
MySQL数据库双机热备的配置方法详解

MySQL数据库双机热备的配置方法详解

《MySQL数据库双机热备的配置方法详解》在企业级应用中,数据库的高可用性和数据的安全性是至关重要的,MySQL作为最流行的开源关系型数据库管理系统之一,提供了多种方式来实现高可用性,其中双机热备（M... 目录1. 环境准备1.1 安装mysql1.2 配置MySQL1.2.1 主服务器配置1.2.2 从
阅读更多...
C++中unordered_set哈希集合的实现

C++中unordered_set哈希集合的实现

《C++中unordered_set哈希集合的实现》std::unordered_set是C++标准库中的无序关联容器,基于哈希表实现,具有元素唯一性和无序性特点,本文就来详细的介绍一下unorder... 目录一、概述二、头文件与命名空间三、常用方法与示例1. 构造与析构2. 迭代器与遍历3. 容量相关4
阅读更多...
C++中悬垂引用(Dangling Reference) 的实现

C++中悬垂引用(Dangling Reference) 的实现

《C++中悬垂引用(DanglingReference)的实现》C++中的悬垂引用指引用绑定的对象被销毁后引用仍存在的情况,会导致访问无效内存,下面就来详细的介绍一下产生的原因以及如何避免,感兴趣... 目录悬垂引用的产生原因1. 引用绑定到局部变量，变量超出作用域后销毁2. 引用绑定到动态分配的对象，对象
阅读更多...
SpringBoot基于注解实现数据库字段回填的完整方案

SpringBoot基于注解实现数据库字段回填的完整方案

《SpringBoot基于注解实现数据库字段回填的完整方案》这篇文章主要为大家详细介绍了SpringBoot如何基于注解实现数据库字段回填的相关方法,文中的示例代码讲解详细,感兴趣的小伙伴可以了解... 目录数据库表pom.XMLRelationFieldRelationFieldMapping基础的一些代
阅读更多...
Java HashMap的底层实现原理深度解析

Java HashMap的底层实现原理深度解析

《JavaHashMap的底层实现原理深度解析》HashMap基于数组+链表+红黑树结构,通过哈希算法和扩容机制优化性能,负载因子与树化阈值平衡效率,是Java开发必备的高效数据结构,本文给大家介绍... 目录一、概述：HashMap的宏观结构二、核心数据结构解析1. 数组（桶数组）2. 链表节点（Node
阅读更多...
Java AOP面向切面编程的概念和实现方式

Java AOP面向切面编程的概念和实现方式

《JavaAOP面向切面编程的概念和实现方式》AOP是面向切面编程,通过动态代理将横切关注点（如日志、事务）与核心业务逻辑分离,提升代码复用性和可维护性,本文给大家介绍JavaAOP面向切面编程的概... 目录一、AOP 是什么？二、AOP 的核心概念与实现方式核心概念实现方式三、Spring AOP 的关
阅读更多...
Python版本信息获取方法详解与实战

Python版本信息获取方法详解与实战

《Python版本信息获取方法详解与实战》在Python开发中,获取Python版本号是调试、兼容性检查和版本控制的重要基础操作,本文详细介绍了如何使用sys和platform模块获取Python的主... 目录1. python版本号获取基础2. 使用sys模块获取版本信息2.1 sys模块概述2.1.1
阅读更多...
Python实现字典转字符串的五种方法

Python实现字典转字符串的五种方法

《Python实现字典转字符串的五种方法》本文介绍了在Python中如何将字典数据结构转换为字符串格式的多种方法,首先可以通过内置的str()函数进行简单转换;其次利用ison.dumps()函数能够... 目录1、使用json模块的dumps方法：2、使用str方法：3、使用循环和字符串拼接：4、使用字符
阅读更多...
Python版本与package版本兼容性检查方法总结

Python版本与package版本兼容性检查方法总结

《Python版本与package版本兼容性检查方法总结》：本文主要介绍Python版本与package版本兼容性检查方法的相关资料,文中提供四种检查方法,分别是pip查询、conda管理、PyP... 目录引言为什么会出现兼容性问题方法一：用 pip 官方命令查询可用版本方法二：conda 管理包环境方法
阅读更多...

Copyright China编程 23002807@qq.com

沪ICP备2023033072号-2