added bank support

modeling_caption_bert.py

@@ -1,17 +1,30 @@
 # ============================================================================
-# CaptionBERT-8192: HuggingFace AutoModel-Compatible Implementation
+# CaptionBERT-8192: HuggingFace AutoModel with Alignment Bank
 #
 # Usage:
 #   from transformers import AutoModel, AutoTokenizer
 #   model = AutoModel.from_pretrained("AbstractPhil/geolip-captionbert-8192",
 #                                      trust_remote_code=True)
-#   tokenizer = AutoTokenizer.from_pretrained("AbstractPhil/geolip-captionbert-8192")
+#   tokenizer = AutoTokenizer.from_pretrained("AbstractPhil/geolip-captionbert-8192",
+#                                              trust_remote_code=True)
 #   inputs = tokenizer("A cat on a windowsill", return_tensors="pt",
 #                       padding=True, truncation=True, max_length=512)
 #   outputs = model(**inputs)
-#   embedding = outputs.last_hidden_state  # (B, 768) L2-normalized
+#
+#   # Core embedding (consensus-distilled, L2-normalized)
+#   embedding = outputs.last_hidden_state      # (B, 768)
+#
+#   # Enriched embedding (with geometric context from 5-expert bank)
+#   enriched = outputs.enriched                # (B, 768 + bank_dim)
+#
+#   # Token-level representations (pre-pooling, for sequence tasks)
+#   tokens = outputs.token_embeddings          # (B, L, 384)
+#
+#   # Geometric diagnostics
+#   geo = outputs.geometric_context            # dict with expert cos, anchors, etc.
 # ============================================================================
 
+import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -30,8 +43,14 @@ class CaptionBertConfig(PretrainedConfig):
         num_hidden_layers=6,
         intermediate_size=1536,
         output_dim=768,
-        hidden_dropout_prob=0.1,
+        hidden_dropout_prob=0.0,
         pad_token_id=0,
+        # Alignment bank
+        bank_enabled=True,
+        bank_n_experts=5,
+        bank_n_anchors=512,
+        bank_dim=128,
+        bank_cv_target=0.082,
         **kwargs,
     ):
         super().__init__(pad_token_id=pad_token_id, **kwargs)
@@ -43,18 +62,151 @@ class CaptionBertConfig(PretrainedConfig):
         self.intermediate_size = intermediate_size
         self.output_dim = output_dim
         self.hidden_dropout_prob = hidden_dropout_prob
+        self.bank_enabled = bank_enabled
+        self.bank_n_experts = bank_n_experts
+        self.bank_n_anchors = bank_n_anchors
+        self.bank_dim = bank_dim
+        self.bank_cv_target = bank_cv_target
+
+
+class AlignmentBank(nn.Module):
+    """
+    Geometric interface layer preserving 5-expert differentiation structure.
+
+    Trained post-hoc on frozen encoder via GPA + whitened Procrustes.
+    Stores per-expert rotation matrices, whiteners, and means that encode
+    how each expert's geometric perspective differs from the consensus center.
+
+    Provides geometric context annotations (128-dim) alongside the core
+    768-dim consensus embedding for downstream heads.
+    """
+    def __init__(self, d_embed=768, n_experts=5, n_anchors=512, d_bank=128):
+        super().__init__()
+        self.d_embed = d_embed
+        self.n_experts = n_experts
+        self.n_anchors = n_anchors
+        self.d_bank = d_bank
+
+        # Per-expert Procrustes components (the differentiation structure)
+        self.expert_rotations = nn.ParameterList([
+            nn.Parameter(torch.eye(d_embed)) for _ in range(n_experts)])
+        self.expert_whiteners = nn.ParameterList([
+            nn.Parameter(torch.eye(d_embed)) for _ in range(n_experts)])
+        self.expert_means = nn.ParameterList([
+            nn.Parameter(torch.zeros(d_embed)) for _ in range(n_experts)])
+
+        # Consensus landmarks on the hypersphere
+        self.anchors = nn.Parameter(
+            F.normalize(torch.randn(n_anchors, d_embed), dim=-1))
+
+        # Geometric context projection
+        n_cross = n_experts * (n_experts - 1) // 2
+        geo_dim = n_experts + n_experts + n_cross + 1 + n_experts + n_anchors
+        self.geo_proj = nn.Sequential(
+            nn.Linear(geo_dim, d_bank * 2), nn.GELU(), nn.LayerNorm(d_bank * 2),
+            nn.Linear(d_bank * 2, d_bank), nn.LayerNorm(d_bank))
+
+        # Calibrated consensus targets (preserved from training)
+        self.register_buffer("target_cv", torch.tensor(0.082))
+        self.register_buffer("target_cross_cos_mean", torch.tensor(0.0))
+        self.register_buffer("target_cross_cos_std", torch.tensor(0.0))
+        self.register_buffer("target_disagreement_ratio", torch.tensor(0.0))
+
+    def forward(self, embedding):
+        B = embedding.shape[0]
+        emb = embedding.float()
+
+        # Full whitened Procrustes per expert: center → whiten → normalize → rotate
+        expert_consistency = []
+        expert_recon = []
+        expert_projected = []
+        for i in range(self.n_experts):
+            R = self.expert_rotations[i]
+            W = self.expert_whiteners[i]
+            mu = self.expert_means[i]
+            centered = emb - mu
+            whitened = centered @ W
+            whitened_n = F.normalize(whitened, dim=-1)
+            in_expert = whitened_n @ R.T
+            back = in_expert @ R
+            cos = F.cosine_similarity(whitened_n, back, dim=-1)
+            recon = (whitened_n - back).pow(2).mean(dim=-1)
+            expert_consistency.append(cos)
+            expert_recon.append(recon)
+            expert_projected.append(in_expert)
+
+        expert_cos = torch.stack(expert_consistency, dim=-1)
+        expert_mse = torch.stack(expert_recon, dim=-1)
+
+        # Cross-expert differentiation (10 pairs for 5 experts)
+        cross_cos = []
+        for i in range(self.n_experts):
+            for j in range(i + 1, self.n_experts):
+                cc = F.cosine_similarity(
+                    expert_projected[i], expert_projected[j], dim=-1)
+                cross_cos.append(cc)
+        cross_features = torch.stack(cross_cos, dim=-1)
+
+        # Per-sample disagreement
+        per_sample_agreement = expert_cos.mean(dim=-1)
+        per_sample_disagreement = expert_cos.std(dim=-1)
+        disagreement_ratio = per_sample_disagreement / (per_sample_agreement + 1e-8)
+
+        # Expert norm ratios
+        expert_norms = []
+        for i in range(self.n_experts):
+            W = self.expert_whiteners[i]; mu = self.expert_means[i]
+            whitened = (emb - mu) @ W
+            expert_norms.append(whitened.norm(dim=-1))
+        norm_ratio = torch.stack(expert_norms, dim=-1)
+        norm_ratio = norm_ratio / (norm_ratio.mean(dim=-1, keepdim=True) + 1e-8)
+
+        # Anchor distances
+        anchors_n = F.normalize(self.anchors, dim=-1)
+        anchor_cos = emb @ anchors_n.T
+
+        # Geometric context vector
+        geo_input = torch.cat([
+            expert_cos, expert_mse, cross_features,
+            disagreement_ratio.unsqueeze(-1), norm_ratio, anchor_cos
+        ], dim=-1)
+        geo_context = self.geo_proj(geo_input)
+        enriched = torch.cat([embedding, geo_context], dim=-1)
+
+        # Diagnostics
+        diagnostics = {
+            "expert_cos_mean": expert_cos.mean().item(),
+            "expert_cos_std": expert_cos.std().item(),
+            "cross_expert_cos": cross_features.mean().item(),
+            "cross_expert_cos_std": cross_features.std().item(),
+            "anchor_max_cos": anchor_cos.max(dim=-1).values.mean().item(),
+            "anchor_mean_cos": anchor_cos.mean().item(),
+            "disagreement_ratio": disagreement_ratio.mean().item(),
+            "norm_ratio_spread": norm_ratio.std(dim=-1).mean().item(),
+        }
+
+        return enriched, geo_context, diagnostics
 
 
 class CaptionBertModel(PreTrainedModel):
     """
-    Consensus-distilled caption encoder.
+    Consensus-distilled caption encoder with geometric alignment bank.
 
-    Produces L2-normalized 768-dim embeddings in the geometric consensus
-    space of 5 BERT-family models (BERT, ModernBERT, RoBERTa, ALBERT, DistilBERT).
+    The encoder produces L2-normalized 768-dim embeddings in the geometric
+    consensus space of 5 BERT-family models (BERT, ModernBERT, RoBERTa,
+    ALBERT, DistilBERT), aligned via Generalized Procrustes Analysis.
 
-    Output:
-        last_hidden_state: (B, output_dim) L2-normalized embedding
-        pooler_output:     (B, output_dim) same as last_hidden_state (for compatibility)
+    The alignment bank annotates each embedding with 128-dim geometric
+    context from the 5-expert differentiation structure — per-expert
+    consistency, cross-expert disagreement, and anchor distances.
+
+    Output fields:
+        last_hidden_state:   (B, 768)         L2-normalized consensus embedding
+        pooler_output:       (B, 768)         same (HF compatibility)
+        token_embeddings:    (B, L, 384)      pre-pooling token representations
+        enriched:            (B, 896)         embedding + bank geometric context
+        geometric_context:   dict             expert cos, cross-expert, anchors, etc.
+        hidden_states:       tuple            per-layer outputs (if requested)
     """
     config_class = CaptionBertConfig
 
@@ -62,7 +214,7 @@ class CaptionBertModel(PreTrainedModel):
         super().__init__(config)
         self.config = config
 
-        # Embeddings
+        # ── Encoder ──
         self.token_emb = nn.Embedding(
             config.vocab_size, config.hidden_size,
             padding_idx=config.pad_token_id)
@@ -71,7 +223,6 @@ class CaptionBertModel(PreTrainedModel):
         self.emb_norm = nn.LayerNorm(config.hidden_size)
         self.emb_drop = nn.Dropout(config.hidden_dropout_prob)
 
-        # Transformer encoder
         encoder_layer = nn.TransformerEncoderLayer(
             d_model=config.hidden_size,
             nhead=config.num_attention_heads,
@@ -85,7 +236,6 @@ class CaptionBertModel(PreTrainedModel):
             encoder_layer, num_layers=config.num_hidden_layers,
             enable_nested_tensor=False)
 
-        # Output projection to consensus space
         self.output_proj = nn.Sequential(
             nn.Linear(config.hidden_size, config.hidden_size),
             nn.GELU(),
@@ -93,6 +243,17 @@ class CaptionBertModel(PreTrainedModel):
             nn.Linear(config.hidden_size, config.output_dim),
         )
 
+        # ── Alignment Bank ──
+        if getattr(config, 'bank_enabled', False):
+            self.bank = AlignmentBank(
+                d_embed=config.output_dim,
+                n_experts=config.bank_n_experts,
+                n_anchors=config.bank_n_anchors,
+                d_bank=config.bank_dim,
+            )
+        else:
+            self.bank = None
+
         self.post_init()
 
     def forward(self, input_ids=None, attention_mask=None,
@@ -100,39 +261,43 @@ class CaptionBertModel(PreTrainedModel):
         B, L = input_ids.shape
         device = input_ids.device
 
-        # Embed
+        # ── Encode ──
         positions = torch.arange(L, device=device).unsqueeze(0)
         x = self.token_emb(input_ids) + self.pos_emb(positions)
         x = self.emb_drop(self.emb_norm(x))
 
-        # Transformer with padding mask
         if attention_mask is not None:
             key_padding_mask = ~attention_mask.bool()
         else:
             key_padding_mask = (input_ids == self.config.pad_token_id)
 
-        # Layer-by-layer for hidden state capture
         hidden_states = [x] if output_hidden_states else None
         for layer in self.encoder.layers:
             x = layer(x, src_key_padding_mask=key_padding_mask)
             if output_hidden_states:
                 hidden_states.append(x)
 
-        # Mean pool over non-padding tokens
+        # ── Pool + Project ──
         if attention_mask is not None:
             mask = attention_mask.unsqueeze(-1).float()
         else:
             mask = (~key_padding_mask).unsqueeze(-1).float()
         pooled = (x * mask).sum(1) / mask.sum(1).clamp(min=1)
-
-        # Project and normalize
         embedding = F.normalize(self.output_proj(pooled), dim=-1)
 
-        # Return in HuggingFace-compatible format
+        # ── Alignment Bank ──
+        enriched = None
+        geo_diagnostics = None
+        if self.bank is not None:
+            enriched, _, geo_diagnostics = self.bank(embedding)
+
+        # ── Output ──
         result = {
-            'last_hidden_state': embedding,       # (B, 768) pooled, normalized
-            'pooler_output': embedding,            # same, for compatibility
-            'token_embeddings': x,                 # (B, L, 384) pre-pooling sequence
+            'last_hidden_state': embedding,       # (B, 768)
+            'pooler_output': embedding,            # (B, 768) compat
+            'token_embeddings': x,                 # (B, L, 384)
+            'enriched': enriched,                  # (B, 896) or None
+            'geometric_context': geo_diagnostics,  # dict or None
         }
         if output_hidden_states:
             result['hidden_states'] = tuple(hidden_states)
@@ -141,29 +306,14 @@ class CaptionBertModel(PreTrainedModel):
 
     def encode(self, texts, tokenizer=None, max_length=512, batch_size=128,
                device=None):
-        """
-        Convenience method: raw text → L2-normalized embeddings.
-
-        Args:
-            texts: str or list of str
-            tokenizer: AutoTokenizer instance (loads default if None)
-            max_length: max token length
-            batch_size: encoding batch size
-            device: torch device
-
-        Returns:
-            (N, 768) L2-normalized tensor
-        """
+        """Convenience: raw text → L2-normalized (N, 768) embeddings."""
         if isinstance(texts, str):
             texts = [texts]
-
         if tokenizer is None:
             from transformers import AutoTokenizer
             tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
-
         if device is None:
             device = next(self.parameters()).device
-
         self.eval()
         all_emb = []
         with torch.no_grad():