seemore：从零开始实现一个视觉语言模型
TL;DR：在这篇博客中，我使用纯 PyTorch 从零开始实现了一个视觉语言模型，它包含一个图像编码器、一个多模态投影模块和一个解码器语言模型。你可以将其视为 Pixtral、GPT-4 或 Claude 3 中所展示的语言模型视觉能力的一个简化版本。名称“seemore”是我向 Andrej Karpathy 的项目“makemore”致敬的方式，因为我在这里使用了字符级别的自回归语言模型，非常类似于他的 nanoGPT/makemore 实现。我的目标是让你通过阅读这篇博客并逐步查看代码库中的代码，对它的工作原理有一个直观的理解。
GitHub 仓库提供了端到端实现：https://github.com/AviSoori1x/seemore
动机
由于 GPT-4、Grok 1.5、Claude 3 和 Google Gemini 所展示的能力，视觉语言模型已成为机器学习社区的重点关注话题。除了这些专有的多模态（主要是视觉-语言）模型之外，还有许多高性能的开源模型，例如 LLaVa、Kosmos（1 和 2）、微软的 GIT 以及最近 Hugging Face 的 Idefics2。
尽管“视觉语言模型”这个术语可能意味着很多东西，但目前这类模型浪潮的特点是它们能够根据图像和文本输入执行指令。本质上，你可以期待一个视觉语言模型为你写一首关于寿司有多棒的诗，同时还能根据给定的图像计算盘子上的寿司卷数量。我想明确这一点，因为还有许多其他类型的视觉语言模型，例如 CLIP 以及最近的变体 SigLIP，它们非常重要，但使用方式截然不同。事实上，我们将探讨这些架构中的组件如何用于当前的视觉语言模型。
为了这篇博客的目的，我将重点关注这种可以进行指令调优以执行有用任务的视觉语言模型。更具体地说，我将在这里阐述一种正在形成并被证明高度通用的常见架构模式。
通用架构
在“seemore”这个我的视觉语言模型（VLM）简单实现中，有3个主要组件。
图像编码器用于从图像中提取视觉特征。在这种情况下，我从零开始实现了 CLIP 中使用的原始视觉 Transformer。这实际上是许多现代 VLM 的流行选择。一个值得注意的例外是 Adept 的 Fuyu 系列模型，它将分块图像直接传递到投影层。
视觉-语言投影器 - 图像嵌入的形状与解码器使用的文本嵌入不同。因此，我们需要“投影”，即改变图像编码器提取的图像特征的维度，以匹配文本嵌入空间中的观察结果。这样，图像特征就变成了解码器的“视觉标记”。这可以是一个单层或一个 MLP。我使用了一个 MLP，因为它值得展示。
一个仅解码器的语言模型。这是最终生成文本的组件。在我的实现中，我通过将投影模块整合到我的解码器中，与 LLaVA 中的做法略有不同。通常不会这样做，并且会保持解码器（通常是已经预训练的模型）的架构不变。
因此，总而言之，图像编码器从给定图像中提取特征，将这些图像嵌入传递给视觉-语言投影器，该投影器将这些图像嵌入投影到文本嵌入空间，然后与文本输入中的文本嵌入连接起来，并由仅解码器语言模型自回归地生成文本。
当你放大来看时，它并没有那么复杂，老实说，相当巧妙。这种方法能奏效也挺令人惊叹的。就像深度学习中的其他一切一样。
让我们从图像编码器开始
如前所述，我选择在这里实现一个类似于 CLIP 中使用的视觉 Transformer。
最近的趋势表明，使用 SigLIP（CLIP 的改进版本，它使用 Sigmoid 损失代替 CLIP 对比学习任务中使用的交叉熵损失）中的视觉 Transformer，视觉语言模型取得了更好的性能。一个使用 SigLIP 视觉 Transformer 的微型视觉语言模型（总参数仅1.6B）表现远超其体量（字面意义上）的优秀例子是 vikhyat 的 moondream 2：https://github.com/vikhyat/moondream。
然而，为了简单起见，我们假设这里使用的是 CLIP 版本，但实现将是相同的。在 seemore 中，我使用与“[CLS]”标记对应的嵌入作为代表整个图像的特征向量。这样做是为了简单起见。然而，选择视觉 Transformer 最后一层的所有特征向量是可能的，而且可能更好。我的假设是这将有助于计数和 OCR 等任务，因为空间信息以更少压缩的方式提供给解码器。苹果最近发表的这篇论文中给出了一些相关的有价值的建议：https://arxiv.org/abs/2403.09611
为了从头开始实现这个视觉 Transformer，我们必须创建一个 PatchEmbeddings 类，它可以接收图像并创建一系列补丁。这个过程对于使 Transformer 架构有效地处理视觉数据至关重要，特别是利用架构后续步骤中的注意力块。这可以很简单地实现如下：
class PatchEmbeddings(nn.Module):
def __init__(self, img_size=96, patch_size=16, hidden_dim=512):
super().__init__()
# Store the input image size
self.img_size = img_size
# Store the size of each patch
self.patch_size = patch_size
# Calculate the total number of patches
self.num_patches = (img_size // patch_size) ** 2
# Create a convolutional layer to extract patch embeddings
# in_channels=3 assumes the input image has 3 color channels (RGB)
# out_channels=hidden_dim sets the number of output channels to match the hidden dimension
# kernel_size=patch_size and stride=patch_size ensure each patch is separately embedded
self.conv = nn.Conv2d(in_channels=3, out_channels=hidden_dim,
kernel_size=patch_size, stride=patch_size)
def forward(self, X):
# Extract patch embeddings from the input image
X = self.conv(X)
# Flatten the spatial dimensions (height and width) of the patch embeddings
# This step flattens the patch dimensions into a single dimension
X = X.flatten(2)
# Transpose the dimensions to obtain the shape [batch_size, num_patches, hidden_dim]
# This step brings the num_patches dimension to the second position
X = X.transpose(1, 2)
return X
在上面的代码中，输入图像通过卷积层被分解成 (img_size // patch_size) ** 2 个补丁，并投影成通道维度（在 PyTorch 实现中，3D 张量通常表示为 [B, T, C]，C 为通道维度）为 512 的向量。
视觉编码器和语言解码器之间的注意力机制
在构建 Transformer 块中的组件时，事情变得有趣起来，例如：注意力头实现、多头注意力、每个 Transformer 块中的多层感知器以及 Transformer 块本身。这些组件在用于“视觉标记”生成的视觉 Transformer 和用于实际文本输出生成的解码器语言模型中几乎是相同的。
唯一关键的区别在于解码器语言模型中每个注意力头应用的掩码。这样做是为了确保自回归语言生成过程的完整性，特别是在仅解码器模型中。这种掩码技术至关重要，因为它会遮蔽当前标记位置之后的所有信息，从而将模型的注意力仅引导到序列的先前部分。这种注意力机制被称为因果自注意力。
上图中，下三角掩码仅应用于解码器模型。请注意，在可视化视觉编码器中每个注意力头的过程时，矩阵 W 中亮蓝色三角形是不存在的。
因此，我以这样的方式实现了这些组件，通过向类构造函数传入一个 is_decoder 布尔参数，它们可以在视觉编码器和语言解码器之间共享。
因果自注意力和多头因果自注意力的代码可以组织如下。多头自注意力并行应用多个注意力头，每个头关注通道（嵌入维度）的不同部分。多头自注意力本质上改善了学习过程，并由于其固有的并行实现而提高了模型训练效率。请注意，我在此实现中使用了 dropout 进行正则化，即防止过拟合。
注意力头的实现如下：
class Head(nn.Module):
def __init__(self, n_embd, head_size, dropout=0.1, is_decoder=False):
super().__init__()
# Linear layer for key projection
self.key = nn.Linear(n_embd, head_size, bias=False)
# Linear layer for query projection
self.query = nn.Linear(n_embd, head_size, bias=False)
# Linear layer for value projection
self.value = nn.Linear(n_embd, head_size, bias=False)
# Dropout layer for regularization
self.dropout = nn.Dropout(dropout)
# Flag indicating whether this head is used in the decoder
self.is_decoder = is_decoder
def forward(self, x):
# Get the batch size (B), sequence length (T), and embedding dimension (C) from the input tensor
B, T, C = x.shape
# Compute key, query, and value projections
k = self.key(x) # Shape: [B, T, head_size]
q = self.query(x) # Shape: [B, T, head_size]
v = self.value(x) # Shape: [B, T, head_size]
# Compute attention scores by taking the dot product of query and key
# and scaling by the square root of the embedding dimension
wei = q @ k.transpose(-2, -1) * (C ** -0.5) # Shape: [B, T, T]
if self.is_decoder:
# If this head is used in the decoder, apply a causal mask to the attention scores
# to prevent attending to future positions
tril = torch.tril(torch.ones(T, T, dtype=torch.bool, device=x.device))
wei = wei.masked_fill(tril == 0, float('-inf'))
# Apply softmax to the attention scores to obtain attention probabilities
wei = F.softmax(wei, dim=-1) # Shape: [B, T, T]
# Apply dropout to the attention probabilities for regularization
wei = self.dropout(wei)
# Perform weighted aggregation of values using the attention probabilities
out = wei @ v # Shape: [B, T, head_size]
return out
多头注意力的实现如下：
class MultiHeadAttention(nn.Module):
def __init__(self, n_embd, num_heads, dropout=0.1, is_decoder=False):
super().__init__()
# Ensure that the embedding dimension is divisible by the number of heads
assert n_embd % num_heads == 0, "n_embd must be divisible by num_heads"
# Create a ModuleList of attention heads
self.heads = nn.ModuleList([
Head(n_embd, n_embd // num_heads, dropout, is_decoder)
for _ in range(num_heads)
])
# Linear layer for projecting the concatenated head outputs
self.proj = nn.Linear(n_embd, n_embd)
# Dropout layer for regularization
self.dropout = nn.Dropout(dropout)
def forward(self, x):
# Apply each attention head to the input tensor
head_outputs = [h(x) for h in self.heads]
# Concatenate the outputs from all heads along the last dimension
out = torch.cat(head_outputs, dim=-1)
# Apply the projection layer to the concatenated outputs
out = self.proj(out)
# Apply dropout to the projected outputs for regularization
out = self.dropout(out)
return out
每个多头注意力模块之后的多层感知器相当简单。请注意，我发现在视觉 Transformer 中 GELU 使用得相当频繁，而在文本 Transformer 中使用 ReLU，所以我用条件逻辑来根据 MLP 的插入位置在这两者之间切换。然而，由于其产生的模型性能，GELU 似乎被用于两者，尽管它的计算成本比 ReLU 高。
class MLP(nn.Module):
def __init__(self, n_embd, dropout=0.1, is_decoder=True):
super().__init__()
# Define the layers of the MLP
layers = [
# First linear layer that expands the input dimension from n_embd to 4 * n_embd
nn.Linear(n_embd, 4 * n_embd),
# Activation function: ReLU if is_decoder is True, else GELU
nn.ReLU() if is_decoder else nn.GELU(),
# Second linear layer that projects the intermediate dimension back to n_embd
nn.Linear(4 * n_embd, n_embd),
# Dropout layer for regularization
nn.Dropout(dropout)
]
# Create a sequential container to hold the layers
self.net = nn.Sequential(*layers)
def forward(self, x):
# Pass the input through the MLP layers
return self.net(x)
多头注意力和 MLP 可以组合成 Transformer 块。如前所述，`is_decoder` 布尔标志将允许我们开启和关闭掩码，从而轻松创建编码器和解码器块。
class Block(nn.Module):
def __init__(self, n_embd, num_heads, dropout=0.1, is_decoder=False):
super().__init__()
# Layer normalization for the input to the attention layer
self.ln1 = nn.LayerNorm(n_embd)
# Multi-head attention module
self.attn = MultiHeadAttention(n_embd, num_heads, dropout, is_decoder)
# Layer normalization for the input to the FFN
self.ln2 = nn.LayerNorm(n_embd)
# Feed-forward neural network (FFN)
self.ffn = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd), # Expand the dimension
nn.GELU(), # Activation function
nn.Linear(4 * n_embd, n_embd), # Project back to the original dimension
)
def forward(self, x):
original_x = x # Save the input for the residual connection
# Apply layer normalization to the input
x = self.ln1(x)
# Apply multi-head attention
attn_output = self.attn(x)
# Add the residual connection (original input) to the attention output
x = original_x + attn_output
# Apply layer normalization to the input to the FFN
x = self.ln2(x)
# Apply the FFN
ffn_output = self.ffn(x)
# Add the residual connection (input to FFN) to the FFN output
x = x + ffn_output
return x
组装视觉编码器
现在，分块逻辑和注意力块可以组合起来创建视觉 Transformer (ViT)。
class ViT(nn.Module):
def __init__(self, img_size, patch_size, num_hiddens, num_heads, num_blks, emb_dropout, blk_dropout):
super().__init__()
# Patch embedding layer to convert the input image into patches
self.patch_embedding = PatchEmbeddings(img_size, patch_size, num_hiddens)
# Learnable classification token
self.cls_token = nn.Parameter(torch.zeros(1, 1, num_hiddens))
# Calculate the number of patches
num_patches = (img_size // patch_size) ** 2
# Learnable position embedding
self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, num_hiddens))
# Dropout layer for the embeddings
self.dropout = nn.Dropout(emb_dropout)
# Stack of transformer blocks
self.blocks = nn.ModuleList([Block(num_hiddens, num_heads, blk_dropout, is_decoder=False) for _ in range(num_blks)])
# Layer normalization for the final representation
self.layer_norm = nn.LayerNorm(num_hiddens)
def forward(self, X):
# Convert the input image into patch embeddings
x = self.patch_embedding(X)
# Expand the classification token to match the batch size
cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
# Concatenate the classification token with the patch embeddings
x = torch.cat((cls_tokens, x), dim=1)
# Add the position embedding to the patch embeddings
x += self.pos_embedding
# Apply dropout to the embeddings
x = self.dropout(x)
# Pass the embeddings through the transformer blocks
for block in self.blocks:
x = block(x)
# Apply layer normalization to the final representation
x = self.layer_norm(x[:, 0])
return x
总的来说，ViT 类封装了视觉 Transformer 模型的架构和前向传播。它接收输入图像，将其转换为补丁嵌入，添加位置信息，并通过一系列 Transformer 块处理嵌入，以生成图像的有意义表示。返回的最终表示是与 CLS 标记对应的嵌入，然后用于在语言解码器中调节文本生成。
视觉-语言投影模块
然而，我们不能直接将其连接到文本嵌入。我们需要将其从视觉 Transformer 的图像嵌入维度投影到文本嵌入的维度。这通过视觉-语言投影器完成。如前所述，这可以是一个单一的可学习层，后接一个非线性激活，或者是一个 MLP。这里我实现了一个 MLP，原因有以下几点：
这是一个旨在理解 VLM 工作原理的实现。所以它比单个投影层更有趣。
目前有一个有趣的趋势是，在 VLM 训练阶段，预训练的视觉编码器和语言解码器都保持冻结。因此，为连接模块分配更多参数可以增强 VLM 的整体泛化能力，并有助于下游的指令微调过程。
这是该投影模块的实现。它与 Transformer 块中使用的 MLP 没有太大区别。
class MultiModalProjector(nn.Module):
def __init__(self, n_embd, image_embed_dim, dropout=0.1):
super().__init__()
# Define the projection network
self.net = nn.Sequential(
# Linear layer to expand the image embedding dimension
nn.Linear(image_embed_dim, 4 * image_embed_dim),
# GELU activation function
nn.GELU(),
# Linear layer to project the expanded image embeddings to the text embedding dimension
nn.Linear(4 * image_embed_dim, n_embd),
# Dropout layer for regularization
nn.Dropout(dropout)
)
def forward(self, x):
# Pass the input through the projection network
x = self.net(x)
return x
构建解码器语言模型
我们需要看的最后一个组件是解码器语言模型。在这里，我仍然保持了现代 VLM 架构的范围，但在实现上有所偏离。我已将投影模块集成到解码器模型类实现中。这是因为我从头开始构建所有内容，并希望保留 Andrej Karpathy 的 makemore 中的因果语言模型架构。在这种实现中，没有简单的方法可以直接输入重塑后的嵌入，所以我不得不做一些即兴创作。请记住，在使用 Hugging Face API 或任何其他允许使用预训练大型语言模型的现代库时，您可以直接将嵌入作为输入传递给模型（例如，使用 inputs_embeds 参数：https://huggingface.co/docs/transformers/en/model_doc/gpt2#transformers.GPT2Model.forward.inputs_embeds）。
也就是说，我在这里所做的是一个有趣的练习，它让你可以通过相当简单的代码看到：
图像嵌入如何使用视觉语言投影器重塑以匹配文本嵌入。
然后与标记嵌入连接。
随后与位置嵌入结合，最终用于计算损失函数（并最终生成文本）。
本质上，文本生成是以初始图像输入为条件的。这可以通过多种方式修改，以与交错的文本和图像配合使用，这对于使用微调 VLM 的多轮对话（即聊天场景）将非常有用。苹果的这篇论文中提供了许多有用的技巧：https://arxiv.org/pdf/2403.09611.pdf。
此解码器实现的关键部分如下所示。请注意，is_decoder 标志是如何设置为“True”的，以便使用自注意力块的掩码版本，从而在语言解码器中实现因果缩放点积自注意力。有关完整实现，请参阅上面链接的 GitHub 仓库。
class DecoderLanguageModel(nn.Module):
def __init__(self, n_embd, image_embed_dim, vocab_size, num_heads, n_layer, use_images=False):
super().__init__()
self.use_images = use_images
# Token embedding table
self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
# Position embedding table
self.position_embedding_table = nn.Embedding(1000, n_embd)
if use_images:
# Image projection layer to align image embeddings with text embeddings
self.image_projection = MultiModalProjector(n_embd, image_embed_dim)
# Stack of transformer decoder blocks
self.blocks = nn.Sequential(*[Block(n_embd, num_heads, is_decoder=True) for _ in range(n_layer)])
# Final layer normalization
self.ln_f = nn.LayerNorm(n_embd)
# Language modeling head
self.lm_head = nn.Linear(n_embd, vocab_size)
def forward(self, idx, image_embeds=None, targets=None):
# Get token embeddings from the input indices
tok_emb = self.token_embedding_table(idx)
if self.use_images and image_embeds is not None:
# Project and concatenate image embeddings with token embeddings
img_emb = self.image_projection(image_embeds).unsqueeze(1)
tok_emb = torch.cat([img_emb, tok_emb], dim=1)
# Get position embeddings
pos_emb = self.position_embedding_table(torch.arange(tok_emb.size(1), device=device)).unsqueeze(0)
# Add position embeddings to token embeddings
x = tok_emb + pos_emb
# Pass through the transformer decoder blocks
x = self.blocks(x)
# Apply final layer normalization
x = self.ln_f(x)
# Get the logits from the language modeling head
logits = self.lm_head(x)
if targets is not None:
if self.use_images and image_embeds is not None:
# Prepare targets by concatenating a dummy target for the image embedding
batch_size = idx.size(0)
targets = torch.cat([torch.full((batch_size, 1), -100, dtype=torch.long, device=device), targets], dim=1)
# Compute the cross-entropy loss
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-100)
return logits, loss
return logits
def generate(self, idx, image_embeds, max_new_tokens):
# Autoregressive character geneneration conditioned on visual token and preceding tokens. Refer to repo
将所有组件整合到 Seemore：简单的视觉语言模型中
现在我们已经有了三个关键组件，我们可以将它们整合到一个视觉语言模型中。完整的实现如下所示。如果你删除用于错误处理的 assert 语句，这看起来非常简单。回到博客开头我给出的概述，这里发生的一切是：
从视觉编码器获取图像特征（这里是视觉 Transformer，但它可以是任何能够从图像输入生成特征的模型，例如 ResNet 或传统的卷积神经网络（不用说性能可能会受到影响））。
一个投影模块，用于将图像标记投影到与解码器文本嵌入相同的嵌入空间（此实现中，该投影器已集成到解码器中）。
一个解码器语言模型，用于根据前面的图像生成文本。
class VisionLanguageModel(nn.Module):
def __init__(self, n_embd, image_embed_dim, vocab_size, n_layer, img_size, patch_size, num_heads, num_blks, emb_dropout, blk_dropout):
super().__init__()
# Set num_hiddens equal to image_embed_dim
num_hiddens = image_embed_dim
# Assert that num_hiddens is divisible by num_heads
assert num_hiddens % num_heads == 0, "num_hiddens must be divisible by num_heads"
# Initialize the vision encoder (ViT)
self.vision_encoder = ViT(img_size, patch_size, num_hiddens, num_heads, num_blks, emb_dropout, blk_dropout)
# Initialize the language model decoder (DecoderLanguageModel)
self.decoder = DecoderLanguageModel(n_embd, image_embed_dim, vocab_size, num_heads, n_layer, use_images=True)
def forward(self, img_array, idx, targets=None):
# Get the image embeddings from the vision encoder
image_embeds = self.vision_encoder(img_array)
# Check if the image embeddings are valid
if image_embeds.nelement() == 0 or image_embeds.shape[1] == 0:
raise ValueError("Something is wrong with the ViT model. It's returning an empty tensor or the embedding dimension is empty.")
if targets is not None:
# If targets are provided, compute the logits and loss
logits, loss = self.decoder(idx, image_embeds, targets)
return logits, loss
else:
# If targets are not provided, compute only the logits
logits = self.decoder(idx, image_embeds)
return logits
def generate(self, img_array, idx, max_new_tokens):
# Get the image embeddings from the vision encoder
image_embeds = self.vision_encoder(img_array)
# Check if the image embeddings are valid
if image_embeds.nelement() == 0 or image_embeds.shape[1] == 0:
raise ValueError("Something is wrong with the ViT model. It's returning an empty tensor or the embedding dimension is empty.")
# Generate new tokens using the language model decoder
generated_tokens = self.decoder.generate(idx, image_embeds, max_new_tokens)
return generated_tokens
现在我们已经实现了我们计划实现的所有内容
该仓库（此处：https://github.com/AviSoori1x/seemore）包含一些模拟数据，大部分从零开始实现的数据加载器，以及一个简单的带有交叉熵损失计算的训练循环。请注意，在这个简单的例子中，我们正在端到端地训练整个系统，非常类似于 Microsoft Research 的 Kosmos-1（https://arxiv.org/pdf/2302.14045.pdf）。为了方便起见，我将其保留在此处。实际上，通常观察到的序列是：
从 SigLIP 或 CLIP 获取预训练的视觉编码器（两者都有不同大小）。冻结权重（即在训练的反向传播中不更新）。
获取预训练的仅解码器语言模型，例如从 TinyLLaMA、Phi-2 等到 Llama 3（甚至更大，如 GPT-4 和 Grok 1.5 等）。冻结权重。
实现一个投影模块，并训练一个 VLM 模块，就像我们这里所做的，但只更新这个投影模块的权重。这实际上是预训练阶段。
然后在指令微调期间，将投影模块和解码器语言模型都解冻，并在反向传播中更新两者的权重。
我是在 Databricks 上使用单个 T4 GPU 和 MLFlow 跟踪损失（在训练过程中）开发的。有关深度学习中 MLFlow 的更多详细信息，请参见此处：https://mlflow.org.cn/blog/deep-learning-part-1）。我希望以这种方式设置，以便我可以轻松地在 Databricks 上扩展到任何大小的 GPU 集群，如果我决定将其应用于更注重性能的实现。但是，你可以在任何地方运行它，无论是否有 GPU。请注意，即使是只有 90 个样本的玩具训练循环，在 CPU 上也会非常慢。
请查看仓库（https://github.com/AviSoori1x/seemore），自己运行代码，玩得开心！
PS：还有一些新的方法，例如混合模态早期融合模型，例如 https://arxiv.org/abs/2405.09818。我计划将来实现一个简单的版本。
感谢阅读！
