Figure out ways to improve my model architecture while figuring out if the multiple heads help or not. You cannot use the dot product as it is expensive to calculate. Try and improve the speed and memory footprint of the model.
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, PretrainedConfig

# 1. Efficient Additive Attention
class EfficientAdditiveAttention(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super(EfficientAdditiveAttention, self).__init__()
        self.query_proj = nn.Linear(input_dim, hidden_dim, bias=False)
        self.key_proj = nn.Linear(input_dim, hidden_dim, bias=False)
        self.energy_proj = nn.Linear(hidden_dim, 1, bias=False)
        self.activation = nn.ReLU()

    def forward(self, queries, keys, mask=None):
        """
        queries: (B, H, Q, D)
        keys: (B, H, K, D)
        mask: (B, 1, Q, K) or None
        Returns:
            attention_scores: (B, H, Q, K)
        """
        # Project queries and keys
        proj_queries = self.query_proj(queries)  # (B, H, Q, hidden_dim)
        proj_keys = self.key_proj(keys)          # (B, H, K, hidden_dim)

        # Expand queries and keys for additive attention
        proj_queries = proj_queries.unsqueeze(3)  # (B, H, Q, 1, hidden_dim)
        proj_keys = proj_keys.unsqueeze(2)          # (B, H, 1, K, hidden_dim)

        # Compute additive energy
        energy = self.activation(proj_queries + proj_keys)  # (B, H, Q, K, hidden_dim)

        # Compute attention scores
        scores = self.energy_proj(energy).squeeze(-1)  # (B, H, Q, K)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        attention_weights = F.softmax(scores, dim=-1)  # (B, H, Q, K)

        return attention_weights

# 2. Trainable Memory
class TrainableMemory(nn.Module):
    def __init__(self, memory_size, memory_dim, output_dim):
        super(TrainableMemory, self).__init__()
        self.memory = nn.Parameter(torch.randn(memory_size, memory_dim))
        self.memory_proj = nn.Linear(memory_dim, output_dim, bias=False)

    def forward(self, batch_size):
        """
        Returns:
            memory_output: (B, M, D_out)
        """
        # Expand memory to (B, M, D_mem)
        memory = self.memory.unsqueeze(0).expand(batch_size, -1, -1)  # (B, M, D_mem)
        projected_memory = self.memory_proj(memory)                   # (B, M, D_out)
        return projected_memory

# 3. Parallel Guidance
class ParallelGuidance(nn.Module):
    def __init__(self, input_dim, guidance_dim):
        super(ParallelGuidance, self).__init__()
        self.guidance_layer = nn.Sequential(
            nn.Linear(input_dim, guidance_dim),
            nn.ReLU(),
            nn.Linear(guidance_dim, input_dim)
        )

    def forward(self, x):
        """
        x: (B, Q+M, D)
        Returns:
            guided: (B, Q+M, D)
        """
        return self.guidance_layer(x)  # (B, Q+M, D)

# 4. Multi-Head Additive Attention
class Attention(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_heads=1):
        super(Attention, self).__init__()
        self.input_dim = input_dim
        self.num_heads = num_heads
        self.hidden_dim = hidden_dim
        self.head_dim = input_dim // num_heads
        assert self.head_dim * num_heads == self.input_dim, "input_dim must be divisible by num_heads"

        self.additive_attention = EfficientAdditiveAttention(input_dim, hidden_dim)
        self.value_proj = nn.Linear(input_dim, input_dim, bias=False)
        self.output_proj = nn.Linear(input_dim, input_dim, bias=False)

        # Gating mechanism
        self.gate = nn.Sequential(
            nn.Linear(input_dim, input_dim),
            nn.Sigmoid()
        )

    def forward(self, input_seq, mask=None):
        """
        input_seq: (B, Q, D)
        mask: (B, Q, K) or None
        Returns:
            output: (B, Q, D)
        """
        B, Q, D = input_seq.size()

        # Reshape for multi-head: (B, H, Q, D')
        input_seq = input_seq.view(B, Q, self.num_heads, self.head_dim).transpose(1, 2)  # (B, H, Q, D')

        # Compute attention weights
        attention_weights = self.additive_attention(input_seq, input_seq, mask)  # (B, H, Q, K)

        # Project values
        values = self.value_proj(input_seq)  # (B, H, Q, D')

        # Weighted sum of values
        attn_output = torch.matmul(attention_weights, values)  # (B, H, Q, D')

        # Concatenate heads: (B, Q, D)
        attn_output = attn_output.transpose(1, 2).contiguous().view(B, Q, D)

        # Gating
        gate = self.gate(attn_output)  # (B, Q, D)
        gated_output = gate * attn_output  # (B, Q, D)

        # Final projection
        output = self.output_proj(gated_output)  # (B, Q, D)

        return output 

# 5. Transformer Layer
class TransformerLayer(nn.Module):
    def __init__(self, embed_size, num_heads, hidden_size, memory_size, memory_dim, guidance_dim, guidance_hidden_dim=512):
        super(TransformerLayer, self).__init__()
        self.self_attention = Attention(embed_size, hidden_size, num_heads=num_heads)
        
        # Feed-Forward Network (FFN)
        self.feed_forward = nn.Sequential(
            nn.Linear(embed_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, embed_size)
        )
        
        # Trainable Memory
        self.memory = TrainableMemory(memory_size, memory_dim, embed_size)
        
        # Parallel Guidance
        self.guidance = ParallelGuidance(embed_size, guidance_dim)
        
        # Layer Normalization
        self.layer_norm1 = nn.LayerNorm(embed_size)
        self.layer_norm2 = nn.LayerNorm(embed_size)
        self.layer_norm3 = nn.LayerNorm(embed_size)

    def forward(self, x, mask=None):
        """
        x: (B, Q, D)
        mask: (B, Q, K) or None
        Returns:
            output: (B, Q+M, D)
        """
        # ======================
        # Self-Attention SubLayer
        # ======================
        attn_output = self.self_attention(x, mask)  # (B, Q, D)
        x = self.layer_norm1(x + attn_output)       # (B, Q, D)

        # ======================
        # Feed-Forward SubLayer
        # ======================
        ffn_output = self.feed_forward(x)           # (B, Q, D)
        x = self.layer_norm2(x + ffn_output)        # (B, Q, D)

        # ======================
        # Incorporate Trainable Memory
        # ======================
        memory_output = self.memory(x.size(0))      # (B, M, D)
        combined = torch.cat([x, memory_output], dim=1)  # (B, Q+M, D)

        # ======================
        # Parallel Guidance SubLayer
        # ======================
        guided = self.guidance(combined)            # (B, Q+M, D)

        # ======================
        # Gating Mechanism
        # ======================
        gate = torch.sigmoid(guided)                # (B, Q+M, D)
        gated_output = gate * guided                # (B, Q+M, D)

        # ======================
        # Final Residual Connection and LayerNorm
        # ======================
        output = self.layer_norm3(gated_output + combined)  # (B, Q+M, D)

        return output

# 6. Language Model Configuration
class LanguageModelConfig(PretrainedConfig):
    model_type = "gated_attention_transformer"

    def __init__(self, vocab_size=30522, embed_size=768, num_layers=6, num_heads=8, hidden_size=2048, memory_size=128, memory_dim=256, guidance_dim=512, **kwargs):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.embed_size = embed_size
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.hidden_size = hidden_size
        self.memory_size = memory_size
        self.memory_dim = memory_dim
        self.guidance_dim = guidance_dim

# 7. Language Model
class LanguageModel(PreTrainedModel):
    config_class = LanguageModelConfig

    def __init__(self, config):
        super().__init__(config)
        self.token_embedding = nn.Embedding(config.vocab_size, config.embed_size, padding_idx=0)
        # Optionally add position embeddings
        # self.position_embedding = nn.Embedding(4096, config.embed_size)
        
        self.transformer = nn.ModuleList(
            [TransformerLayer(
                embed_size=config.embed_size,
                num_heads=config.num_heads,
                hidden_size=config.hidden_size,
                memory_size=config.memory_size,
                memory_dim=config.memory_dim,
                guidance_dim=config.guidance_dim
            ) for _ in range(config.num_layers)]
        )
        self.decoder = nn.Linear(config.embed_size, config.vocab_size)

        # Initialize weights
        self.init_weights()

    def forward(self, input_ids, attention_mask=None):
        """
        input_ids: (B, Q)
        attention_mask: (B, Q) where 1 indicates valid tokens and 0 indicates padding
        Returns:
            logits: (B, Q+M, V)
        """
        batch_size, seq_len = input_ids.size()

        # Embeddings
        token_embeddings = self.token_embedding(input_ids)  # (B, Q, D)
        
        # Optionally add position embeddings
        # position_ids = torch.arange(seq_len, dtype=torch.long, device=input_ids.device).unsqueeze(0).expand(batch_size, -1)
        # pos_embeddings = self.position_embedding(position_ids)
        # embeddings = token_embeddings + pos_embeddings
        embeddings = token_embeddings

        out = embeddings
        for layer in self.transformer:
            out = layer(out, mask=attention_mask)  # (B, Q+M, D)

        # Decoder
        logits = self.decoder(out)  # (B, Q+M, V)

        return logits

    def _init_weights(self, module):
        """Initialize the weights."""
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            if module.bias is not None:
                module.bias.data.zero_()
            if module.weight is not None:
                module.weight.data.fill_(1.0)