Initial upload: BitMar Epoch 1 - 99,686,013 tokens processed

README.md

@@ -0,0 +1,43 @@
+---
+language: en
+license: mit
+tags:
+- bitmar
+- multimodal
+- babylm
+- cross-modal
+datasets:
+- babylm_multimodal
+metrics:
+- bleu
+- cross_modal_similarity
+---
+
+# BitMar 100M Token Model
+
+This model was trained on exactly 100 million tokens as part of the BabyLM challenge.
+
+## Training Details
+- Total tokens: 100,000,000
+- Epochs completed: 1
+- Tokens processed: 99,686,013
+- Cross-modal similarity: 0.3418
+
+## Model Architecture
+- Text encoder: 4 layers, 128 hidden size
+- Vision encoder: DiNOv2 features compressed to 128
+- Episodic memory: 32 slots
+
+## Usage
+```python
+from transformers import AutoModel, AutoTokenizer
+
+model = AutoModel.from_pretrained("euhidaman/bitmar-attention-multimodal")
+tokenizer = AutoTokenizer.from_pretrained("euhidaman/bitmar-attention-multimodal")
+```
+
+
+## Training Status
+- **Status**: In Progress (Epoch 1)
+- **Tokens Processed**: 99,686,013
+- **Best Cross-modal Similarity**: 0.3418

config.json

@@ -0,0 +1,42 @@
+{
+  "architectures": ["BitMarModel"],
+  "auto_map": {
+    "AutoConfig": "modeling_bitmar.BitMarConfig",
+    "AutoModel": "modeling_bitmar.BitMarModel"
+  },
+  "model_type": "bitmar",
+  "vocab_size": 50257,
+  "text_encoder_dim": 128,
+  "text_encoder_layers": 4,
+  "text_encoder_heads": 4,
+  "text_decoder_dim": 128,
+  "text_decoder_layers": 4,
+  "text_decoder_heads": 4,
+  "vision_encoder_dim": 768,
+  "vision_latent_size": 128,
+  "vision_hidden_size": 64,
+  "vision_compression_method": "learned_compression",
+  "vision_spatial_pooling": true,
+  "vision_pool_size": 2,
+  "fusion_hidden_size": 128,
+  "fusion_num_heads": 4,
+  "fusion_num_layers": 2,
+  "memory_size": 32,
+  "episode_dim": 128,
+  "memory_alpha": 0.2,
+  "direct_writing": true,
+  "memory_compression": true,
+  "max_seq_len": 256,
+  "dropout": 0.15,
+  "torch_dtype": "float32",
+  "transformers_version": "4.36.0",
+  "use_cache": true,
+  "tie_word_embeddings": true,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-5,
+  "pad_token_id": 50256,
+  "bos_token_id": 50256,
+  "eos_token_id": 50256,
+  "sep_token_id": null,
+  "decoder_start_token_id": null
+}

modeling_bitmar.py

@@ -0,0 +1,829 @@
+"""
+BitMar Model for Hugging Face Transformers
+BitNet-quantized Vision-Language Episodic Memory Transformer
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import logging
+import math
+import os
+import pickle
+import gzip
+from typing import Dict, List, Optional, Tuple, Union
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import CausalLMOutput, BaseModelOutput
+
+logger = logging.getLogger(__name__)
+
+
+class BitMarConfig(PretrainedConfig):
+    """Configuration class for BitMar model"""
+    
+    model_type = "bitmar"
+    
+    def __init__(
+        self,
+        vocab_size: int = 50257,
+        text_encoder_dim: int = 128,
+        text_encoder_layers: int = 4,
+        text_encoder_heads: int = 4,
+        text_decoder_dim: int = 128,
+        text_decoder_layers: int = 4,
+        text_decoder_heads: int = 4,
+        vision_encoder_dim: int = 768,
+        vision_latent_size: int = 128,
+        vision_hidden_size: int = 64,
+        vision_compression_method: str = "learned_compression",
+        vision_spatial_pooling: bool = True,
+        vision_pool_size: int = 2,
+        fusion_hidden_size: int = 128,
+        fusion_num_heads: int = 4,
+        fusion_num_layers: int = 2,
+        memory_size: int = 32,
+        episode_dim: int = 128,
+        memory_alpha: float = 0.2,
+        direct_writing: bool = True,
+        memory_compression: bool = True,
+        max_seq_len: int = 256,
+        dropout: float = 0.15,
+        initializer_range: float = 0.02,
+        layer_norm_epsilon: float = 1e-5,
+        use_cache: bool = True,
+        tie_word_embeddings: bool = True,
+        pad_token_id: int = 50256,
+        bos_token_id: int = 50256,
+        eos_token_id: int = 50256,
+        **kwargs
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs
+        )
+        
+        self.vocab_size = vocab_size
+        self.text_encoder_dim = text_encoder_dim
+        self.text_encoder_layers = text_encoder_layers
+        self.text_encoder_heads = text_encoder_heads
+        self.text_decoder_dim = text_decoder_dim
+        self.text_decoder_layers = text_decoder_layers
+        self.text_decoder_heads = text_decoder_heads
+        self.vision_encoder_dim = vision_encoder_dim
+        self.vision_latent_size = vision_latent_size
+        self.vision_hidden_size = vision_hidden_size
+        self.vision_compression_method = vision_compression_method
+        self.vision_spatial_pooling = vision_spatial_pooling
+        self.vision_pool_size = vision_pool_size
+        self.fusion_hidden_size = fusion_hidden_size
+        self.fusion_num_heads = fusion_num_heads
+        self.fusion_num_layers = fusion_num_layers
+        self.memory_size = memory_size
+        self.episode_dim = episode_dim
+        self.memory_alpha = memory_alpha
+        self.direct_writing = direct_writing
+        self.memory_compression = memory_compression
+        self.max_seq_len = max_seq_len
+        self.dropout = dropout
+        self.initializer_range = initializer_range
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.use_cache = use_cache
+        self.tie_word_embeddings = tie_word_embeddings
+
+
+class BitNetLinear(nn.Module):
+    """1.58-bit Linear layer following BitNet b1.58 architecture"""
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        
+        self.weight = nn.Parameter(torch.randn(out_features, in_features))
+        self.bias = nn.Parameter(torch.zeros(out_features)) if bias else None
+        
+        self.register_buffer('weight_scale', torch.ones(1))
+        self.register_buffer('input_scale', torch.ones(1))
+
+    def quantize_weights_1_58_bit(self, weight: torch.Tensor) -> torch.Tensor:
+        scale = weight.abs().mean()
+        self.weight_scale.data = scale.clamp(min=1e-5, max=1e3)
+        
+        weight_norm = torch.clamp(weight / self.weight_scale, min=-10.0, max=10.0)
+        threshold = 2.0 / 3.0
+        
+        quantized = torch.zeros_like(weight_norm)
+        quantized[weight_norm > threshold] = 1.0
+        quantized[weight_norm < -threshold] = -1.0
+        
+        return quantized
+
+    def quantize_activations_8bit(self, x: torch.Tensor) -> torch.Tensor:
+        x_clamped = torch.clamp(x, min=-1e6, max=1e6)
+        x_min, x_max = x_clamped.min(), x_clamped.max()
+        
+        range_val = x_max - x_min
+        if range_val < 1e-8:
+            return x_clamped
+            
+        scale = range_val / 255.0
+        self.input_scale.data = scale.clamp(min=1e-8, max=1e3)
+        
+        zero_point = (-x_min / scale).round().clamp(0, 255)
+        quantized = ((x_clamped / scale) + zero_point).round().clamp(0, 255)
+        dequantized = scale * (quantized - zero_point)
+        
+        return dequantized
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.training:
+            weight = self.quantize_weights_1_58_bit(self.weight)
+            x = self.quantize_activations_8bit(x)
+        else:
+            weight = self.weight
+            
+        output = F.linear(x, weight, self.bias)
+        return output
+
+
+class BitNetMLP(nn.Module):
+    """BitNet MLP block with 1.58-bit quantization"""
+
+    def __init__(self, dim: int, hidden_dim: int, dropout: float = 0.1):
+        super().__init__()
+        self.up_proj = BitNetLinear(dim, hidden_dim)
+        self.gate_proj = BitNetLinear(dim, hidden_dim)
+        self.down_proj = BitNetLinear(hidden_dim, dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate = torch.sigmoid(self.gate_proj(x))
+        up = F.silu(self.up_proj(x))
+        return self.dropout(self.down_proj(gate * up))
+
+
+class BitNetAttention(nn.Module):
+    """Multi-head attention with BitNet quantization"""
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        dropout: float = 0.1,
+        bias: bool = True
+    ):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        
+        assert self.head_dim * num_heads == dim
+        
+        self.q_proj = BitNetLinear(dim, dim, bias=bias)
+        self.k_proj = BitNetLinear(dim, dim, bias=bias)
+        self.v_proj = BitNetLinear(dim, dim, bias=bias)
+        self.out_proj = BitNetLinear(dim, dim, bias=bias)
+        
+        self.dropout = nn.Dropout(dropout)
+        self.scale = self.head_dim ** -0.5
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        mask: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        B, L, D = query.shape
+        
+        q = self.q_proj(query).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(key).view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(value).view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)
+        
+        attn_weights = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+        
+        if mask is not None:
+            attn_weights = attn_weights.masked_fill(mask.unsqueeze(1).unsqueeze(1) == 0, float('-inf'))
+        
+        attn_weights = F.softmax(attn_weights, dim=-1)
+        attn_weights = self.dropout(attn_weights)
+        
+        attn_output = torch.matmul(attn_weights, v)
+        attn_output = attn_output.transpose(1, 2).contiguous().view(B, L, D)
+        attn_output = self.out_proj(attn_output)
+        
+        return attn_output, attn_weights
+
+
+class BitNetTransformerBlock(nn.Module):
+    """BitNet Transformer block with quantized components"""
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        dropout: float = 0.1
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.attention = BitNetAttention(dim, num_heads, dropout)
+        self.norm2 = nn.LayerNorm(dim)
+        self.mlp = BitNetMLP(dim, int(dim * mlp_ratio), dropout)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Self-attention with residual
+        norm_x = self.norm1(x)
+        attn_out, attn_weights = self.attention(norm_x, norm_x, norm_x, mask)
+        x = x + attn_out
+        
+        # MLP with residual
+        x = x + self.mlp(self.norm2(x))
+        
+        return x, attn_weights
+
+
+class BitNetTextEncoder(nn.Module):
+    """BitNet-based text encoder"""
+
+    def __init__(
+        self,
+        vocab_size: int,
+        dim: int,
+        num_layers: int,
+        num_heads: int,
+        max_seq_len: int = 512,
+        dropout: float = 0.1
+    ):
+        super().__init__()
+        self.dim = dim
+        self.embedding = nn.Embedding(vocab_size, dim)
+        self.pos_embedding = nn.Embedding(max_seq_len, dim)
+        self.dropout = nn.Dropout(dropout)
+        
+        self.layers = nn.ModuleList([
+            BitNetTransformerBlock(dim, num_heads, dropout=dropout)
+            for _ in range(num_layers)
+        ])
+        
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+        B, L = input_ids.shape
+        
+        # Token embeddings + positional embeddings
+        positions = torch.arange(L, device=input_ids.device).unsqueeze(0)
+        x = self.embedding(input_ids) + self.pos_embedding(positions)
+        x = self.dropout(x)
+        
+        # Apply transformer layers
+        attention_weights = []
+        for layer in self.layers:
+            x, attn = layer(x, attention_mask)
+            attention_weights.append(attn)
+        
+        x = self.norm(x)
+        return x, attention_weights
+
+
+class BitNetTextDecoder(nn.Module):
+    """BitNet-based text decoder with causal masking"""
+
+    def __init__(
+        self,
+        vocab_size: int,
+        dim: int,
+        num_layers: int,
+        num_heads: int,
+        max_seq_len: int = 512,
+        dropout: float = 0.1
+    ):
+        super().__init__()
+        self.dim = dim
+        self.max_seq_len = max_seq_len
+        self.embedding = nn.Embedding(vocab_size, dim)
+        self.pos_embedding = nn.Embedding(max_seq_len, dim)
+        self.dropout = nn.Dropout(dropout)
+        
+        self.layers = nn.ModuleList([
+            BitNetTransformerBlock(dim, num_heads, dropout=dropout)
+            for _ in range(num_layers)
+        ])
+        
+        self.norm = nn.LayerNorm(dim)
+        self.lm_head = BitNetLinear(dim, vocab_size, bias=False)
+        
+        # Create causal mask
+        self.register_buffer(
+            "causal_mask",
+            torch.tril(torch.ones(max_seq_len, max_seq_len)).unsqueeze(0).unsqueeze(0)
+        )
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None
+    ) -> Dict[str, torch.Tensor]:
+        
+        if input_ids is not None:
+            B, L = input_ids.shape
+            positions = torch.arange(L, device=input_ids.device).unsqueeze(0)
+            x = self.embedding(input_ids) + self.pos_embedding(positions)
+        else:
+            x = inputs_embeds
+            B, L, _ = x.shape
+        
+        x = self.dropout(x)
+        
+        # Create causal mask
+        causal_mask = self.causal_mask[:, :, :L, :L]
+        if attention_mask is not None:
+            causal_mask = causal_mask * attention_mask.unsqueeze(1).unsqueeze(2)
+        
+        # Apply transformer layers
+        attention_weights = []
+        for layer in self.layers:
+            x, attn = layer(x, causal_mask)
+            attention_weights.append(attn)
+        
+        x = self.norm(x)
+        logits = self.lm_head(x)
+        
+        outputs = {"logits": logits, "hidden_states": x, "attentions": attention_weights}
+        
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+            outputs["loss"] = loss
+        
+        return outputs
+
+
+class EpisodicMemory(nn.Module):
+    """Episodic Memory mechanism inspired by Larimar"""
+
+    def __init__(
+        self,
+        memory_size: int,
+        episode_dim: int,
+        alpha: float = 0.1,
+        direct_writing: bool = True,
+        observation_noise_std: float = 1e-6,
+        external_storage: bool = False,
+        memory_storage_path: str = None,
+        compression_enabled: bool = True,
+        lazy_loading: bool = False
+    ):
+        super().__init__()
+        self.memory_size = memory_size
+        self.episode_dim = episode_dim
+        self.alpha = alpha
+        self.direct_writing = direct_writing
+        self.observation_noise_std = observation_noise_std
+        self.external_storage = external_storage
+        self.memory_storage_path = memory_storage_path
+        self.compression_enabled = compression_enabled
+        self.lazy_loading = lazy_loading
+        
+        # Initialize memory
+        self.register_buffer('memory', torch.randn(memory_size, episode_dim))
+        self.register_buffer('write_head', torch.zeros(1, dtype=torch.long))
+        self.register_buffer('memory_age', torch.zeros(memory_size))
+        
+        # Statistics
+        self.register_buffer('episode_mean', torch.zeros(episode_dim))
+        self.register_buffer('episode_std', torch.ones(episode_dim))
+        self.register_buffer('update_count', torch.zeros(1))
+
+    def write_memory(self, episode: torch.Tensor) -> torch.Tensor:
+        batch_size = episode.size(0)
+        
+        if self.direct_writing:
+            # Direct writing to memory
+            for i in range(batch_size):
+                write_pos = self.write_head.item()
+                self.memory[write_pos] = episode[i].detach()
+                self.memory_age[write_pos] = 0
+                self.write_head = (self.write_head + 1) % self.memory_size
+        
+        # Add observation noise
+        if self.observation_noise_std > 0:
+            noise = torch.randn_like(episode) * self.observation_noise_std
+            episode = episode + noise
+        
+        return episode
+
+    def read_memory(self, query: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size, query_dim = query.shape
+        
+        # Compute similarities
+        similarities = F.cosine_similarity(
+            query.unsqueeze(1), 
+            self.memory.unsqueeze(0), 
+            dim=-1
+        )
+        
+        # Apply softmax to get attention weights
+        attention_weights = F.softmax(similarities / 0.1, dim=-1)
+        
+        # Weighted sum of memory
+        retrieved = torch.sum(
+            attention_weights.unsqueeze(-1) * self.memory.unsqueeze(0), 
+            dim=1
+        )
+        
+        return retrieved, attention_weights
+
+    def forward(self, episode: torch.Tensor, mode: str = "read_write") -> Tuple[torch.Tensor, torch.Tensor]:
+        if mode == "write":
+            return self.write_memory(episode), torch.zeros(episode.size(0), self.memory_size, device=episode.device)
+        elif mode == "read":
+            return self.read_memory(episode)
+        else:  # read_write
+            # Write to memory
+            written_episode = self.write_memory(episode)
+            # Read from memory
+            retrieved, attention_weights = self.read_memory(episode)
+            return retrieved, attention_weights
+
+
+class CrossModalFusion(nn.Module):
+    """Cross-modal fusion module for text and vision features"""
+
+    def __init__(
+        self,
+        text_dim: int,
+        vision_dim: int,
+        hidden_dim: int,
+        num_heads: int = 8,
+        num_layers: int = 2
+    ):
+        super().__init__()
+        self.text_dim = text_dim
+        self.vision_dim = vision_dim
+        self.hidden_dim = hidden_dim
+        
+        # Project to same dimension
+        self.text_proj = BitNetLinear(text_dim, hidden_dim)
+        self.vision_proj = BitNetLinear(vision_dim, hidden_dim)
+        
+        # Cross-attention layers
+        self.cross_attention = nn.ModuleList([
+            BitNetAttention(hidden_dim, num_heads)
+            for _ in range(num_layers)
+        ])
+        
+        self.norm = nn.LayerNorm(hidden_dim)
+
+    def forward(
+        self,
+        text_features: torch.Tensor,
+        vision_features: torch.Tensor
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        
+        # Project to same dimension
+        text_proj = self.text_proj(text_features)
+        vision_proj = self.vision_proj(vision_features)
+        
+        # Cross-modal attention
+        fused_features = text_proj
+        attention_maps = {}
+        
+        for i, cross_attn in enumerate(self.cross_attention):
+            fused_features, attn_weights = cross_attn(
+                fused_features, vision_proj, vision_proj
+            )
+            attention_maps[f'cross_attn_{i}'] = attn_weights
+        
+        fused_features = self.norm(fused_features)
+        
+        return fused_features, attention_maps
+
+
+class VisionEncoder(nn.Module):
+    """Quantized Vision Encoder for DiNOv2 features"""
+
+    def __init__(
+        self,
+        input_dim: int = 768,
+        hidden_dim: int = 512,
+        output_dim: int = 768,
+        num_layers: int = 2
+    ):
+        super().__init__()
+        
+        layers = []
+        layers.append(BitNetLinear(input_dim, hidden_dim))
+        layers.append(nn.ReLU())
+        
+        for _ in range(num_layers - 1):
+            layers.append(BitNetLinear(hidden_dim, hidden_dim))
+            layers.append(nn.ReLU())
+        
+        layers.append(BitNetLinear(hidden_dim, output_dim))
+        
+        self.encoder = nn.Sequential(*layers)
+
+    def forward(self, vision_features: torch.Tensor) -> torch.Tensor:
+        return self.encoder(vision_features)
+
+
+class BitMarModel(PreTrainedModel):
+    """
+    BitMar: BitNet-quantized Vision-Language Episodic Memory Transformer
+    Compatible with Hugging Face Transformers
+    """
+    
+    config_class = BitMarConfig
+    base_model_prefix = "bitmar"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["BitNetTransformerBlock", "EpisodicMemory"]
+    
+    def __init__(self, config: BitMarConfig):
+        super().__init__(config)
+        self.config = config
+        
+        # Text encoder
+        self.text_encoder = BitNetTextEncoder(
+            vocab_size=config.vocab_size,
+            dim=config.text_encoder_dim,
+            num_layers=config.text_encoder_layers,
+            num_heads=config.text_encoder_heads,
+            max_seq_len=config.max_seq_len,
+            dropout=config.dropout
+        )
+        
+        # Text decoder
+        self.text_decoder = BitNetTextDecoder(
+            vocab_size=config.vocab_size,
+            dim=config.text_decoder_dim,
+            num_layers=config.text_decoder_layers,
+            num_heads=config.text_decoder_heads,
+            max_seq_len=config.max_seq_len,
+            dropout=config.dropout
+        )
+        
+        # Vision encoder
+        self.vision_encoder = VisionEncoder(
+            input_dim=config.vision_encoder_dim,
+            hidden_dim=config.vision_hidden_size,
+            output_dim=config.vision_latent_size
+        )
+        
+        # Cross-modal fusion
+        self.cross_modal_fusion = CrossModalFusion(
+            text_dim=config.text_encoder_dim,
+            vision_dim=config.vision_latent_size,
+            hidden_dim=config.fusion_hidden_size,
+            num_heads=config.fusion_num_heads,
+            num_layers=config.fusion_num_layers
+        )
+        
+        # Episodic memory
+        self.episodic_memory = EpisodicMemory(
+            memory_size=config.memory_size,
+            episode_dim=config.episode_dim,
+            alpha=config.memory_alpha,
+            direct_writing=config.direct_writing,
+            compression_enabled=config.memory_compression
+        )
+        
+        # Initialize weights
+        self.post_init()
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        if isinstance(module, (nn.Linear, BitNetLinear)):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+    def encode_text(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+        return self.text_encoder(input_ids, attention_mask)
+
+    def encode_vision(self, vision_features: torch.Tensor) -> torch.Tensor:
+        return self.vision_encoder(vision_features)
+
+    def create_episode(
+        self,
+        text_features: torch.Tensor,
+        vision_latent: torch.Tensor,
+        attention_weights: Dict[str, torch.Tensor]
+    ) -> torch.Tensor:
+        # Simple concatenation for episode creation
+        # Average pool text features
+        text_pooled = text_features.mean(dim=1)  # [B, D]
+        vision_pooled = vision_latent.mean(dim=1)  # [B, D]
+        
+        # Concatenate and project to episode dimension
+        episode = torch.cat([text_pooled, vision_pooled], dim=-1)
+        
+        # Project to episode dimension if needed
+        if episode.size(-1) != self.config.episode_dim:
+            if not hasattr(self, 'episode_proj'):
+                self.episode_proj = nn.Linear(episode.size(-1), self.config.episode_dim).to(episode.device)
+            episode = self.episode_proj(episode)
+        
+        return episode
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        vision_features: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        mode: str = "train",
+        step: int = 0,
+        has_vision: Optional[torch.Tensor] = None,
+        **kwargs
+    ) -> Union[Tuple, CausalLMOutput]:
+        
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        # Handle missing vision features
+        if vision_features is None:
+            batch_size = input_ids.size(0) if input_ids is not None else 1
+            vision_features = torch.zeros(batch_size, 196, self.config.vision_encoder_dim, device=self.device)
+        
+        # Encode text
+        text_features, text_attentions = self.encode_text(input_ids, attention_mask)
+        
+        # Encode vision
+        vision_latent = self.encode_vision(vision_features)
+        
+        # Cross-modal fusion
+        fused_features, fusion_attentions = self.cross_modal_fusion(text_features, vision_latent)
+        
+        # Create episode for memory
+        episode = self.create_episode(text_features, vision_latent, fusion_attentions)
+        
+        # Episodic memory interaction
+        retrieved_memory, memory_weights = self.episodic_memory(episode, mode="read_write")
+        
+        # Text generation with decoder
+        decoder_outputs = self.text_decoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=labels
+        )
+        
+        # Prepare outputs
+        loss = decoder_outputs.get("loss", None)
+        logits = decoder_outputs["logits"]
+        hidden_states = decoder_outputs["hidden_states"] if output_hidden_states else None
+        attentions = decoder_outputs["attentions"] if output_attentions else None
+        
+        if return_dict:
+            return CausalLMOutput(
+                loss=loss,
+                logits=logits,
+                hidden_states=hidden_states,
+                attentions=attentions,
+            )
+        else:
+            outputs = (logits,)
+            if loss is not None:
+                outputs = (loss,) + outputs
+            if hidden_states is not None:
+                outputs = outputs + (hidden_states,)
+            if attentions is not None:
+                outputs = outputs + (attentions,)
+            return outputs
+
+    def generate(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        vision_features: Optional[torch.FloatTensor] = None,
+        max_length: int = 100,
+        temperature: float = 0.7,
+        top_p: float = 0.9,
+        do_sample: bool = True,
+        **kwargs
+    ) -> torch.LongTensor:
+        """Simple generation method"""
+        
+        batch_size = input_ids.size(0)
+        device = input_ids.device
+        
+        # Handle missing vision features
+        if vision_features is None:
+            vision_features = torch.zeros(batch_size, 196, self.config.vision_encoder_dim, device=device)
+        
+        generated = input_ids.clone()
+        
+        for _ in range(max_length - input_ids.size(1)):
+            # Get model outputs
+            with torch.no_grad():
+                outputs = self.forward(
+                    input_ids=generated,
+                    attention_mask=attention_mask,
+                    vision_features=vision_features,
+                    return_dict=True
+                )
+            
+            # Get next token logits
+            next_token_logits = outputs.logits[:, -1, :] / temperature
+            
+            if do_sample:
+                # Apply top-p sampling
+                sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
+                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                
+                # Remove tokens with cumulative probability above the threshold
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                sorted_indices_to_remove[..., 0] = 0
+                
+                indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+                next_token_logits[indices_to_remove] = float('-inf')
+                
+                # Sample from the filtered distribution
+                probs = F.softmax(next_token_logits, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+            else:
+                # Greedy decoding
+                next_token = next_token_logits.argmax(dim=-1, keepdim=True)
+            
+            # Append to generated sequence
+            generated = torch.cat([generated, next_token], dim=-1)
+            
+            # Update attention mask
+            if attention_mask is not None:
+                attention_mask = torch.cat([
+                    attention_mask, 
+                    torch.ones(batch_size, 1, device=device)
+                ], dim=-1)
+            
+            # Stop if EOS token is generated
+            if (next_token == self.config.eos_token_id).all():
+                break
+        
+        return generated
+
+    def prepare_inputs_for_generation(
+        self, 
+        input_ids, 
+        past_key_values=None, 
+        attention_mask=None, 
+        vision_features=None,
+        **kwargs
+    ):
+        """Prepare inputs for generation"""
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "vision_features": vision_features,
+            "use_cache": kwargs.get("use_cache", True),
+        }
+
+
+# Register the model with transformers
+from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
+
+AutoConfig.register("bitmar", BitMarConfig)
+AutoModel.register(BitMarConfig, BitMarModel)
+AutoModelForCausalLM.register(BitMarConfig, BitMarModel)
+
+
+def count_parameters(model: nn.Module) -> Dict[str, int]:
+    """Count model parameters"""
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    
+    return {
+        "total_parameters": total_params,
+        "trainable_parameters": trainable_params,
+        "non_trainable_parameters": total_params - trainable_params
+    }
+
+
+def create_bitmar_model(config: Dict) -> BitMarModel:
+    """Create BitMar model from config dictionary"""
+    bitmar_config = BitMarConfig(**config)
+    model = BitMarModel(bitmar_config)
+    return model

pytorch_model.bin

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

tokenizer.json

@@ -0,0 +1,48 @@
+{
+  "version": "1.0",
+  "truncation": null,
+  "padding": null,
+  "added_tokens": [
+    {
+      "id": 50256,
+      "content": "<|endoftext|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": true,
+      "special": true
+    }
+  ],
+  "normalizer": {
+    "type": "NFC"
+  },
+  "pre_tokenizer": {
+    "type": "ByteLevel",
+    "add_prefix_space": false,
+    "trim_offsets": true,
+    "use_regex": true
+  },
+  "post_processor": {
+    "type": "ByteLevel",
+    "add_prefix_space": false,
+    "trim_offsets": true,
+    "use_regex": true
+  },
+  "decoder": {
+    "type": "ByteLevel",
+    "add_prefix_space": false,
+    "trim_offsets": true,
+    "use_regex": true
+  },
+  "model": {
+    "type": "BPE",
+    "dropout": null,
+    "unk_token": null,
+    "continuing_subword_prefix": null,
+    "end_of_word_suffix": null,
+    "fuse_unk": false,
+    "byte_fallback": false,
+    "vocab": {},
+    "merges": []
+  }
+}

tokenizer_config.json

@@ -0,0 +1,15 @@
+{
+  "tokenizer_class": "GPT2Tokenizer",
+  "auto_map": {
+    "AutoTokenizer": ["transformers", "GPT2Tokenizer"]
+  },
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>",
+  "add_prefix_space": false,
+  "model_max_length": 1024,
+  "special_tokens_map_file": null,
+  "name_or_path": "gpt2",
+  "tokenizer_type": "GPT2Tokenizer"
+}

training_metadata.json

@@ -0,0 +1,261 @@
+{
+  "epoch": 0,
+  "global_step": 99498,
+  "tokens_processed": 99686013,
+  "target_tokens": 100000000,
+  "best_similarity": 0.34183505177497864,
+  "training_config": {
+    "model": {
+      "vocab_size": 50257,
+      "text_encoder_dim": 128,
+      "text_encoder_layers": 4,
+      "text_encoder_heads": 4,
+      "text_decoder_dim": 128,
+      "text_decoder_layers": 4,
+      "text_decoder_heads": 4,
+      "vision_encoder_dim": 768,
+      "vision_latent_size": 128,
+      "vision_hidden_size": 64,
+      "vision_compression_method": "learned_compression",
+      "vision_spatial_pooling": true,
+      "vision_pool_size": 2,
+      "fusion_hidden_size": 128,
+      "fusion_num_heads": 4,
+      "fusion_num_layers": 2,
+      "memory_size": 32,
+      "episode_dim": 128,
+      "memory_alpha": 0.2,
+      "direct_writing": true,
+      "memory_compression": true,
+      "enable_adaptive_training": true,
+      "max_seq_len": 256,
+      "dropout": 0.15
+    },
+    "token_constraints": {
+      "total_tokens": 100000000,
+      "caption_tokens": 50000000,
+      "text_tokens": 50000000,
+      "enforce_exact_count": true,
+      "uniform_sampling": true,
+      "alignment_priority": "perfect_alignment",
+      "preserve_image_caption_pairs": true,
+      "strict_alignment_validation": true
+    },
+    "vision_feature_reduction": {
+      "enabled": true,
+      "method": "learned_compression",
+      "target_dim": 64,
+      "spatial_pooling": true,
+      "pool_method": "attention",
+      "hidden_dim": 128,
+      "learnable": true,
+      "preserve_spatial_info": true
+    },
+    "data": {
+      "dataset_dir": "../babylm_dataset",
+      "text_encoder_name": "gpt2",
+      "max_seq_length": 256,
+      "count_tokens": true,
+      "target_caption_tokens": 50000000,
+      "target_text_tokens": 50000000,
+      "token_counting_method": "gpt2",
+      "batch_size": 64,
+      "num_workers": 6,
+      "pin_memory": true,
+      "persistent_workers": true,
+      "mix_ratio": 0.5,
+      "shuffle_datasets": true,
+      "ensure_alignment": true,
+      "validate_alignment": true,
+      "alignment_verification": "strict",
+      "never_break_pairs": true,
+      "alignment_check_frequency": 1000,
+      "use_validation": false,
+      "train_only": true
+    },
+    "attention_analysis": {
+      "track_top_k": 5,
+      "log_every_n_steps": 200,
+      "viz_every_n_epochs": 3,
+      "save_head_patterns": true,
+      "analyze_memory_attention": true,
+      "analyze_cross_modal": true,
+      "track_token_alignment": true
+    },
+    "adaptive_training": {
+      "enabled": true,
+      "similarity_window_size": 200,
+      "drop_threshold": 0.12,
+      "min_steps_between_interventions": 800,
+      "freeze_duration_steps": 1500,
+      "loss_rebalance_factor": 2.0,
+      "similarity_smoothing_alpha": 0.15
+    },
+    "training": {
+      "max_epochs": 10,
+      "accumulate_grad_batches": 2,
+      "gradient_clip_val": 0.3,
+      "val_check_interval": 1000,
+      "scheduler": "cosine_with_restarts",
+      "min_lr": 5e-05,
+      "warmup_steps": 1000,
+      "learning_rate": 0.0002,
+      "weight_decay": 0.02,
+      "optimizer": "adamw8bit",
+      "scheduler_config": {
+        "T_0": 1000,
+        "T_mult": 2,
+        "eta_min_ratio": 0.1
+      },
+      "cross_modal_loss_weight": 1.5,
+      "text_generation_loss_weight": 1.0,
+      "memory_regularization_weight": 0.1,
+      "alignment_consistency_weight": 0.5,
+      "track_token_usage": true,
+      "log_token_progress": true,
+      "stop_at_token_limit": false,
+      "validate_alignment_every_n_steps": 500,
+      "log_alignment_metrics": true,
+      "alignment_loss_scaling": "adaptive"
+    },
+    "wandb": {
+      "project": "bitmar-100M-attention-epochs",
+      "entity": "babylm-ntust",
+      "api_key": null,
+      "log_every_n_steps": 100,
+      "log_attention": true,
+      "log_memory": true,
+      "log_gradients": true,
+      "log_token_usage": true,
+      "log_cross_modal_similarity": true,
+      "log_alignment_quality": true,
+      "log_caption_image_matching": true,
+      "save_code": true,
+      "create_plots": true,
+      "plot_attention_heatmaps": true,
+      "plot_memory_usage": true,
+      "plot_token_distribution": true,
+      "plot_alignment_metrics": true,
+      "log_memory_evolution": true,
+      "plot_memory_evolution_heatmap": true,
+      "plot_memory_diversity": true,
+      "plot_memory_access_patterns": true,
+      "memory_visualization_frequency": 5000,
+      "memory_snapshot_frequency": 10000,
+      "track_memory_metrics": [
+        "memory_diversity_score",
+        "memory_specialization_score",
+        "memory_usage_entropy",
+        "cross_modal_memory_ratio",
+        "memory_slot_utilization",
+        "memory_update_frequency",
+        "memory_retrieval_accuracy"
+      ]
+    },
+    "evaluation": {
+      "metrics": [
+        "bleu",
+        "rouge",
+        "cross_modal_similarity",
+        "memory_efficiency"
+      ],
+      "generate_samples": true,
+      "num_samples": 20,
+      "max_generation_length": 32,
+      "temperature": 0.8,
+      "top_p": 0.9,
+      "evaluate_alignment": true,
+      "alignment_metrics": [
+        "cosine_similarity",
+        "retrieval_accuracy",
+        "caption_image_matching",
+        "cross_modal_retrieval"
+      ],
+      "alignment_threshold": 0.8,
+      "validate_pairs_during_eval": true
+    },
+    "output": {
+      "checkpoint_dir": "checkpoints_100M_dataset",
+      "log_dir": "logs_100M_dataset",
+      "attention_dir": "attention_100M_dataset",
+      "memory_dir": "memory_100M_dataset",
+      "results_dir": "results_100M_dataset",
+      "token_logs_dir": "token_logs_100M_dataset"
+    },
+    "memory_optimization": {
+      "use_gradient_checkpointing": true,
+      "use_fp16": true,
+      "use_int8_vision": false,
+      "empty_cache_frequency": 10,
+      "max_memory_slots_in_ram": 16,
+      "compress_episodic_memory": true,
+      "vision_feature_caching": false,
+      "vision_batch_processing": true,
+      "tie_word_embeddings": true,
+      "use_shared_attention": false
+    },
+    "performance_targets": {
+      "max_model_size_mb": 50,
+      "target_cross_modal_similarity": 0.75,
+      "target_text_generation_quality": 0.6,
+      "memory_efficiency_threshold": 0.8
+    },
+    "flops_tracking": {
+      "enabled": true,
+      "log_frequency": 100,
+      "save_statistics": true,
+      "estimate_theoretical": true,
+      "track_peak_performance": true,
+      "log_to_wandb": true,
+      "detailed_breakdown": true,
+      "memory_bandwidth_tracking": false,
+      "efficiency_analysis": true,
+      "track_components": [
+        "attention",
+        "feedforward",
+        "layer_norm",
+        "embeddings",
+        "vision_encoder",
+        "cross_modal_fusion"
+      ]
+    },
+    "token_tracking": {
+      "log_frequency": 1000,
+      "save_token_distribution": true,
+      "monitor_caption_text_ratio": true,
+      "enforce_token_limits": false,
+      "early_stopping_on_limit": false,
+      "track_alignment_quality": true,
+      "log_misaligned_samples": true,
+      "alignment_quality_threshold": 0.7,
+      "save_alignment_statistics": true,
+      "correlate_flops_with_tokens": true,
+      "log_computational_efficiency": true,
+      "track_throughput_vs_quality": true
+    },
+    "huggingface_hub": {
+      "enabled": true,
+      "repo_id": "euhidaman/bitmar-attention-multimodal",
+      "private": true,
+      "upload_after_epoch": true,
+      "upload_final_model": true,
+      "commit_message_template": "BitMar 100M tokens - Epoch {epoch} - {tokens_processed:,} tokens processed",
+      "create_model_card": true,
+      "model_card_template": "---\nlanguage: en\nlicense: mit\ntags:\n- bitmar\n- multimodal\n- babylm\n- cross-modal\ndatasets:\n- babylm_multimodal\nmetrics:\n- bleu\n- cross_modal_similarity\n---\n\n# BitMar 100M Token Model\n\nThis model was trained on exactly 100 million tokens as part of the BabyLM challenge.\n\n## Training Details\n- Total tokens: 100,000,000\n- Epochs completed: {epoch}\n- Tokens processed: {tokens_processed:,}\n- Cross-modal similarity: {best_similarity:.4f}\n\n## Model Architecture\n- Text encoder: {text_encoder_layers} layers, {text_encoder_dim} hidden size\n- Vision encoder: DiNOv2 features compressed to {vision_latent_size}\n- Episodic memory: {memory_size} slots\n\n## Usage\n```python\nfrom transformers import AutoModel, AutoTokenizer\n\nmodel = AutoModel.from_pretrained(\"{repo_id}\")\ntokenizer = AutoTokenizer.from_pretrained(\"{repo_id}\")\n```\n"
+    },
+    "attention_sinks": {
+      "enabled": true,
+      "attention_sink_size": 4,
+      "attention_sink_window_size": 1020,
+      "inject_to_text_encoder": true,
+      "inject_to_text_decoder": true,
+      "position_shift_enabled": true,
+      "cache_compression": true,
+      "adaptive_window_size": false,
+      "memory_efficient_attention": true,
+      "preserve_episodic_memory": true,
+      "preserve_quantization": true,
+      "preserve_cross_modal_fusion": true
+    }
+  }
+}