《DeepSeek大模型高性能核心技术与多模态融合开发(人工智能技术丛书)》(王晓华)【摘要 书评 试读】- 京东图书
图像特征token作为图文多模态特征之一,在视觉与语言的联合表示学习中扮演着至关重要的角色。它们能够有效地捕捉图像的局部细节和全局信息,为跨模态检索、视觉问答以及图像描述生成等任务提供了丰富的特征支持。
在图像特征token的研究中,我们不断探索如何更有效地提取、表示和利用这些特征。从最初的底层视觉特征,如边缘、纹理和颜色,到后来的高层语义特征,如物体、场景和动作,图像特征token的表达能力逐渐增强,为模型提供了更丰富和更准确的输入信息。
但是,随着数据规模的扩大和模型复杂度的提升,图像特征token的处理也面临着前所未有的挑战。如何在保证特征表达能力的同时,降低计算成本、提高处理效率,成为了当前研究的重要课题。
12.1.1 Pixel-Shuffle的token压缩
作为图像token压缩的一种高效策略,我们的核心目标在于直接削减token的数量,同时确保压缩后的token集合仍能充分承载原始图像的关键信息,维持信息的完整性和表达力。在这一理念的指引下,Pixel-Shuffle技术应运而生,成为实现这一目标的具体手段。
Pixel-Shuffle,顾名思义,是一种通过“像素重组”来实现数据维度转换的创新方法。其核心思想在于巧妙地利用通道与空间之间的转换关系,通过牺牲一定的空间分辨率来换取通道数的增加。这一转换过程不仅有效地降低了数据的空间维度,还使得每个token能够蕴含更为丰富和紧凑的图像特征,从而在减少token数目的同时,保证了图像信息的容量和表达质量。
具体来说,Pixel-Shuffle技术通过一系列精心设计的操作,将原始图像数据在通道维度上进行扩展,并在空间维度上进行相应的缩减。这种维度变换不仅有助于减少数据的冗余度,还能提升模型对关键特征的敏感度和捕捉能力。经过Pixel-Shuffle处理后的图像token,不仅数量得到了有效控制,而且每个token都蕴含了更为精炼和有价值的图像信息,为后续的模型处理提供了更优质和更高效的输入。其实现如下:
# 定义一个名为pixel_shuffle的函数,它接受一个输入张量x和一个缩放因子scale_factor,默认值为2。
def pixel_shuffle(x, scale_factor=2):
# 使用einops.rearrange函数将输入张量x从(batch, channel, height, width)格式重新排列为(batch, height, width, channel)格式。
x = einops.rearrange(x,"b c h w -> b h w c")
# 获取重新排列后张量的维度信息,分别是批次大小n,宽度w,高度h和通道数c。
n, w, h, c = x.size()
# 再次使用einops.rearrange函数将张量按照(scale_factor, scale_factor)的块进行重组,并将通道数扩展为原来的(scale_factor * scale_factor)倍。
x = einops.rearrange(x,"b h (w s) c -> b h w (c s)",s = scale_factor)
# 调整张量的维度顺序,使其变为(batch, width, height, channel)格式,并确保数据在内存中是连续的。
x = x.permute(0, 2, 1, 3).contiguous()
# 调整张量的形状,使其高度和宽度分别除以缩放因子,通道数乘以缩放因子的平方。
x = x.view(n, int(h // scale_factor), int(w // scale_factor),int(c * (scale_factor * scale_factor)))
# 再次调整张量的维度顺序,使其变为(batch, height, width, channel)格式,并确保数据在内存中是连续的。
x = x.permute(0, 2, 1, 3).contiguous()
# 最后,使用einops.rearrange函数将张量重新排列回(batch, channel, height, width)格式。
x = einops.rearrange(x,"b h w c-> b c h w")
# 返回处理后的张量。
return x
上面代码定义了一个名为pixel_shuffle的函数,它接受一个输入张量x和一个缩放因子scale_factor,通过重新排列和调整张量的形状,实现了像素洗牌操作,这是一种用于图像上采样的技术,可以增加图像的空间分辨率而不增加计算量。
下面采用Patch_Embedding的方式,来计算图像的token数。我们首先计算没有经过变换前的图像,之后对比经过重排的图像,完整代码如下所示。
import torch
import einops
class PatchEmbedding(torch.nn.Module):
def __init__(self, image_size, in_channels, patch_size = 14 , embed_dim = 312, dropout=0.):
super(PatchEmbedding, self).__init__()
# patch_embed相当于做了一个卷积
self.patch_embed = torch.nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size, bias=False)
self.drop = torch.nn.Dropout(dropout)
def forward(self, x):
# x[4, 3, 224, 224]
x = self.patch_embed(x)
# x [4, 16, 32, 32]
# x:[n,embed_dim,h',w']
x = x.flatten(2) # 将x拉直,h'和w'合并 [n,embed,h'*w'] #x [4, 16, 1024]
x = x.permute(0, 2, 1) # [n,h'*w',embed] #x [4, 1024, 16]
return x
if __name__ == '__main__':
image = torch.randn(size=(2,3,224,224))
image_token = PatchEmbedding(224,3)(image)
print(image_token.shape)
image = pixel_shuffle(image)
image_token = PatchEmbedding(112,12)(image)
print(image_token.shape)
结果如下:
torch.Size([2, 256, 312])
torch.Size([2, 64, 312])
对比未经过变换前的token化图像,经过转换后的图像在整体维度上缩小了3/4,而只有原有的1/4。其缩减的大小则是由缩放比例scale_factor得到,我们可以根据需要调整不同大小的scale_factor,从而获取到不同大小维度的图像token。
整体来看,Pixel-Shuffle技术的实现过程并不复杂,且具有较高的灵活性和可扩展性。它可以轻松地与其他图像处理技术相结合,形成更为强大和全面的图像处理流水线。此外,Pixel-Shuffle还具有良好的泛化能力,能够适用于不同类型的图像数据和深度学习模型,为图像token压缩领域带来了更为广阔的应用前景。
12.1.2 Cross-layer Token Fusion压缩
Cross-layer Token Fusion策略是一种高效的方法,它通过细致评估各个token对模型效率和准确性的综合贡献,巧妙地在特定网络层上实施token fusion,从而有效突破了传统模型的局限性。在多模态模型中,token的合并以及相似token的精准识别,均建立在余弦相似度的坚实基础之上,确保了融合的准确性和有效性。
Cross-layer Token Fusion的具体实现如下所示。
def bipartite_soft_matching(
metric: torch.Tensor, r: int,
class_token: bool = False, distill_token: bool = False,
):
protected = 0
if class_token:
protected += 1
if distill_token:
protected += 1
# 我们最多只能减少50%的令牌
t = metric.shape[1]
r = min(r, (t - protected) // 2)
if r <= 0:
return metric, metric
with torch.no_grad():
# 归一化相似度矩阵
metric = metric / metric.norm(dim=-1, keepdim=True)
a, b = metric[..., ::2, :], metric[..., 1::2, :]
scores = a @ b.transpose(-1, -2)
if class_token:
scores[..., 0, :] = -math.inf
if distill_token:
scores[..., :, 0] = -math.inf
# 找到匹配分数最高的令牌对
node_max, node_idx = scores.max(dim=-1)
edge_idx = node_max.argsort(dim=-1, descending=True)[..., None]
unm_idx = edge_idx[..., r:, :] # 未合并的令牌
src_idx = edge_idx[..., :r, :] # 要合并的令牌
dst_idx = node_idx[..., None].gather(dim=-2, index=src_idx)
if class_token:
# 确保类别令牌在开始位置
unm_idx = unm_idx.sort(dim=1)[0]
def merge(x: torch.Tensor, mode="mean") -> torch.Tensor:
"""
合并操作,将指定的令牌合并。
"""
src, dst = x[..., ::2, :], x[..., 1::2, :]
n, t1, c = src.shape
unm = src.gather(dim=-2, index=unm_idx.expand(n, t1 - r, c))
src = src.gather(dim=-2, index=src_idx.expand(n, r, c))
dst = dst.scatter_reduce(-2, dst_idx.expand(n, r, c), src, reduce=mode)
if distill_token:
return torch.cat([unm[:, :1], dst[:, :1], unm[:, 1:], dst[:, 1:]], dim=1)
else:
return torch.cat([unm, dst], dim=1)
def unmerge(x: torch.Tensor) -> torch.Tensor:
"""
取消合并操作,恢复原始令牌。
"""
unm_len = unm_idx.shape[1]
unm, dst = x[..., :unm_len, :], x[..., unm_len:, :]
n, _, c = unm.shape
src = dst.gather(dim=-2, index=dst_idx.expand(n, r, c))
out = torch.zeros(n, metric.shape[1], c, device=x.device, dtype=x.dtype)
out[..., 1::2, :] = dst
out.scatter_(dim=-2, index=(2 * unm_idx).expand(n, unm_len, c), src=unm)
out.scatter_(dim=-2, index=(2 * src_idx).expand(n, r, c), src=src)
return out
return merge, unmerge
其中,bipartite_soft_matching函数的输入参数解释如下:[yx1] [晓王2]
- metric:输入的hidden_state张量,尺寸为[batch, tokens, channels]。
- r:要移除的令牌数量,最多为总令牌数的50%。
- class_token(可选):布尔值,指示是否有类别令牌。默认为False。
- distill_token(可选):布尔值,指示是否有蒸馏令牌。默认为False。
在具体使用上,这个函数是一种用于缩减图像处理中的image_token令牌数量的方案。通过应用平衡匹配集(ToMe),它能够在减少令牌(token)数量的同时保持信息的完整性。下面是使用bipartite_soft_matching进行token缩减的代码示例:
# 创建一个随机embedding张量,用于模拟token被torch.Embedding类计算的结果
embedding = torch.randn(size=(2, 48, 312))
# 调用bipartite_soft_matching函数,指定要移除的令牌数量r为18,不包含类别令牌
merge, unmerge = bipartite_soft_matching(embedding, 18, class_token=False)
# 使用merge函数处理embedding
hidden_states = merge(embedding)
# 打印处理后的hidden_states的形状
print(hidden_states.shape) # 结果:torch.Size([2, 30, 312]),其中30 = 48 - 18
12.1.3 AvgPool的token压缩
AvgPoolProjector是一种用于取代corss_attention的图像token压缩算法,AvgPoolProjector通过自适应平均池化技术,在保留关键视觉信息的同时,有效缩减了图片token的数量。这一方法不仅简化了模型的复杂度,还提升了训练效率,使得在有限的计算资源下也能实现高效的视觉-文本模态对齐。更重要的是,由于其无参特性,AvgPoolProjector避免了繁琐的参数调优过程,且在实际应用中表现出色,为视觉与语言的跨模态理解提供了强有力的工具。
class AvgPoolProjector(nn.Module):
def __init__(
self,
layer_num: int = 2,
query_num: int = 36, #这里是输出的seq_length
mm_hidden_size: int = 384, #图片经过patch_embedding后的d_model,也就是输入的维度
llm_hidden_size: int = 384, #语言模型的d_model,也就是输出的维度
):
super().__init__()
self.layer_num = layer_num
self.query_num = query_num
self.mm_hidden_size = mm_hidden_size
self.llm_hidden_size = llm_hidden_size
self.build_net()
def build_net(self):
hw = int(self.query_num ** 0.5)
sampler = nn.AdaptiveAvgPool2d((hw, hw))
self.sampler = sampler
modules = [nn.Linear(self.mm_hidden_size, self.llm_hidden_size)]
for _ in range(1, self.layer_num):
modules.append(nn.GELU())
modules.append(nn.Linear(self.llm_hidden_size, self.llm_hidden_size))
modules.append(torch.nn.RMSNorm(self.llm_hidden_size))
self.mlp_projector = nn.Sequential(*modules)
def forward(self, visual_feat: torch.Tensor):
batch_size, seq_len, h_dim = visual_feat.shape # 576
# 计算平方根并断言它必须是整数
root = seq_len ** 0.5
assert float(int(root)) == root, f"{seq_len}的平方根不是整数"
hw = int(seq_len ** 0.5) # 24
shaped_visual_feat = rearrange(visual_feat, "b (h w) d -> b d h w", h=hw,w=hw) # torch.Size([64, 1024, 24, 24])
pooled_visual_feat = self.sampler(shaped_visual_feat) # torch.Size([64, 1024, 12, 12])
reshaped_visual_feat = rearrange(pooled_visual_feat, "b d h w -> b (h w) d") # [64, 144, 1024]
output_feat = self.mlp_projector(reshaped_visual_feat) # [64, 144, 4096])
return output_feat
其具体使用示例如下:
embedding = torch.randn(size=(2, 49, 384))
avg_pool = AvgPoolProjector()
pooled = avg_pool(embedding)
print(pooled.shape)
AvgPoolProjector的好处显而易见:它能够在减少计算负担的同时,保持模型的性能。通过直接在patch级别进行下采样,它有效地避免了语义信息的损失,确保了视觉特征与文本之间的准确对应。此外,其实现的简洁性也大大提升了模型的易用性和可扩展性。
总的来说,AvgPoolProjector以简洁高效的方式实现了视觉token的压缩,不仅提升了模型的训练效率和性能,还为视觉与文本的跨模态交互提供了新的可能性。无论是在资源受限的场景下,还是在追求高性能的应用中,AvgPoolProjector都展现出了其独特的优势和潜力。