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)