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, )