Nagi-ovo

# image_encoder: 可以是 ResNet 或 Vision Transformer 等模型
# text_encoder:  可以是 CBOW 或 Text Transformer 等模型
# I[n, h, w, c]: 输入的图像 (minibatch)，包含 n 张图像，每张维度为 h*w*c
# T[n, l]:       输入的文本小批量，包含 n 段文本，每段长度为 l (通常是 token 数量)
# W_i[d_i, d_e]: 学习到的图像投影矩阵，将 d_i 维图像特征映射到 d_e 维嵌入空间
# W_t[d_t, d_e]: 学习到的文本投影矩阵，将 d_t 维文本特征映射到 d_e 维嵌入空间
# t:             学习到的温度参数 (一个标量)

I_f = image_encoder(I)  # 输出形状：[n, d_i]
T_f = text_encoder(T)  # 输出形状：[n, d_t]

# --- 计算联合多模态嵌入 ---
# 将图像特征通过投影矩阵 W_i 映射到 d_e 维空间
image_embedding_projected = np.dot(I_f, W_i)
# 将文本特征通过投影矩阵 W_t 映射到 d_e 维空间
text_embedding_projected = np.dot(T_f, W_t)
# 对投影后的图像和文本嵌入进行 L2 归一化 (使其模长为 1)，使点积等于余弦相似度
I_e = l2_normalize(image_embedding_projected, axis=1)  # 输出形状： [n, d_e]
T_e = l2_normalize(text_embedding_projected, axis=1)  # 输出形状： [n, d_e]

# --- 计算缩放后的成对余弦相似度 ---
# 计算 n 个图像嵌入和 n 个文本嵌入两两之间的点积。
# 由于 I_e 和 T_e 已经 L2 归一化，点积等价于余弦相似度。
# T_e.T 是 T_e 的转置，形状为 [d_e, n]。
# np.dot(I_e, T_e.T) 得到一个 [n, n] 的矩阵，(i, j) 位置的值是第 i 张图和第 j 段文本的余弦相似度。
# 乘以 np.exp(t) 对相似度进行缩放，t 是可学习的温度参数。
logits = np.dot(I_e, T_e.T) * np.exp(t)  # 输出形状：[n, n]

# --- 计算对称损失函数 ---
# 创建目标标签，对于一个包含 n 个配对样本的批次，正确的配对在 logits 矩阵的对角线上。正好存储了每一行（每一张图片）对应的正确文本的索引。
# labels = [0, 1, 2, ..., n-1]
labels = np.arange(n)

# 计算图像到文本的对比损失：
# 对 logits 矩阵（相似度矩阵）的每一行 (代表一张图像)，计算其与所有文本相似度得分的交叉熵损失。
# 目标是预测正确的文本索引 (即 labels)。axis=0 指示按行计算损失（不同框架实现可能略有差异，但概念如此）。
loss_i = cross_entropy_loss(logits, labels, axis=0)

# 计算文本到图像的对比损失：
# 对 logits 矩阵的每一列 (代表一段文本)，计算其与所有图像相似度得分的交叉熵损失。
# 目标是预测正确的图像索引 (即 labels)。axis=1 指示按列计算损失（概念上是这样）。
loss_t = cross_entropy_loss(logits, labels, axis=1)

# 最终的对称损失是两个方向损失的平均值。
loss = (loss_i + loss_t) / 2

softmax(x_i) = exp(x_i) / sum(exp(x_j))

safe_softmax(x_i) = exp(x_i - max(x)) / sum(exp(x_j - max(x)))

class SiglipVisionConfig:

    def __init__(
        self,
        hidden_size=768,                # 嵌入向量大小
        intermediate_size=3072,         # FFN 中线性层尺寸
        num_hidden_layers=12,           # ViT 层数
        num_attention_heads=12,         # MHA 的注意力头数
        num_channels=3,                 # 图片的通道数，即 RGB
        image_size=224,                 # 接收的图像特征尺寸
        patch_size=16,                  # 图片被分割的块数
        layer_norm_eps=1e-6,
        attention_dropout=0.0,
        num_image_tokens: int = None,   # 图像分割为 patch 后的 token 总数
        # (image_size / patch_size)^2 = (224 // 16) ** 2 = 196
        **kwargs
    ):
        super().__init__()

        self.hidden_size = hidden_size
        ...

class SiglipVisionEmbeddings(nn.Module):
    def __init__(self, config: SiglipVisionConfig):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.image_size = config.image_size
        self.patch_size = config.patch_size

        self.patch_embedding = nn.Conv2d(
            in_channels=config.num_channels,
            out_channels=self.embed_dim,
            kernel_size=self.patch_size,
            stride=self.patch_size,
            padding="valid", # This indicates no padding is added
        )

        self.num_patches = (self.image_size // self.patch_size) ** 2 # 对应前面
        self.num_positions = self.num_patches
        self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
        self.register_buffer(
            "position_ids",
            torch.arange(self.num_positions).expand((1, -1)),
            persistent=False, # 该 buffer 是否作为模型的永久部分保存
        )

    def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
        _, _, height, width = pixel_values.shape # [Batch_Size, Channels, Height, Width]
        # Convolve the `patch_size` kernel over the image, with no overlapping patches since the stride is equal to the kernel size
        # The output of the convolution will have shape [Batch_Size, Embed_Dim, Num_Patches_H, Num_Patches_W]
        # where Num_Patches_H = height // patch_size and Num_Patches_W = width // patch_size
        patch_embeds = self.patch_embedding(pixel_values)  
        # [Batch_Size, Embed_Dim, Num_Patches_H, Num_Patches_W] -> [Batch_Size, Embed_Dim, Num_Patches]
        # where Num_Patches = Num_Patches_H * Num_Patches_W
        embeddings = patch_embeds.flatten(2)
        # [Batch_Size, Embed_Dim, Num_Patches] -> [Batch_Size, Num_Patches, Embed_Dim]
        embeddings = embeddings.transpose(1, 2)
        # Add position embeddings to each patch. Each positional encoding is a vector of size [Embed_Dim]
        embeddings = embeddings + self.position_embedding(self.position_ids)
        # [Batch_Size, Num_Patches, Embed_Dim]
        return embeddings

class SiglipEncoderLayer(nn.Module):
    def __init__(self, config: SiglipVisionConfig):
        super().__init__()
        self.embed_dim = config.hidden_size
        self.self_attn = SiglipAttention(config)
        self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
        self.mlp = SiglipMLP(config)
        self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

    # Ignore copy
    def forward(
        self,
        hidden_states: torch.Tensor
    ) -> torch.Tensor:
        # residual: [Batch_Size, Num_Patches, Embed_Dim]
        residual = hidden_states
        # [Batch_Size, Num_Patches, Embed_Dim] -> [Batch_Size, Num_Patches, Embed_Dim]
        hidden_states = self.layer_norm1(hidden_states)
        # [Batch_Size, Num_Patches, Embed_Dim] -> [Batch_Size, Num_Patches, Embed_Dim]
        hidden_states, _ = self.self_attn(hidden_states=hidden_states)
        # [Batch_Size, Num_Patches, Embed_Dim]
        hidden_states = residual + hidden_states
        # residual: [Batch_Size, Num_Patches, Embed_Dim] 
        residual = hidden_states
        # [Batch_Size, Num_Patches, Embed_Dim] -> [Batch_Size, Num_Patches, Embed_Dim]
        hidden_states = self.layer_norm2(hidden_states)
        # [Batch_Size, Num_Patches, Embed_Dim] -> [Batch_Size, Num_Patches, Embed_Dim]
        hidden_states = self.mlp(hidden_states)
        # [Batch_Size, Num_Patches, Embed_Dim]
        hidden_states = residual + hidden_states
        
        return hidden_states

class SiglipEncoder(nn.Module):
    def __init__(self, config: SiglipVisionConfig):
        super().__init__()
        self.config = config
        self.layers = nn.ModuleList(
            [SiglipEncoderLayer(config) for _ in range(config.num_hidden_layers)]
        )

    # Ignore copy
    def forward(
        self,
        inputs_embeds: torch.Tensor
    ) -> torch.Tensor:
        # inputs_embeds: [Batch_Size, Num_Patches, Embed_Dim]
        hidden_states = inputs_embeds

        for encoder_layer in self.layers:
            # [Batch_Size, Num_Patches, Embed_Dim] -> [Batch_Size, Num_Patches, Embed_Dim]
            hidden_states = encoder_layer(hidden_states)

        return hidden_states

class SiglipVisionTransformer(nn.Module):
    def __init__(self, config: SiglipVisionConfig):
        super().__init__()
        self.config = config
        embed_dim = config.hidden_size

        self.embeddings = SiglipVisionEmbeddings(config)
        self.encoder = SiglipEncoder(config)
        self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
        # pixel_values: [Batch_Size, Channels, Height, Width] -> [Batch_Size, Num_Patches, Embed_Dim]
        hidden_states = self.embeddings(pixel_values)

        last_hidden_state = self.encoder(inputs_embeds=hidden_states)

        last_hidden_state = self.post_layernorm(last_hidden_state)

        return last_hidden_state

class SiglipVisionModel(nn.Module):

    def __init__(self, config: SiglipVisionConfig):
        super().__init__()
        self.config = config
        self.vision_model = SiglipVisionTransformer(config)

    def forward(self, pixel_values) -> Tuple:
        # [Batch_Size, Channels, Height, Width] -> [Batch_Size, Num_Patches, Embed_Dim]
        return self.vision_model(pixel_values=pixel_values) 

class PaliGemmaProcessor:

    IMAGE_TOKEN = "<image>"

    def __init__(self, tokenizer, num_image_tokens: int, image_size: int):
        super().__init__()

        self.image_seq_length = num_image_tokens
        self.image_size = image_size
		# tokenizer.add_tokens(EXTRA_TOKENS)
        # Tokenizer described here: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/paligemma/README.md#tokenizer
        tokens_to_add = {"additional_special_tokens": [self.IMAGE_TOKEN]}
        tokenizer.add_special_tokens(tokens_to_add)
        self.image_token_id = tokenizer.convert_tokens_to_ids(self.IMAGE_TOKEN)
        # We will add the BOS and EOS tokens ourselves
        tokenizer.add_bos_token = False
        tokenizer.add_eos_token = False

        self.tokenizer = tokenizer

    def __call__(
        self,
        text: List[str],            # 输入的文本 prompt
        images: List[Image.Image],  # 输入的 PIL.Image
        # 标准的分词器填充和截断策略
        padding: str = "longest",
        truncation: bool = True,
    ) -> dict:
        assert len(images) == 1 and len(text) == 1, f"Received {len(images)} images for {len(text)} prompts." # 这里只实现了一对图像——文本

        pixel_values = process_images(
            images,
            size=(self.image_size, self.image_size),
            resample=Image.Resampling.BICUBIC,  # 选择双三次插值算法进行图像重采样，在调整图像尺寸时保留更多细节与纹理，提高视觉质量
            rescale_factor=1 / 255.0,           # 像素值重缩放到 [0, 1]
            image_mean=IMAGENET_STANDARD_MEAN,  # 标准化均值 (近似 [0.5, 0.5, 0.5]) 
            image_std=IMAGENET_STANDARD_STD,    # 标准化标准差 (近似 [0.5, 0.5, 0.5])
        )
        # 这会将像素值转换到 [-1, 1]
        # Convert the list of numpy arrays to a single numpy array with shape [Batch_Size, Channel, Height, Width]
        pixel_values = np.stack(pixel_values, axis=0) # 堆叠成一个 batch
        # Convert the numpy array to a PyTorch tensor
        pixel_values = torch.tensor(pixel_values)

   图像 (PIL.Image) → 调整大小 → NumPy 数组 → 像素值缩放 (0-1) → 标准化 (均值/标准差) → 转置通道维度 → PyTorch 张量

        # Prepend a `self.image_seq_length` number of image tokens to the prompt
        input_strings = [
            add_image_tokens_to_prompt(
                prefix_prompt=prompt,
                bos_token=self.tokenizer.bos_token,
                image_seq_len=self.image_seq_length, # <image>占位符的数量
                image_token=self.IMAGE_TOKEN,        # "<image>"字符串
            )
            for prompt in text
        ]

        # 将包含<image>占位符的字符串序列转换为模型可读的 input_ids 和 attention_mask
        inputs = self.tokenizer(
            input_strings,
            return_tensors="pt",
            padding=padding,
            truncation=truncation,
        )
		# 返回一个字典，包含处理好的图像张量和文本相关的 token 序列
        return_data = {"pixel_values": pixel_values, **inputs}

        return return_data```

为了便于理解，这里用一个例子$$^{[2]}$$来可视化最后的形态：

```python
from transformers.utils.attention_visualizer import AttentionMaskVisualizer

visualizer = AttentionMaskVisualizer("google/paligemma2-3b-mix-224")
visualizer("<img> What is in this image?")

class PaliGemmaForConditionalGeneration(nn.Module):
    def __init__(self, config: PaliGemmaConfig):
        super().__init__()
        self.config = config
        self.vision_tower = SiglipVisionModel(config.vision_config)
        self.multi_modal_projector = PaliGemmaMultiModalProjector(config)
        self.vocab_size = config.vocab_size

        language_model = GemmaForCausalLM(config.text_config)
        self.language_model = language_model

        self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1

	# weight tying
    def tie_weights(self):
        return self.language_model.tie_weights()

def forward(
        self,
        input_ids: torch.LongTensor = None,
        pixel_values: torch.FloatTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        kv_cache: Optional[KVCache] = None,
    ) -> Tuple:

        # Make sure the input is right-padded
        assert torch.all(attention_mask == 1), "The input cannot be padded"

        # 1. Extra the input embeddings
        # shape: (Batch_Size, Seq_Len, Hidden_Size)
        inputs_embeds = self.language_model.get_input_embeddings()(input_ids)

        # 2. Merge text and images
        # [Batch_Size, Channels, Height, Width] -> [Batch_Size, Num_Patches, Embed_Dim]
        selected_image_feature = self.vision_tower(pixel_values.to(inputs_embeds.dtype))
 

		# [Batch_Size, Num_Patches, Embed_Dim] -> [Batch_Size, Num_Patches, Hidden_Size]
		image_features = self.multi_modal_projector(selected_image_feature)

class PaliGemmaMultiModalProjector(nn.Module):
    def __init__(self, config: PaliGemmaConfig):
        super().__init__()
        self.linear = nn.Linear(config.vision_config.hidden_size, config.vision_config.projection_dim, bias=True)

    def forward(self, image_features):
        # [Batch_Size, Num_Patches, Embed_Dim] -> [Batch_Size, Num_Patches, Projection_Dim]
        hidden_states = self.linear(image_features)
        return hidden_states

        # 3. Merge the embeddings of the text tokens and the image tokens
        inputs_embeds, attention_mask, position_ids = self._merge_input_ids_with_image_features(image_features, inputs_embeds, input_ids, attention_mask, kv_cache)

# Shape: [Batch_Size, Seq_Len]. True for image tokens
image_mask = input_ids == self.config.image_token_index # "<image>": 257152

# Shape: [Batch_Size, Seq_Len]. True for text tokens
text_mask = (input_ids != self.config.image_token_index) & (input_ids != self.pad_token_id)

text_mask_expanded = text_mask.unsqueeze(-1).expand(-1, -1, embed_dim)
image_mask_expanded = image_mask.unsqueeze(-1).expand(-1, -1, embed_dim)

# 放置文本嵌入
final_embedding = torch.where(text_mask_expanded, inputs_embeds, final_embedding)

# 放置图像嵌入
final_embedding = final_embedding.masked_scatter(image_mask_expanded, scaled_image_features)

		# 4. 输入语言模型
        outputs = self.language_model(
            attention_mask=attention_mask,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            kv_cache=kv_cache,
        )

        return outputs

class GemmaForCausalLM(nn.Module):
    # ... (其他代码) ...

    def forward(
        self,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        kv_cache: Optional[KVCache] = None,
    ) -> Tuple:

        # ... (模型前向传播) ...
        outputs = self.model( # self.model 是 GemmaModel
            attention_mask=attention_mask,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            kv_cache=kv_cache,
        )

        hidden_states = outputs # 这是来自 GemmaModel 的最后一个隐藏状态
        logits = self.lm_head(hidden_states) # lm_head 是一个线性层，输出维度为词汇表大小
        logits = logits.float()

        return_data = {
            "logits": logits,
        }

        if kv_cache is not None:
            # 返回更新后的缓存
            return_data["kv_cache"] = kv_cache

        return return_data```
## Prospectives

- VLM 如何更有效地利用图片信息，而不会忽略图片只关注文本
- VLM + RL，提升对于图片中信息的推理能力以完成视觉推理任务，及其可解释性

## Reference

- \[1] [Generalized Visual Language Models | Lil'Log](https://lilianweng.github.io/posts/2022-06-09-vlm/)
- \[2] [PaliGemma – Google's Cutting-Edge Open Vision Language Model](https://huggingface.co/blog/paligemma)
- \[3] [Transformers/PaliGemma](https://huggingface.co/docs/transformers/main/en//model_doc/paligemma)
- \[4] [Groundlight.ai: How does a Vision-Language-Model (VLM) work?](https://www.groundlight.ai/blog/how-vlm-works-tokens)