Upload folder using huggingface_hub

config.json

@@ -0,0 +1,26 @@
+{
+  "_name_or_path": "laion/CLIP-ViT-B-32-laion2B-s34B-b79K",
+  "architectures": [
+    "BayesVLMModel"
+  ],
+  "auto_map": {
+    "AutoModel": "modeling_bayesvlm_clip.BayesVLMModel",
+    "AutoProcessor": "transformers.CLIPProcessor"
+  },
+  "initializer_factor": 1.0,
+  "logit_scale_init_value": 2.6592,
+  "model_type": "clip",
+  "projection_dim": 512,
+  "text_config": {
+    "dropout": 0.0,
+    "hidden_act": "gelu",
+    "model_type": "clip_text_model"
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.40.2",
+  "vision_config": {
+    "dropout": 0.0,
+    "hidden_act": "gelu",
+    "model_type": "clip_vision_model"
+  }
+}

modeling_bayesvlm_clip.py

@@ -0,0 +1,495 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from transformers import CLIPModel, CLIPTextModelWithProjection, CLIPVisionModelWithProjection
+from transformers.modeling_outputs import ModelOutput
+
+
+def _as_optional_tensor(tensor: torch.Tensor | None) -> torch.Tensor | None:
+    return tensor if tensor is not None else None
+
+
+def _diag_cov(
+    activations: torch.Tensor,
+    a_inv: torch.Tensor,
+    b_diag: torch.Tensor,
+    add_bias: bool,
+) -> torch.Tensor | None:
+    if a_inv.numel() == 0 or b_diag.numel() == 0:
+        return None
+
+    if add_bias:
+        ones = torch.ones_like(activations[:, :1])
+        activations = torch.cat([activations, ones], dim=-1)
+
+    quad = torch.einsum("ij,jk,ik->i", activations, a_inv, activations)[:, None]
+    return quad * b_diag
+
+
+def _std_from_var(var: torch.Tensor | None) -> torch.Tensor | None:
+    if var is None:
+        return None
+    return torch.sqrt(var)
+
+def _get_output(outputs, name: str, index: int):
+    if hasattr(outputs, name):
+        return getattr(outputs, name)
+    if isinstance(outputs, (tuple, list)) and len(outputs) > index:
+        return outputs[index]
+    return None
+
+def _normalize_mean_and_var(
+    mean: torch.Tensor,
+    var: torch.Tensor,
+    eps: float = 1e-6,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    r2 = (mean**2).sum(dim=-1, keepdim=True).clamp_min(eps)
+    r = torch.sqrt(r2)
+    normalized = mean / r
+
+    # Delta-method approximation with diagonal covariance.
+    y2 = normalized**2
+    sum_y2v = (y2 * var).sum(dim=-1, keepdim=True)
+    norm_var = (var - 2 * y2 * var + y2 * sum_y2v) / r2
+    norm_var = norm_var.clamp_min(0)
+    return normalized, norm_var
+
+
+@dataclass
+class BayesVLMEmbeddingOutput(ModelOutput):
+    mean: torch.FloatTensor | None = None
+    var: torch.FloatTensor | None = None
+    std: torch.FloatTensor | None = None
+
+
+@dataclass
+class BayesVLMTextModelOutput(ModelOutput):
+    text_embeds: torch.FloatTensor | None = None
+    text_embeds_var: torch.FloatTensor | None = None
+    text_embeds_std: torch.FloatTensor | None = None
+    last_hidden_state: torch.FloatTensor | None = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class BayesVLMVisionModelOutput(ModelOutput):
+    image_embeds: torch.FloatTensor | None = None
+    image_embeds_var: torch.FloatTensor | None = None
+    image_embeds_std: torch.FloatTensor | None = None
+    last_hidden_state: torch.FloatTensor | None = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class BayesVLMOutput(ModelOutput):
+    loss: torch.FloatTensor | None = None
+    logits_per_image: torch.FloatTensor | None = None
+    logits_per_text: torch.FloatTensor | None = None
+    logits_per_image_var: torch.FloatTensor | None = None
+    logits_per_text_var: torch.FloatTensor | None = None
+    logits_per_image_std: torch.FloatTensor | None = None
+    logits_per_text_std: torch.FloatTensor | None = None
+    text_embeds: torch.FloatTensor | None = None
+    image_embeds: torch.FloatTensor | None = None
+    text_embeds_var: torch.FloatTensor | None = None
+    image_embeds_var: torch.FloatTensor | None = None
+    text_embeds_std: torch.FloatTensor | None = None
+    image_embeds_std: torch.FloatTensor | None = None
+    text_model_output: Optional[ModelOutput] = None
+    vision_model_output: Optional[ModelOutput] = None
+
+
+class BayesVLMTextModel(CLIPTextModelWithProjection):
+    def __init__(self, config):
+        super().__init__(config)
+        hidden = int(config.hidden_size)
+        proj = int(config.projection_dim)
+        self.register_buffer("a_inv", torch.zeros(hidden, hidden))
+        self.register_buffer("b_diag", torch.zeros(proj))
+
+    def set_covariance(self, a_inv: torch.Tensor, b_inv: torch.Tensor) -> None:
+        self.a_inv = a_inv
+        self.b_diag = torch.diagonal(b_inv)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if not return_dict:
+            return super().forward(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        text_outputs = self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        pooled_output = _get_output(text_outputs, "pooler_output", 1)
+        last_hidden_state = _get_output(text_outputs, "last_hidden_state", 0)
+        hidden_states = _get_output(text_outputs, "hidden_states", 2)
+        attentions = _get_output(text_outputs, "attentions", 3)
+        text_embeds = self.text_projection(pooled_output)
+
+        text_var = _diag_cov(
+            pooled_output,
+            self.a_inv,
+            self.b_diag,
+            add_bias=self.text_projection.bias is not None,
+        )
+        if text_var is None:
+            text_var = torch.zeros_like(text_embeds)
+        text_std = _std_from_var(text_var)
+
+        return BayesVLMTextModelOutput(
+            text_embeds=text_embeds,
+            text_embeds_var=text_var,
+            text_embeds_std=text_std,
+            last_hidden_state=last_hidden_state,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
+
+
+class BayesVLMVisionModel(CLIPVisionModelWithProjection):
+    def __init__(self, config):
+        super().__init__(config)
+        hidden = int(config.hidden_size)
+        proj = int(config.projection_dim)
+        self.register_buffer("a_inv", torch.zeros(hidden, hidden))
+        self.register_buffer("b_diag", torch.zeros(proj))
+
+    def set_covariance(self, a_inv: torch.Tensor, b_inv: torch.Tensor) -> None:
+        self.a_inv = a_inv
+        self.b_diag = torch.diagonal(b_inv)
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if not return_dict:
+            return super().forward(
+                pixel_values=pixel_values,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        vision_outputs = self.vision_model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        pooled_output = _get_output(vision_outputs, "pooler_output", 1)
+        last_hidden_state = _get_output(vision_outputs, "last_hidden_state", 0)
+        hidden_states = _get_output(vision_outputs, "hidden_states", 2)
+        attentions = _get_output(vision_outputs, "attentions", 3)
+        image_embeds = self.visual_projection(pooled_output)
+
+        image_var = _diag_cov(
+            pooled_output,
+            self.a_inv,
+            self.b_diag,
+            add_bias=self.visual_projection.bias is not None,
+        )
+        if image_var is None:
+            image_var = torch.zeros_like(image_embeds)
+        image_std = _std_from_var(image_var)
+
+        return BayesVLMVisionModelOutput(
+            image_embeds=image_embeds,
+            image_embeds_var=image_var,
+            image_embeds_std=image_std,
+            last_hidden_state=last_hidden_state,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
+
+
+class BayesVLMModel(CLIPModel):
+    def __init__(self, config):
+        super().__init__(config)
+        text_hidden = int(config.text_config.hidden_size)
+        vision_hidden = int(config.vision_config.hidden_size)
+        proj = int(config.projection_dim)
+        self.register_buffer("text_a_inv", torch.zeros(text_hidden, text_hidden))
+        self.register_buffer("text_b_diag", torch.zeros(proj))
+        self.register_buffer("image_a_inv", torch.zeros(vision_hidden, vision_hidden))
+        self.register_buffer("image_b_diag", torch.zeros(proj))
+
+    def set_covariances(
+        self,
+        image_a_inv: torch.Tensor,
+        image_b_inv: torch.Tensor,
+        text_a_inv: torch.Tensor,
+        text_b_inv: torch.Tensor,
+    ) -> None:
+        self.image_a_inv = image_a_inv
+        self.image_b_diag = torch.diagonal(image_b_inv)
+        self.text_a_inv = text_a_inv
+        self.text_b_diag = torch.diagonal(text_b_inv)
+
+    def _expected_logits_and_var(
+        self,
+        image_embeds: torch.Tensor,
+        text_embeds: torch.Tensor,
+        image_acts: torch.Tensor,
+        text_acts: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor | None]:
+        scale = self.logit_scale.exp()
+
+        if self.image_a_inv.numel() == 0 or self.text_a_inv.numel() == 0:
+            image_norm = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
+            text_norm = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)
+            logits = image_norm @ text_norm.t()
+            logits = logits * scale
+            return logits, None
+
+        image_diag_cov = _diag_cov(
+            image_acts,
+            self.image_a_inv,
+            self.image_b_diag,
+            add_bias=self.visual_projection.bias is not None,
+        )
+        text_diag_cov = _diag_cov(
+            text_acts,
+            self.text_a_inv,
+            self.text_b_diag,
+            add_bias=self.text_projection.bias is not None,
+        )
+
+        norm_image = image_embeds**2 + image_diag_cov
+        norm_text = text_embeds**2 + text_diag_cov
+        expect_norm_image = norm_image.sum(dim=-1, keepdim=True)
+        expect_norm_text = norm_text.sum(dim=-1, keepdim=True)
+
+        expected_similarity = torch.matmul(
+            image_embeds / torch.sqrt(expect_norm_image),
+            (text_embeds / torch.sqrt(expect_norm_text)).t(),
+        )
+
+        term1 = torch.matmul(norm_image, text_diag_cov.t())
+        term2 = torch.matmul(image_diag_cov, (text_embeds**2).t())
+        variance_similarity = (term1 + term2) / (expect_norm_image * expect_norm_text.t())
+
+        logits_mean = expected_similarity * scale
+        logits_var = variance_similarity * (scale**2)
+        return logits_mean, logits_var
+
+    def get_text_features(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        return_std: bool = False,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        text_outputs = self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        pooled_output = _get_output(text_outputs, "pooler_output", 1)
+        text_embeds = self.text_projection(pooled_output)
+
+        text_var = _diag_cov(
+            pooled_output,
+            self.text_a_inv,
+            self.text_b_diag,
+            add_bias=self.text_projection.bias is not None,
+        )
+        if text_var is None:
+            text_var = torch.zeros_like(text_embeds)
+        text_std = _std_from_var(text_var)
+
+        if not return_dict and not return_std:
+            return text_embeds
+
+        return BayesVLMEmbeddingOutput(mean=text_embeds, var=text_var, std=text_std)
+
+    def get_image_features(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        return_std: bool = False,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        vision_outputs = self.vision_model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        pooled_output = _get_output(vision_outputs, "pooler_output", 1)
+        image_embeds = self.visual_projection(pooled_output)
+
+        image_var = _diag_cov(
+            pooled_output,
+            self.image_a_inv,
+            self.image_b_diag,
+            add_bias=self.visual_projection.bias is not None,
+        )
+        if image_var is None:
+            image_var = torch.zeros_like(image_embeds)
+        image_std = _std_from_var(image_var)
+
+        if not return_dict and not return_std:
+            return image_embeds
+
+        return BayesVLMEmbeddingOutput(mean=image_embeds, var=image_var, std=image_std)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        return_loss: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if not return_dict:
+            return super().forward(
+                input_ids=input_ids,
+                pixel_values=pixel_values,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                return_loss=return_loss,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        text_outputs = self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        vision_outputs = self.vision_model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+
+        text_pooled = _get_output(text_outputs, "pooler_output", 1)
+        image_pooled = _get_output(vision_outputs, "pooler_output", 1)
+
+        text_embeds = self.text_projection(text_pooled)
+        image_embeds = self.visual_projection(image_pooled)
+
+        text_var = _diag_cov(
+            text_pooled,
+            self.text_a_inv,
+            self.text_b_diag,
+            add_bias=self.text_projection.bias is not None,
+        )
+        image_var = _diag_cov(
+            image_pooled,
+            self.image_a_inv,
+            self.image_b_diag,
+            add_bias=self.visual_projection.bias is not None,
+        )
+        if text_var is None:
+            text_var = torch.zeros_like(text_embeds)
+        if image_var is None:
+            image_var = torch.zeros_like(image_embeds)
+
+        text_std = _std_from_var(text_var)
+        image_std = _std_from_var(image_var)
+
+        logits_mean, logits_var = self._expected_logits_and_var(
+            image_embeds,
+            text_embeds,
+            image_pooled,
+            text_pooled,
+        )
+
+        text_embeds, text_var = _normalize_mean_and_var(text_embeds, text_var)
+        image_embeds, image_var = _normalize_mean_and_var(image_embeds, image_var)
+        text_std = _std_from_var(text_var)
+        image_std = _std_from_var(image_var)
+
+        logits_per_image = logits_mean
+        logits_per_text = logits_mean.t() if logits_mean is not None else None
+
+        if logits_var is None and logits_mean is not None:
+            logits_var = torch.zeros_like(logits_mean)
+        logits_per_image_var = _as_optional_tensor(logits_var)
+        logits_per_text_var = logits_var.t() if logits_var is not None else None
+
+        logits_per_image_std = _std_from_var(logits_per_image_var)
+        logits_per_text_std = _std_from_var(logits_per_text_var)
+
+        loss = None
+        if return_loss and logits_per_image is not None and logits_per_text is not None:
+            labels = torch.arange(logits_per_image.shape[0], device=logits_per_image.device)
+            loss_i = torch.nn.functional.cross_entropy(logits_per_image, labels)
+            loss_t = torch.nn.functional.cross_entropy(logits_per_text, labels)
+            loss = (loss_i + loss_t) / 2
+
+        return BayesVLMOutput(
+            loss=loss,
+            logits_per_image=logits_per_image,
+            logits_per_text=logits_per_text,
+            logits_per_image_var=logits_per_image_var,
+            logits_per_text_var=logits_per_text_var,
+            logits_per_image_std=logits_per_image_std,
+            logits_per_text_std=logits_per_text_std,
+            text_embeds=text_embeds,
+            image_embeds=image_embeds,
+            text_embeds_var=text_var,
+            image_embeds_var=image_var,
+            text_embeds_std=text_std,
+            image_embeds_std=image_std,
+            text_model_output=text_outputs,
+            vision_model_output=vision_outputs,
+        )

preprocessor_config.json

@@ -0,0 +1,45 @@
+{
+  "_valid_processor_keys": [
+    "images",
+    "do_resize",
+    "size",
+    "resample",
+    "do_center_crop",
+    "crop_size",
+    "do_rescale",
+    "rescale_factor",
+    "do_normalize",
+    "image_mean",
+    "image_std",
+    "do_convert_rgb",
+    "return_tensors",
+    "data_format",
+    "input_data_format"
+  ],
+  "crop_size": {
+    "height": 224,
+    "width": 224
+  },
+  "do_center_crop": true,
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "image_processor_type": "CLIPImageProcessor",
+  "image_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "processor_class": "CLIPProcessor",
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "shortest_edge": 224
+  }
+}

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:77f85aa5629f2cb0abe974fd5da165f1fcce6823b4088c0ce5df845c2e89e2a1
+size 610746962

special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

text/config.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "BayesVLMTextModel"
+  ],
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "modeling_bayesvlm_clip.BayesVLMTextModel",
+    "AutoProcessor": "transformers.CLIPProcessor"
+  },
+  "bos_token_id": 49406,
+  "dropout": 0.0,
+  "eos_token_id": 49407,
+  "hidden_act": "gelu",
+  "hidden_size": 512,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 2048,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 77,
+  "model_type": "clip_text_model",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 12,
+  "pad_token_id": 1,
+  "projection_dim": 512,
+  "torch_dtype": "float32",
+  "transformers_version": "4.40.2",
+  "vocab_size": 49408
+}

text/modeling_bayesvlm_clip.py

@@ -0,0 +1,495 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from transformers import CLIPModel, CLIPTextModelWithProjection, CLIPVisionModelWithProjection
+from transformers.modeling_outputs import ModelOutput
+
+
+def _as_optional_tensor(tensor: torch.Tensor | None) -> torch.Tensor | None:
+    return tensor if tensor is not None else None
+
+
+def _diag_cov(
+    activations: torch.Tensor,
+    a_inv: torch.Tensor,
+    b_diag: torch.Tensor,
+    add_bias: bool,
+) -> torch.Tensor | None:
+    if a_inv.numel() == 0 or b_diag.numel() == 0:
+        return None
+
+    if add_bias:
+        ones = torch.ones_like(activations[:, :1])
+        activations = torch.cat([activations, ones], dim=-1)
+
+    quad = torch.einsum("ij,jk,ik->i", activations, a_inv, activations)[:, None]
+    return quad * b_diag
+
+
+def _std_from_var(var: torch.Tensor | None) -> torch.Tensor | None:
+    if var is None:
+        return None
+    return torch.sqrt(var)
+
+def _get_output(outputs, name: str, index: int):
+    if hasattr(outputs, name):
+        return getattr(outputs, name)
+    if isinstance(outputs, (tuple, list)) and len(outputs) > index:
+        return outputs[index]
+    return None
+
+def _normalize_mean_and_var(
+    mean: torch.Tensor,
+    var: torch.Tensor,
+    eps: float = 1e-6,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    r2 = (mean**2).sum(dim=-1, keepdim=True).clamp_min(eps)
+    r = torch.sqrt(r2)
+    normalized = mean / r
+
+    # Delta-method approximation with diagonal covariance.
+    y2 = normalized**2
+    sum_y2v = (y2 * var).sum(dim=-1, keepdim=True)
+    norm_var = (var - 2 * y2 * var + y2 * sum_y2v) / r2
+    norm_var = norm_var.clamp_min(0)
+    return normalized, norm_var
+
+
+@dataclass
+class BayesVLMEmbeddingOutput(ModelOutput):
+    mean: torch.FloatTensor | None = None
+    var: torch.FloatTensor | None = None
+    std: torch.FloatTensor | None = None
+
+
+@dataclass
+class BayesVLMTextModelOutput(ModelOutput):
+    text_embeds: torch.FloatTensor | None = None
+    text_embeds_var: torch.FloatTensor | None = None
+    text_embeds_std: torch.FloatTensor | None = None
+    last_hidden_state: torch.FloatTensor | None = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class BayesVLMVisionModelOutput(ModelOutput):
+    image_embeds: torch.FloatTensor | None = None
+    image_embeds_var: torch.FloatTensor | None = None
+    image_embeds_std: torch.FloatTensor | None = None
+    last_hidden_state: torch.FloatTensor | None = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class BayesVLMOutput(ModelOutput):
+    loss: torch.FloatTensor | None = None
+    logits_per_image: torch.FloatTensor | None = None
+    logits_per_text: torch.FloatTensor | None = None
+    logits_per_image_var: torch.FloatTensor | None = None
+    logits_per_text_var: torch.FloatTensor | None = None
+    logits_per_image_std: torch.FloatTensor | None = None
+    logits_per_text_std: torch.FloatTensor | None = None
+    text_embeds: torch.FloatTensor | None = None
+    image_embeds: torch.FloatTensor | None = None
+    text_embeds_var: torch.FloatTensor | None = None
+    image_embeds_var: torch.FloatTensor | None = None
+    text_embeds_std: torch.FloatTensor | None = None
+    image_embeds_std: torch.FloatTensor | None = None
+    text_model_output: Optional[ModelOutput] = None
+    vision_model_output: Optional[ModelOutput] = None
+
+
+class BayesVLMTextModel(CLIPTextModelWithProjection):
+    def __init__(self, config):
+        super().__init__(config)
+        hidden = int(config.hidden_size)
+        proj = int(config.projection_dim)
+        self.register_buffer("a_inv", torch.zeros(hidden, hidden))
+        self.register_buffer("b_diag", torch.zeros(proj))
+
+    def set_covariance(self, a_inv: torch.Tensor, b_inv: torch.Tensor) -> None:
+        self.a_inv = a_inv
+        self.b_diag = torch.diagonal(b_inv)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if not return_dict:
+            return super().forward(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        text_outputs = self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        pooled_output = _get_output(text_outputs, "pooler_output", 1)
+        last_hidden_state = _get_output(text_outputs, "last_hidden_state", 0)
+        hidden_states = _get_output(text_outputs, "hidden_states", 2)
+        attentions = _get_output(text_outputs, "attentions", 3)
+        text_embeds = self.text_projection(pooled_output)
+
+        text_var = _diag_cov(
+            pooled_output,
+            self.a_inv,
+            self.b_diag,
+            add_bias=self.text_projection.bias is not None,
+        )
+        if text_var is None:
+            text_var = torch.zeros_like(text_embeds)
+        text_std = _std_from_var(text_var)
+
+        return BayesVLMTextModelOutput(
+            text_embeds=text_embeds,
+            text_embeds_var=text_var,
+            text_embeds_std=text_std,
+            last_hidden_state=last_hidden_state,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
+
+
+class BayesVLMVisionModel(CLIPVisionModelWithProjection):
+    def __init__(self, config):
+        super().__init__(config)
+        hidden = int(config.hidden_size)
+        proj = int(config.projection_dim)
+        self.register_buffer("a_inv", torch.zeros(hidden, hidden))
+        self.register_buffer("b_diag", torch.zeros(proj))
+
+    def set_covariance(self, a_inv: torch.Tensor, b_inv: torch.Tensor) -> None:
+        self.a_inv = a_inv
+        self.b_diag = torch.diagonal(b_inv)
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if not return_dict:
+            return super().forward(
+                pixel_values=pixel_values,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        vision_outputs = self.vision_model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        pooled_output = _get_output(vision_outputs, "pooler_output", 1)
+        last_hidden_state = _get_output(vision_outputs, "last_hidden_state", 0)
+        hidden_states = _get_output(vision_outputs, "hidden_states", 2)
+        attentions = _get_output(vision_outputs, "attentions", 3)
+        image_embeds = self.visual_projection(pooled_output)
+
+        image_var = _diag_cov(
+            pooled_output,
+            self.a_inv,
+            self.b_diag,
+            add_bias=self.visual_projection.bias is not None,
+        )
+        if image_var is None:
+            image_var = torch.zeros_like(image_embeds)
+        image_std = _std_from_var(image_var)
+
+        return BayesVLMVisionModelOutput(
+            image_embeds=image_embeds,
+            image_embeds_var=image_var,
+            image_embeds_std=image_std,
+            last_hidden_state=last_hidden_state,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
+
+
+class BayesVLMModel(CLIPModel):
+    def __init__(self, config):
+        super().__init__(config)
+        text_hidden = int(config.text_config.hidden_size)
+        vision_hidden = int(config.vision_config.hidden_size)
+        proj = int(config.projection_dim)
+        self.register_buffer("text_a_inv", torch.zeros(text_hidden, text_hidden))
+        self.register_buffer("text_b_diag", torch.zeros(proj))
+        self.register_buffer("image_a_inv", torch.zeros(vision_hidden, vision_hidden))
+        self.register_buffer("image_b_diag", torch.zeros(proj))
+
+    def set_covariances(
+        self,
+        image_a_inv: torch.Tensor,
+        image_b_inv: torch.Tensor,
+        text_a_inv: torch.Tensor,
+        text_b_inv: torch.Tensor,
+    ) -> None:
+        self.image_a_inv = image_a_inv
+        self.image_b_diag = torch.diagonal(image_b_inv)
+        self.text_a_inv = text_a_inv
+        self.text_b_diag = torch.diagonal(text_b_inv)
+
+    def _expected_logits_and_var(
+        self,
+        image_embeds: torch.Tensor,
+        text_embeds: torch.Tensor,
+        image_acts: torch.Tensor,
+        text_acts: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor | None]:
+        scale = self.logit_scale.exp()
+
+        if self.image_a_inv.numel() == 0 or self.text_a_inv.numel() == 0:
+            image_norm = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
+            text_norm = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)
+            logits = image_norm @ text_norm.t()
+            logits = logits * scale
+            return logits, None
+
+        image_diag_cov = _diag_cov(
+            image_acts,
+            self.image_a_inv,
+            self.image_b_diag,
+            add_bias=self.visual_projection.bias is not None,
+        )
+        text_diag_cov = _diag_cov(
+            text_acts,
+            self.text_a_inv,
+            self.text_b_diag,
+            add_bias=self.text_projection.bias is not None,
+        )
+
+        norm_image = image_embeds**2 + image_diag_cov
+        norm_text = text_embeds**2 + text_diag_cov
+        expect_norm_image = norm_image.sum(dim=-1, keepdim=True)
+        expect_norm_text = norm_text.sum(dim=-1, keepdim=True)
+
+        expected_similarity = torch.matmul(
+            image_embeds / torch.sqrt(expect_norm_image),
+            (text_embeds / torch.sqrt(expect_norm_text)).t(),
+        )
+
+        term1 = torch.matmul(norm_image, text_diag_cov.t())
+        term2 = torch.matmul(image_diag_cov, (text_embeds**2).t())
+        variance_similarity = (term1 + term2) / (expect_norm_image * expect_norm_text.t())
+
+        logits_mean = expected_similarity * scale
+        logits_var = variance_similarity * (scale**2)
+        return logits_mean, logits_var
+
+    def get_text_features(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        return_std: bool = False,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        text_outputs = self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        pooled_output = _get_output(text_outputs, "pooler_output", 1)
+        text_embeds = self.text_projection(pooled_output)
+
+        text_var = _diag_cov(
+            pooled_output,
+            self.text_a_inv,
+            self.text_b_diag,
+            add_bias=self.text_projection.bias is not None,
+        )
+        if text_var is None:
+            text_var = torch.zeros_like(text_embeds)
+        text_std = _std_from_var(text_var)
+
+        if not return_dict and not return_std:
+            return text_embeds
+
+        return BayesVLMEmbeddingOutput(mean=text_embeds, var=text_var, std=text_std)
+
+    def get_image_features(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        return_std: bool = False,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        vision_outputs = self.vision_model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        pooled_output = _get_output(vision_outputs, "pooler_output", 1)
+        image_embeds = self.visual_projection(pooled_output)
+
+        image_var = _diag_cov(
+            pooled_output,
+            self.image_a_inv,
+            self.image_b_diag,
+            add_bias=self.visual_projection.bias is not None,
+        )
+        if image_var is None:
+            image_var = torch.zeros_like(image_embeds)
+        image_std = _std_from_var(image_var)
+
+        if not return_dict and not return_std:
+            return image_embeds
+
+        return BayesVLMEmbeddingOutput(mean=image_embeds, var=image_var, std=image_std)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        return_loss: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if not return_dict:
+            return super().forward(
+                input_ids=input_ids,
+                pixel_values=pixel_values,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                return_loss=return_loss,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        text_outputs = self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        vision_outputs = self.vision_model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+
+        text_pooled = _get_output(text_outputs, "pooler_output", 1)
+        image_pooled = _get_output(vision_outputs, "pooler_output", 1)
+
+        text_embeds = self.text_projection(text_pooled)
+        image_embeds = self.visual_projection(image_pooled)
+
+        text_var = _diag_cov(
+            text_pooled,
+            self.text_a_inv,
+            self.text_b_diag,
+            add_bias=self.text_projection.bias is not None,
+        )
+        image_var = _diag_cov(
+            image_pooled,
+            self.image_a_inv,
+            self.image_b_diag,
+            add_bias=self.visual_projection.bias is not None,
+        )
+        if text_var is None:
+            text_var = torch.zeros_like(text_embeds)
+        if image_var is None:
+            image_var = torch.zeros_like(image_embeds)
+
+        text_std = _std_from_var(text_var)
+        image_std = _std_from_var(image_var)
+
+        logits_mean, logits_var = self._expected_logits_and_var(
+            image_embeds,
+            text_embeds,
+            image_pooled,
+            text_pooled,
+        )
+
+        text_embeds, text_var = _normalize_mean_and_var(text_embeds, text_var)
+        image_embeds, image_var = _normalize_mean_and_var(image_embeds, image_var)
+        text_std = _std_from_var(text_var)
+        image_std = _std_from_var(image_var)
+
+        logits_per_image = logits_mean
+        logits_per_text = logits_mean.t() if logits_mean is not None else None
+
+        if logits_var is None and logits_mean is not None:
+            logits_var = torch.zeros_like(logits_mean)
+        logits_per_image_var = _as_optional_tensor(logits_var)
+        logits_per_text_var = logits_var.t() if logits_var is not None else None
+
+        logits_per_image_std = _std_from_var(logits_per_image_var)
+        logits_per_text_std = _std_from_var(logits_per_text_var)
+
+        loss = None
+        if return_loss and logits_per_image is not None and logits_per_text is not None:
+            labels = torch.arange(logits_per_image.shape[0], device=logits_per_image.device)
+            loss_i = torch.nn.functional.cross_entropy(logits_per_image, labels)
+            loss_t = torch.nn.functional.cross_entropy(logits_per_text, labels)
+            loss = (loss_i + loss_t) / 2
+
+        return BayesVLMOutput(
+            loss=loss,
+            logits_per_image=logits_per_image,
+            logits_per_text=logits_per_text,
+            logits_per_image_var=logits_per_image_var,
+            logits_per_text_var=logits_per_text_var,
+            logits_per_image_std=logits_per_image_std,
+            logits_per_text_std=logits_per_text_std,
+            text_embeds=text_embeds,
+            image_embeds=image_embeds,
+            text_embeds_var=text_var,
+            image_embeds_var=image_var,
+            text_embeds_std=text_std,
+            image_embeds_std=image_std,
+            text_model_output=text_outputs,
+            vision_model_output=vision_outputs,
+        )

text/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ff0e694ffe9a4f158d729c7d81291ef05c8db29960e7b4baab382845fefe5245
+size 255875046

tokenizer_config.json

@@ -0,0 +1,31 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "49406": {
+      "content": "<|startoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49407": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|startoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "do_lower_case": true,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "model_max_length": 77,
+  "pad_token": "<|endoftext|>",
+  "processor_class": "CLIPProcessor",
+  "tokenizer_class": "CLIPTokenizer",
+  "unk_token": "<|endoftext|>"
+}

vision/config.json

@@ -0,0 +1,26 @@
+{
+  "architectures": [
+    "BayesVLMVisionModel"
+  ],
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "modeling_bayesvlm_clip.BayesVLMVisionModel",
+    "AutoProcessor": "transformers.CLIPProcessor"
+  },
+  "dropout": 0.0,
+  "hidden_act": "gelu",
+  "hidden_size": 768,
+  "image_size": 224,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-05,
+  "model_type": "clip_vision_model",
+  "num_attention_heads": 12,
+  "num_channels": 3,
+  "num_hidden_layers": 12,
+  "patch_size": 32,
+  "projection_dim": 512,
+  "torch_dtype": "float32",
+  "transformers_version": "4.40.2"
+}

vision/modeling_bayesvlm_clip.py

@@ -0,0 +1,495 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from transformers import CLIPModel, CLIPTextModelWithProjection, CLIPVisionModelWithProjection
+from transformers.modeling_outputs import ModelOutput
+
+
+def _as_optional_tensor(tensor: torch.Tensor | None) -> torch.Tensor | None:
+    return tensor if tensor is not None else None
+
+
+def _diag_cov(
+    activations: torch.Tensor,
+    a_inv: torch.Tensor,
+    b_diag: torch.Tensor,
+    add_bias: bool,
+) -> torch.Tensor | None:
+    if a_inv.numel() == 0 or b_diag.numel() == 0:
+        return None
+
+    if add_bias:
+        ones = torch.ones_like(activations[:, :1])
+        activations = torch.cat([activations, ones], dim=-1)
+
+    quad = torch.einsum("ij,jk,ik->i", activations, a_inv, activations)[:, None]
+    return quad * b_diag
+
+
+def _std_from_var(var: torch.Tensor | None) -> torch.Tensor | None:
+    if var is None:
+        return None
+    return torch.sqrt(var)
+
+def _get_output(outputs, name: str, index: int):
+    if hasattr(outputs, name):
+        return getattr(outputs, name)
+    if isinstance(outputs, (tuple, list)) and len(outputs) > index:
+        return outputs[index]
+    return None
+
+def _normalize_mean_and_var(
+    mean: torch.Tensor,
+    var: torch.Tensor,
+    eps: float = 1e-6,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    r2 = (mean**2).sum(dim=-1, keepdim=True).clamp_min(eps)
+    r = torch.sqrt(r2)
+    normalized = mean / r
+
+    # Delta-method approximation with diagonal covariance.
+    y2 = normalized**2
+    sum_y2v = (y2 * var).sum(dim=-1, keepdim=True)
+    norm_var = (var - 2 * y2 * var + y2 * sum_y2v) / r2
+    norm_var = norm_var.clamp_min(0)
+    return normalized, norm_var
+
+
+@dataclass
+class BayesVLMEmbeddingOutput(ModelOutput):
+    mean: torch.FloatTensor | None = None
+    var: torch.FloatTensor | None = None
+    std: torch.FloatTensor | None = None
+
+
+@dataclass
+class BayesVLMTextModelOutput(ModelOutput):
+    text_embeds: torch.FloatTensor | None = None
+    text_embeds_var: torch.FloatTensor | None = None
+    text_embeds_std: torch.FloatTensor | None = None
+    last_hidden_state: torch.FloatTensor | None = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class BayesVLMVisionModelOutput(ModelOutput):
+    image_embeds: torch.FloatTensor | None = None
+    image_embeds_var: torch.FloatTensor | None = None
+    image_embeds_std: torch.FloatTensor | None = None
+    last_hidden_state: torch.FloatTensor | None = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class BayesVLMOutput(ModelOutput):
+    loss: torch.FloatTensor | None = None
+    logits_per_image: torch.FloatTensor | None = None
+    logits_per_text: torch.FloatTensor | None = None
+    logits_per_image_var: torch.FloatTensor | None = None
+    logits_per_text_var: torch.FloatTensor | None = None
+    logits_per_image_std: torch.FloatTensor | None = None
+    logits_per_text_std: torch.FloatTensor | None = None
+    text_embeds: torch.FloatTensor | None = None
+    image_embeds: torch.FloatTensor | None = None
+    text_embeds_var: torch.FloatTensor | None = None
+    image_embeds_var: torch.FloatTensor | None = None
+    text_embeds_std: torch.FloatTensor | None = None
+    image_embeds_std: torch.FloatTensor | None = None
+    text_model_output: Optional[ModelOutput] = None
+    vision_model_output: Optional[ModelOutput] = None
+
+
+class BayesVLMTextModel(CLIPTextModelWithProjection):
+    def __init__(self, config):
+        super().__init__(config)
+        hidden = int(config.hidden_size)
+        proj = int(config.projection_dim)
+        self.register_buffer("a_inv", torch.zeros(hidden, hidden))
+        self.register_buffer("b_diag", torch.zeros(proj))
+
+    def set_covariance(self, a_inv: torch.Tensor, b_inv: torch.Tensor) -> None:
+        self.a_inv = a_inv
+        self.b_diag = torch.diagonal(b_inv)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if not return_dict:
+            return super().forward(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        text_outputs = self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        pooled_output = _get_output(text_outputs, "pooler_output", 1)
+        last_hidden_state = _get_output(text_outputs, "last_hidden_state", 0)
+        hidden_states = _get_output(text_outputs, "hidden_states", 2)
+        attentions = _get_output(text_outputs, "attentions", 3)
+        text_embeds = self.text_projection(pooled_output)
+
+        text_var = _diag_cov(
+            pooled_output,
+            self.a_inv,
+            self.b_diag,
+            add_bias=self.text_projection.bias is not None,
+        )
+        if text_var is None:
+            text_var = torch.zeros_like(text_embeds)
+        text_std = _std_from_var(text_var)
+
+        return BayesVLMTextModelOutput(
+            text_embeds=text_embeds,
+            text_embeds_var=text_var,
+            text_embeds_std=text_std,
+            last_hidden_state=last_hidden_state,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
+
+
+class BayesVLMVisionModel(CLIPVisionModelWithProjection):
+    def __init__(self, config):
+        super().__init__(config)
+        hidden = int(config.hidden_size)
+        proj = int(config.projection_dim)
+        self.register_buffer("a_inv", torch.zeros(hidden, hidden))
+        self.register_buffer("b_diag", torch.zeros(proj))
+
+    def set_covariance(self, a_inv: torch.Tensor, b_inv: torch.Tensor) -> None:
+        self.a_inv = a_inv
+        self.b_diag = torch.diagonal(b_inv)
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if not return_dict:
+            return super().forward(
+                pixel_values=pixel_values,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        vision_outputs = self.vision_model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        pooled_output = _get_output(vision_outputs, "pooler_output", 1)
+        last_hidden_state = _get_output(vision_outputs, "last_hidden_state", 0)
+        hidden_states = _get_output(vision_outputs, "hidden_states", 2)
+        attentions = _get_output(vision_outputs, "attentions", 3)
+        image_embeds = self.visual_projection(pooled_output)
+
+        image_var = _diag_cov(
+            pooled_output,
+            self.a_inv,
+            self.b_diag,
+            add_bias=self.visual_projection.bias is not None,
+        )
+        if image_var is None:
+            image_var = torch.zeros_like(image_embeds)
+        image_std = _std_from_var(image_var)
+
+        return BayesVLMVisionModelOutput(
+            image_embeds=image_embeds,
+            image_embeds_var=image_var,
+            image_embeds_std=image_std,
+            last_hidden_state=last_hidden_state,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
+
+
+class BayesVLMModel(CLIPModel):
+    def __init__(self, config):
+        super().__init__(config)
+        text_hidden = int(config.text_config.hidden_size)
+        vision_hidden = int(config.vision_config.hidden_size)
+        proj = int(config.projection_dim)
+        self.register_buffer("text_a_inv", torch.zeros(text_hidden, text_hidden))
+        self.register_buffer("text_b_diag", torch.zeros(proj))
+        self.register_buffer("image_a_inv", torch.zeros(vision_hidden, vision_hidden))
+        self.register_buffer("image_b_diag", torch.zeros(proj))
+
+    def set_covariances(
+        self,
+        image_a_inv: torch.Tensor,
+        image_b_inv: torch.Tensor,
+        text_a_inv: torch.Tensor,
+        text_b_inv: torch.Tensor,
+    ) -> None:
+        self.image_a_inv = image_a_inv
+        self.image_b_diag = torch.diagonal(image_b_inv)
+        self.text_a_inv = text_a_inv
+        self.text_b_diag = torch.diagonal(text_b_inv)
+
+    def _expected_logits_and_var(
+        self,
+        image_embeds: torch.Tensor,
+        text_embeds: torch.Tensor,
+        image_acts: torch.Tensor,
+        text_acts: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor | None]:
+        scale = self.logit_scale.exp()
+
+        if self.image_a_inv.numel() == 0 or self.text_a_inv.numel() == 0:
+            image_norm = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
+            text_norm = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)
+            logits = image_norm @ text_norm.t()
+            logits = logits * scale
+            return logits, None
+
+        image_diag_cov = _diag_cov(
+            image_acts,
+            self.image_a_inv,
+            self.image_b_diag,
+            add_bias=self.visual_projection.bias is not None,
+        )
+        text_diag_cov = _diag_cov(
+            text_acts,
+            self.text_a_inv,
+            self.text_b_diag,
+            add_bias=self.text_projection.bias is not None,
+        )
+
+        norm_image = image_embeds**2 + image_diag_cov
+        norm_text = text_embeds**2 + text_diag_cov
+        expect_norm_image = norm_image.sum(dim=-1, keepdim=True)
+        expect_norm_text = norm_text.sum(dim=-1, keepdim=True)
+
+        expected_similarity = torch.matmul(
+            image_embeds / torch.sqrt(expect_norm_image),
+            (text_embeds / torch.sqrt(expect_norm_text)).t(),
+        )
+
+        term1 = torch.matmul(norm_image, text_diag_cov.t())
+        term2 = torch.matmul(image_diag_cov, (text_embeds**2).t())
+        variance_similarity = (term1 + term2) / (expect_norm_image * expect_norm_text.t())
+
+        logits_mean = expected_similarity * scale
+        logits_var = variance_similarity * (scale**2)
+        return logits_mean, logits_var
+
+    def get_text_features(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        return_std: bool = False,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        text_outputs = self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        pooled_output = _get_output(text_outputs, "pooler_output", 1)
+        text_embeds = self.text_projection(pooled_output)
+
+        text_var = _diag_cov(
+            pooled_output,
+            self.text_a_inv,
+            self.text_b_diag,
+            add_bias=self.text_projection.bias is not None,
+        )
+        if text_var is None:
+            text_var = torch.zeros_like(text_embeds)
+        text_std = _std_from_var(text_var)
+
+        if not return_dict and not return_std:
+            return text_embeds
+
+        return BayesVLMEmbeddingOutput(mean=text_embeds, var=text_var, std=text_std)
+
+    def get_image_features(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        return_std: bool = False,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        vision_outputs = self.vision_model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        pooled_output = _get_output(vision_outputs, "pooler_output", 1)
+        image_embeds = self.visual_projection(pooled_output)
+
+        image_var = _diag_cov(
+            pooled_output,
+            self.image_a_inv,
+            self.image_b_diag,
+            add_bias=self.visual_projection.bias is not None,
+        )
+        if image_var is None:
+            image_var = torch.zeros_like(image_embeds)
+        image_std = _std_from_var(image_var)
+
+        if not return_dict and not return_std:
+            return image_embeds
+
+        return BayesVLMEmbeddingOutput(mean=image_embeds, var=image_var, std=image_std)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        return_loss: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if not return_dict:
+            return super().forward(
+                input_ids=input_ids,
+                pixel_values=pixel_values,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                return_loss=return_loss,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        text_outputs = self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        vision_outputs = self.vision_model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+
+        text_pooled = _get_output(text_outputs, "pooler_output", 1)
+        image_pooled = _get_output(vision_outputs, "pooler_output", 1)
+
+        text_embeds = self.text_projection(text_pooled)
+        image_embeds = self.visual_projection(image_pooled)
+
+        text_var = _diag_cov(
+            text_pooled,
+            self.text_a_inv,
+            self.text_b_diag,
+            add_bias=self.text_projection.bias is not None,
+        )
+        image_var = _diag_cov(
+            image_pooled,
+            self.image_a_inv,
+            self.image_b_diag,
+            add_bias=self.visual_projection.bias is not None,
+        )
+        if text_var is None:
+            text_var = torch.zeros_like(text_embeds)
+        if image_var is None:
+            image_var = torch.zeros_like(image_embeds)
+
+        text_std = _std_from_var(text_var)
+        image_std = _std_from_var(image_var)
+
+        logits_mean, logits_var = self._expected_logits_and_var(
+            image_embeds,
+            text_embeds,
+            image_pooled,
+            text_pooled,
+        )
+
+        text_embeds, text_var = _normalize_mean_and_var(text_embeds, text_var)
+        image_embeds, image_var = _normalize_mean_and_var(image_embeds, image_var)
+        text_std = _std_from_var(text_var)
+        image_std = _std_from_var(image_var)
+
+        logits_per_image = logits_mean
+        logits_per_text = logits_mean.t() if logits_mean is not None else None
+
+        if logits_var is None and logits_mean is not None:
+            logits_var = torch.zeros_like(logits_mean)
+        logits_per_image_var = _as_optional_tensor(logits_var)
+        logits_per_text_var = logits_var.t() if logits_var is not None else None
+
+        logits_per_image_std = _std_from_var(logits_per_image_var)
+        logits_per_text_std = _std_from_var(logits_per_text_var)
+
+        loss = None
+        if return_loss and logits_per_image is not None and logits_per_text is not None:
+            labels = torch.arange(logits_per_image.shape[0], device=logits_per_image.device)
+            loss_i = torch.nn.functional.cross_entropy(logits_per_image, labels)
+            loss_t = torch.nn.functional.cross_entropy(logits_per_text, labels)
+            loss = (loss_i + loss_t) / 2
+
+        return BayesVLMOutput(
+            loss=loss,
+            logits_per_image=logits_per_image,
+            logits_per_text=logits_per_text,
+            logits_per_image_var=logits_per_image_var,
+            logits_per_text_var=logits_per_text_var,
+            logits_per_image_std=logits_per_image_std,
+            logits_per_text_std=logits_per_text_std,
+            text_embeds=text_embeds,
+            image_embeds=image_embeds,
+            text_embeds_var=text_var,
+            image_embeds_var=image_var,
+            text_embeds_std=text_std,
+            image_embeds_std=image_std,
+            text_model_output=text_outputs,
+            vision_model_output=vision_outputs,
+        )

vision/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6c6c33da51b2f2e691b8fb9a69da3525398a68942e523351fad39b19e449230f
+size 354871602