Finsihed Decoder Implementation, as well as prediction heads and multitask

src/models/decoder.py

@@ -1,28 +1,16 @@
 """
-Transformer Decoder layout (Pre-LN)
-
-Contents:
-- create_causal_mask: utility to build a causal (subsequent) mask
-- TransformerDecoderLayer: one decoder block (masked self-attn, cross-attn, FFN)
-- TransformerDecoder: embedding/pos-encoding + stack of decoder layers + generation helpers
-
-Notes / conventions:
-- Pre-LN (LayerNorm before each sublayer) is assumed for stability (consistent with your encoder).
-- MultiHeadAttention, FeedForward, PositionalEncoding are expected to live in src/models
-  (you already implemented them).
-- Masks use boolean semantics: True = allowed, False = masked.
-- The decoder API supports:
-    - inputs: token ids (LongTensor, (B, T)) or embeddings ((B, T, d_model))
-    - memory: encoder outputs (B, S, d_model)
-    - mask arguments: tgt_mask (causal/padding), memory_mask (encoder padding)
-    - collect_attn: return attention maps per layer if requested
-- Generation helpers (greedy) are skeletons that you can extend to beam search or caching.
-
-TODO status keys:
-- [IMPLEMENT] : core implementation required
-- [OPTIONAL]  : useful enhancement (caching, beam search, advanced scheduling)
-"""
+Transformer Decoder (Pre-LN) - implementation.
+
+Implements:
+- create_causal_mask
+- TransformerDecoderLayer
+- TransformerDecoder (stack + naive greedy decoding)
 
+Conventions:
+- Masks are boolean: True = allowed, False = masked.
+- MultiHeadAttention expects masks broadcastable to (B, num_heads, T_q, T_k).
+- This decoder uses Pre-LN (LayerNorm before each sublayer).
+"""
 from typing import Optional, Tuple, List, Union, Dict
 import math
 import torch
@@ -35,47 +23,34 @@ from .positional_encoding import PositionalEncoding
 
 def create_causal_mask(seq_len: int, device: Optional[torch.device] = None) -> torch.Tensor:
     """
-    Create a square causal mask of shape (seq_len, seq_len).
-    True indicates allowed positions; False indicates masked (future) positions.
-
-    Returns:
-        mask: torch.BoolTensor of shape (seq_len, seq_len)
+    Create a (seq_len, seq_len) causal mask where entry (i, j) is True iff
+    j <= i (query at i may attend to keys up to i).
     """
-    # return a mask with True on and below diagonal, False above diagonal 
-    # The torch.trui function does masking, which is the idea of zeroing all the values in a matrix below a certain diagonal
-    mask = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1)
-    # mask has True above diagonal (to be masked). Want True for allowed, so invert:
-    return ~mask # shape (seq_len, seq_len) or (T, T)
+    # torch.triu(..., diagonal=1) is True above the diagonal. Invert to get allowed positions.
+    mask = ~torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1)
+    return mask  # shape: (T, T)
 
 
 class TransformerDecoderLayer(nn.Module):
     """
-    One decoder layer with:
-      - Masked self-attention (query/key/value = tgt)
-      - Encoder-Decoder cross-attention (query = tgt, key/value = memory)
-      - Position-wise FeedForward
-      - Pre-LN + residuals + dropout
-
-    Args:
-      d_model: model hidden size
-      num_heads: number of attention heads
-      d_ff: ff intermediate size
-      dropout: dropout for residuals / FFN
+    Single decoder layer (Pre-LN):
+      1) Masked self-attention
+      2) Cross-attention (encoder -> decoder)
+      3) Feed-forward
+    Returns the updated tgt and a dict of attention maps.
     """
 
     def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float = 0.1):
         super().__init__()
-        # NOTE: instantiate internal MHA with dropout=0.0 and manage dropout at layer-level
+        # use internal MHA dropout = 0.0; the layer handles dropout after sublayers
         self.self_attn = MultiHeadAttention(d_model=d_model, num_heads=num_heads, dropout=0.0)
         self.cross_attn = MultiHeadAttention(d_model=d_model, num_heads=num_heads, dropout=0.0)
         self.ffn = FeedForward(d_model=d_model, d_ff=d_ff, dropout=dropout)
 
-        # LayerNorms (Pre-LN)
         self.norm1 = nn.LayerNorm(d_model)
         self.norm2 = nn.LayerNorm(d_model)
         self.norm3 = nn.LayerNorm(d_model)
 
-        # Dropouts applied after sublayers (on sublayer outputs before residual add)
         self.dropout1 = nn.Dropout(dropout)
         self.dropout2 = nn.Dropout(dropout)
         self.dropout3 = nn.Dropout(dropout)
@@ -88,46 +63,51 @@ class TransformerDecoderLayer(nn.Module):
         memory_mask: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
         """
-        Forward pass for one decoder layer.
-
         Args:
-            tgt: (batch, tgt_len, d_model)
-            memory: (batch, src_len, d_model)  -- encoder outputs
-            tgt_mask: optional (batch, tgt_len, tgt_len) or (batch, 1, tgt_len, tgt_len)
-            memory_mask: optional (batch, src_len, src_len) or (batch, 1, tgt_len, src_len)
+            tgt: (B, T, d_model)
+            memory: (B, S, d_model)
+            tgt_mask: optional mask for self-attn - shape (B, T, T) or (B, 1, T, T)
+            memory_mask: optional mask for cross-attn - shape (B, S) or (B, 1, S) or (B, 1, T, S)
 
         Returns:
-            output: (batch, tgt_len, d_model)
-            attn_maps: dict with keys 'self' and 'cross' containing attention tensors
-                       shapes: (batch, num_heads, tgt_len, tgt_len) and (batch, num_heads, tgt_len, src_len)
+            (tgt_out, {"self": self_attn_weights, "cross": cross_attn_weights})
         """
-        # TODO [IMPLEMENT] Self-attention (Pre-LN)
-        # x_norm = self.norm1(tgt)
-        # self_out, self_attn = self.self_attn(x_norm, x_norm, x_norm, tgt_mask)
-        # tgt = tgt + self.dropout1(self_out)
+        # Ensure masks are on same device and boolean
+        if tgt_mask is not None:
+            tgt_mask = tgt_mask.to(dtype=torch.bool, device=tgt.device)
+        if memory_mask is not None:
+            memory_mask = memory_mask.to(dtype=torch.bool, device=tgt.device)
+            # If memory_mask is provided as (B, S) (per-key padding), expand to (B, 1, 1, S)
+            if memory_mask.dim() == 2:
+                memory_mask = memory_mask.unsqueeze(1).unsqueeze(1)  # (B,1,1,S)
+            # If it's (B, S, S) or (B, 1, S, S) leave as-is; if (B, T, S) convert to (B,1,T,S)
+            elif memory_mask.dim() == 3 and memory_mask.shape[1] != 1:
+                # assume (B, T, S) -> make (B, 1, T, S)
+                memory_mask = memory_mask.unsqueeze(1)
+
+        # --- Masked self-attention (Pre-LN) ---
+        x_norm = self.norm1(tgt)
+        self_out, self_attn = self.self_attn(x_norm, x_norm, x_norm, tgt_mask)
+        tgt = tgt + self.dropout1(self_out)
 
-        # TODO [IMPLEMENT] Cross-attention (Pre-LN)
-        # x_norm = self.norm2(tgt)
-        # cross_out, cross_attn = self.cross_attn(x_norm, memory, memory, memory_mask)
-        # tgt = tgt + self.dropout2(cross_out)
+        # --- Cross-attention (Pre-LN) ---
+        x_norm = self.norm2(tgt)
+        cross_out, cross_attn = self.cross_attn(x_norm, memory, memory, memory_mask)
+        tgt = tgt + self.dropout2(cross_out)
 
-        # TODO [IMPLEMENT] Feed-forward (Pre-LN)
-        # x_norm = self.norm3(tgt)
-        # ffn_out = self.ffn(x_norm)
-        # tgt = tgt + self.dropout3(ffn_out)
+        # --- Feed-forward (Pre-LN) ---
+        x_norm = self.norm3(tgt)
+        ffn_out = self.ffn(x_norm)
+        tgt = tgt + self.dropout3(ffn_out)
 
-        # TODO [RETURN] Return (tgt, {"self": self_attn, "cross": cross_attn})
-        raise NotImplementedError("TransformerDecoderLayer.forward: implement Pre-LN pipeline")
+        return tgt, {"self": self_attn, "cross": cross_attn}
 
 
 class TransformerDecoder(nn.Module):
     """
-    Full decoder: token embedding + positional encoding + stack of decoder layers.
-    Also supports simple greedy decoding.
+    Decoder stack with token embeddings and positional encoding.
 
-    Args:
-        vocab_size: for token embeddings (if using token ids)
-        d_model, num_layers, num_heads, d_ff, dropout, max_len, pad_token_id: same semantics as encoder
+    Forward returns logits (B, T, vocab_size) by default; if collect_attn=True returns (logits, attn_list).
     """
 
     def __init__(
@@ -146,37 +126,25 @@ class TransformerDecoder(nn.Module):
         self.d_model = d_model
         self.pad_token_id = pad_token_id
 
-        # Token embedding (used if inputs are token ids)
         self.embedding = nn.Embedding(vocab_size, d_model)
-
-        # Positional encoding
         self.pos_encoder = PositionalEncoding(d_model=d_model, max_len=max_len, dropout=dropout)
 
-        # Decoder layers
         self.layers = nn.ModuleList(
-            [
-                TransformerDecoderLayer(d_model=d_model, num_heads=num_heads, d_ff=d_ff, dropout=dropout)
-                for _ in range(num_layers)
-            ]
+            [TransformerDecoderLayer(d_model=d_model, num_heads=num_heads, d_ff=d_ff, dropout=dropout)
+             for _ in range(num_layers)]
         )
 
-        # Final layer norm for Pre-LN stacks
         self.final_norm = nn.LayerNorm(d_model)
-
-        # Output projection to vocabulary (logits)
         self.output_projection = nn.Linear(d_model, vocab_size)
-
-        # Input dropout (after pos encoding)
         self.input_dropout = nn.Dropout(dropout)
 
     def _build_padding_mask_from_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         """
-        Build (batch, seq, seq) boolean mask from input ids and pad_token_id.
-        True = allowed, False = masked.
+        Convert input ids to (B, T, T) boolean mask where True = allowed.
         """
         assert self.pad_token_id is not None, "pad_token_id must be set to build mask from ids"
-        pad_mask = (input_ids != self.pad_token_id)  # (B, S)
-        attn_mask = pad_mask.unsqueeze(1) & pad_mask.unsqueeze(2)  # (B, S, S)
+        pad_mask = (input_ids != self.pad_token_id)  # (B, T)
+        attn_mask = pad_mask.unsqueeze(1) & pad_mask.unsqueeze(2)  # (B, T, T)
         return attn_mask
 
     def forward(
@@ -188,21 +156,13 @@ class TransformerDecoder(nn.Module):
         collect_attn: bool = False,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[Dict[str, torch.Tensor]]]]:
         """
-        Forward pass for the decoder stack.
-
         Args:
-            inputs: token ids (B, T) or embeddings (B, T, d_model)
-            memory: encoder outputs (B, S, d_model)
-            tgt_mask: optional mask for decoder self-attention. If None, a causal mask will be created.
-                      Mask shapes: (B, T, T) or (B, 1, T, T)
-            memory_mask: optional mask over memory (B, S, S) or (B, 1, T, S)
-            collect_attn: if True returns (logits/outputs, [per-layer attn dicts])
-
-        Returns:
-            logits: (B, T, vocab_size) or (B, T, d_model) if you prefer returning hidden states
-            or (logits, attn_list) if collect_attn True
+            inputs: (B, T) token ids or (B, T, d_model) embeddings
+            memory: (B, S, d_model)
+            tgt_mask: optional; if None, will create (causal [+ padding if ids available])
+            memory_mask: optional; if provided as (B, S) will be expanded to (B, 1, 1, S)
         """
-        # Inputs: if token ids, embed and scale; else assume embeddings
+        # Prepare embeddings
         if inputs.dim() == 2:  # token ids
             x = self.embedding(inputs) * math.sqrt(self.d_model)
         elif inputs.dim() == 3:
@@ -210,47 +170,48 @@ class TransformerDecoder(nn.Module):
         else:
             raise ValueError("inputs must be (B, T) token ids or (B, T, d_model) embeddings")
 
-        # Positional encoding + dropout
         x = self.pos_encoder(x)
         x = self.input_dropout(x)
 
-        # Build tgt_mask if not provided: combine causal mask and padding mask if available
-        seq_len = x.size(1)
+        B, T, _ = x.shape
+
+        # Build target mask if not provided: combine causal + padding (if available)
         if tgt_mask is None:
-            # base causal mask (T, T)
-            causal = create_causal_mask(seq_len, device=x.device)  # [TODO implement]
-            # expand to batch dim later if padding present
+            causal = create_causal_mask(T, device=x.device)  # (T, T)
             if inputs.dim() == 2 and self.pad_token_id is not None:
-                padding_mask = self._build_padding_mask_from_ids(inputs)  # (B, T, T)
-                # combine: True only where both causal and padding allow attention
-                # TODO: ensure shapes align; broadcast causal to (1, T, T) then & with padding_mask
-                raise NotImplementedError("tgt_mask construction: combine causal + padding_mask")
+                pad_pairwise = self._build_padding_mask_from_ids(inputs)  # (B, T, T)
+                combined = pad_pairwise & causal.unsqueeze(0)  # (B, T, T)
+                tgt_mask = combined.unsqueeze(1)  # (B, 1, T, T) -> broadcast to heads
             else:
-                # TODO: Broadcast causal to (1, T, T) or (B, 1, T, T) depending on downstream expectations
-                raise NotImplementedError("tgt_mask construction: broadcast causal mask for batch")
+                # No per-batch padding info: broadcast causal to (1, 1, T, T)
+                tgt_mask = causal.unsqueeze(0).unsqueeze(1)  # (1, 1, T, T)
+        else:
+            # Ensure boolean and device alignment; accept (B, T, T) or (B,1,T,T) or (1,1,T,T)
+            tgt_mask = tgt_mask.to(dtype=torch.bool, device=x.device)
 
-        # Ensure memory_mask is boolean on correct device if provided
+        # Normalize memory_mask dtype/device and expand simple shapes
         if memory_mask is not None:
             memory_mask = memory_mask.to(dtype=torch.bool, device=x.device)
+            if memory_mask.dim() == 2:  # (B, S) -> (B, 1, 1, S)
+                memory_mask = memory_mask.unsqueeze(1).unsqueeze(1)
+            elif memory_mask.dim() == 3:  # (B, T, S) -> (B, 1, T, S)
+                memory_mask = memory_mask.unsqueeze(1)
 
         attn_list: List[Dict[str, torch.Tensor]] = []
 
-        # Pass through layers
+        # Pass through decoder layers
         for layer in self.layers:
             x, attn = layer(x, memory, tgt_mask=tgt_mask, memory_mask=memory_mask)
             if collect_attn:
                 attn_list.append(attn)
 
-        x = self.final_norm(x)  # Pre-LN final normalization
-
+        x = self.final_norm(x)
         logits = self.output_projection(x)  # (B, T, vocab)
+
         if collect_attn:
             return logits, attn_list
         return logits
 
-    # ---------------------------------------------------------------------
-    # Generation / inference helpers (skeletons)
-    # ---------------------------------------------------------------------
     def greedy_decode(
         self,
         memory: torch.Tensor,
@@ -258,26 +219,32 @@ class TransformerDecoder(nn.Module):
         start_token_id: int,
         end_token_id: Optional[int] = None,
         device: Optional[torch.device] = None,
-    ) -> torch.LongTensor:
+    ) -> torch.Tensor:
         """
-        Greedy autoregressive decoding using the decoder stack.
-
-        Args:
-            memory: encoder outputs (B, S, d_model)
-            max_len: maximum target length to generate
-            start_token_id: BOS token id
-            end_token_id: optional EOS token id to stop early
-        Returns:
-            generated: (B, T_out) long tensor of token ids
+        Naive greedy decoding: repeatedly run the decoder on the growing prefix.
+        Not optimized (recomputes full decoder each step) but simple and correct.
         """
-        # TODO [IMPLEMENT]:
-        #  - Start with tensor of shape (B, 1) filled with start_token_id
-        #  - Repeatedly call decoder.forward in incremental mode (or full forward with causal mask)
-        #  - At each step pick argmax over logits and append to sequence
-        #  - Stop if all sequences produced end_token_id or max_len reached
-        raise NotImplementedError("greedy_decode: implement autoregressive generation loop")
-
-    # Optional: incremental step method with caching of past keys/values for speed
+        if device is None:
+            device = memory.device
+        B = memory.size(0)
+        generated = torch.full((B, 1), start_token_id, dtype=torch.long, device=device)
+
+        for _ in range(max_len - 1):
+            logits = self.forward(generated, memory, collect_attn=False)  # (B, L, V)
+            assert isinstance(logits, torch.Tensor)  # type narrowing
+            next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)  # (B, 1)
+            generated = torch.cat([generated, next_token], dim=1)
+
+            if end_token_id is not None:
+                # stop if all sequences ended
+                if (generated[:, -1] == end_token_id).all():
+                    break
+
+        return generated
+
+    # -----------------------------
+    # Incremental single-step API
+    # -----------------------------
     def step(
         self,
         last_token_ids: torch.LongTensor,
@@ -285,16 +252,141 @@ class TransformerDecoder(nn.Module):
         cache: Optional[Dict] = None,
     ) -> Tuple[torch.Tensor, Dict]:
         """
-        Single-step decoder that returns logits for the next token given last_token_ids.
+        Run one autoregressive step.
 
         Args:
-            last_token_ids: (B, 1) tokens at current time step
-            memory: encoder outputs
-            cache: optional dict storing per-layer cached keys/values
+            last_token_ids: (B, 1) last generated token ids
+            memory: encoder outputs (B, S, d_model)
+            cache: optional dict with previous cached keys/values and 'past_length'.
 
         Returns:
-            logits: (B, vocab_size)
-            new_cache: updated cache
+            logits: (B, vocab_size) logits for the next-token prediction
+            new_cache: updated cache dictionary
         """
-        # TODO [OPTIONAL]: implement fast incremental decoding caching keys/values per layer
-        raise NotImplementedError("step: incremental decoding (optional optimization)")
+        device = memory.device
+        B = last_token_ids.size(0)
+
+        if cache is None:
+            cache = {}
+        past_len = int(cache.get("past_length", 0))
+
+        # 1) Embed last token and add positional encoding for position `past_len`
+        x = self.embedding(last_token_ids) * math.sqrt(self.d_model)  # (B,1,d)
+        # Use positional encoding buffer directly (avoid dropout in pos_encoder)
+        # pos_encoder.pe expected shape (1, max_len, d_model)
+        if hasattr(self.pos_encoder, "pe"):
+            pe = self.pos_encoder.pe  # (1, max_len, d_model)
+            pos_idx = past_len
+            if pos_idx >= pe.size(1):
+                raise RuntimeError(f"pos_idx {pos_idx} exceeds max_len {pe.size(1)}")
+            x = x + pe[:, pos_idx:pos_idx + 1, :].to(device)
+        else:
+            # fallback: call pos_encoder and rely on its dropout (less ideal)
+            x = self.pos_encoder(x)
+
+        # We will update new_cache incrementally
+        new_cache = dict(cache)  # shallow copy
+        new_cache["past_length"] = past_len + 1
+
+        # Optional: memory_mask could be supplied in cache under 'memory_mask'
+        memory_mask = new_cache.get("memory_mask", None)
+        if memory_mask is not None:
+            memory_mask = memory_mask.to(dtype=torch.bool, device=device)
+            # expand (B, S) -> (B,1,1,S) if necessary
+            if memory_mask.dim() == 2:
+                memory_mask = memory_mask.unsqueeze(1).unsqueeze(1)
+            elif memory_mask.dim() == 3:
+                memory_mask = memory_mask.unsqueeze(1)
+
+        # Iterate layers, updating caches and computing output for current token only
+        layer_input = x  # (B,1,d_model)
+        for i, layer in enumerate(self.layers):
+            # -------------------
+            # 1) Self-attention (incremental)
+            # -------------------
+            # Normalize input for pre-LN
+            x_norm = layer.norm1(layer_input)  # (B,1,d)
+
+            # Project Q,K,V for the new token
+            Q_new = layer.self_attn.W_Q(x_norm)  # (B,1,d_model)
+            K_new = layer.self_attn.W_K(x_norm)
+            V_new = layer.self_attn.W_V(x_norm)
+
+            # Reshape into heads: (B, num_heads, 1, d_k)
+            B_, Lq, _ = Q_new.shape
+            num_heads = layer.self_attn.num_heads
+            d_k = layer.self_attn.d_k
+            Qh = Q_new.view(B_, Lq, num_heads, d_k).transpose(1, 2)  # (B, num_heads, 1, d_k)
+            Kh = K_new.view(B_, Lq, num_heads, d_k).transpose(1, 2)
+            Vh = V_new.view(B_, Lq, num_heads, d_k).transpose(1, 2)
+
+            # Retrieve cached keys/values for self-attn (if exist)
+            cache_k = cache.get(f"self_k_{i}", None)
+            cache_v = cache.get(f"self_v_{i}", None)
+            if cache_k is None or cache_v is None:
+                K_all = Kh  # (B, H, 1, d_k)
+                V_all = Vh
+            else:
+                # concat along sequence dim (dim=2)
+                K_all = torch.cat([cache_k.to(device), Kh], dim=2)
+                V_all = torch.cat([cache_v.to(device), Vh], dim=2)
+
+            # Store updated caches
+            new_cache[f"self_k_{i}"] = K_all
+            new_cache[f"self_v_{i}"] = V_all
+
+            # Compute attention for the new token: Query length = 1, Key length = K_all.size(2)
+            attn_out_heads, self_attn_w = layer.self_attn.attention(Qh, K_all, V_all, mask=None)
+            # attn_out_heads: (B, H, 1, d_k)
+            # concat heads, project out
+            attn_out = attn_out_heads.transpose(1, 2).contiguous().view(B_, 1, num_heads * d_k)
+            attn_out = layer.self_attn.W_O(attn_out)  # (B,1,d_model)
+            layer_output = layer_input + layer.dropout1(attn_out)
+
+            # -------------------
+            # 2) Cross-attention (use cached memory projections if available)
+            # -------------------
+            x_norm2 = layer.norm2(layer_output)  # (B,1,d)
+            # Ensure memory K/V are cached per layer
+            mem_k = cache.get(f"mem_k_{i}", None)
+            mem_v = cache.get(f"mem_v_{i}", None)
+            if mem_k is None or mem_v is None:
+                # project memory once for this layer and cache it
+                # memory: (B, S, d_model)
+                MK = layer.cross_attn.W_K(memory)  # (B, S, d_model)
+                MV = layer.cross_attn.W_V(memory)
+                Bm, S, _ = MK.shape
+                MKh = MK.view(Bm, S, layer.cross_attn.num_heads, layer.cross_attn.d_k).transpose(1, 2)  # (B,H,S,d_k)
+                MVh = MV.view(Bm, S, layer.cross_attn.num_heads, layer.cross_attn.d_k).transpose(1, 2)
+                mem_k = MKh
+                mem_v = MVh
+                new_cache[f"mem_k_{i}"] = mem_k
+                new_cache[f"mem_v_{i}"] = mem_v
+            else:
+                mem_k = mem_k.to(device)
+                mem_v = mem_v.to(device)
+
+            Qc = layer.cross_attn.W_Q(x_norm2)  # (B,1,d_model)
+            Qch = Qc.view(B, 1, layer.cross_attn.num_heads, layer.cross_attn.d_k).transpose(1, 2)  # (B,H,1,d_k)
+
+            cross_out_heads, cross_attn_w = layer.cross_attn.attention(Qch, mem_k, mem_v, mask=memory_mask)
+            cross_out = cross_out_heads.transpose(1, 2).contiguous().view(B, 1, layer.cross_attn.num_heads * layer.cross_attn.d_k)
+            cross_out = layer.cross_attn.W_O(cross_out)  # (B,1,d_model)
+            layer_output = layer_output + layer.dropout2(cross_out)
+
+            # -------------------
+            # 3) Feed-forward (incremental)
+            # -------------------
+            x_norm3 = layer.norm3(layer_output)
+            ffn_out = layer.ffn(x_norm3)  # (B,1,d_model)
+            layer_output = layer_output + layer.dropout3(ffn_out)
+
+            # Prepare for next layer
+            layer_input = layer_output
+
+        # Final norm + output projection (for this single time step)
+        out_norm = self.final_norm(layer_input)  # (B,1,d_model)
+        logits = self.output_projection(out_norm)  # (B,1,vocab)
+        logits = logits.squeeze(1)  # (B, vocab)
+
+        return logits, new_cache

src/models/heads.py

@@ -0,0 +1,151 @@
+"""
+Prediction heads for Transformer models.
+
+Includes:
+- ClassificationHead: sequence-level classification with simple pooling (mean/cls/max).
+- TokenClassificationHead: per-token classification (e.g., NER).
+- LMHead: language-modeling head mapping hidden states to vocabulary logits. Optional weight tying to an Embedding.
+- ProjectionHead: small projection MLP for representation learning / contrastive heads.
+
+Keep these heads minimal, well-tested, and easy to compose on top of encoder/decoder outputs.
+"""
+from typing import Optional, Literal
+
+import torch
+import torch.nn as nn
+
+
+class ClassificationHead(nn.Module):
+    """
+    Sequence-level classification head.
+
+    Args:
+        d_model: hidden size from encoder/decoder
+        num_labels: number of output classes
+        pooler: one of 'mean', 'cls', 'max' - how to pool the sequence
+        dropout: dropout probability before final linear layer
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_labels: int,
+        pooler: Literal["mean", "cls", "max"] = "mean",
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        assert pooler in ("mean", "cls", "max"), "pooler must be 'mean'|'cls'|'max'"
+        self.pooler = pooler
+        self.dropout = nn.Dropout(dropout)
+        self.out_proj = nn.Linear(d_model, num_labels)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x: (batch, seq_len, d_model)
+        returns: (batch, num_labels)
+        """
+        if self.pooler == "mean":
+            pooled = x.mean(dim=1)
+        elif self.pooler == "cls":
+            pooled = x[:, 0, :]
+        else:  # max
+            pooled, _ = x.max(dim=1)
+        pooled = self.dropout(pooled)
+        return self.out_proj(pooled)
+
+
+class TokenClassificationHead(nn.Module):
+    """
+    Per-token classification head. Useful for NER, POS, etc.
+
+    Args:
+        d_model: hidden size
+        num_labels: number of per-token classes
+        dropout: dropout probability applied before the linear layer
+    """
+
+    def __init__(self, d_model: int, num_labels: int, dropout: float = 0.1):
+        super().__init__()
+        self.dropout = nn.Dropout(dropout)
+        self.out_proj = nn.Linear(d_model, num_labels)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x: (batch, seq_len, d_model)
+        returns: (batch, seq_len, num_labels)
+        """
+        x = self.dropout(x)
+        return self.out_proj(x)
+
+
+class LMHead(nn.Module):
+    """
+    Language modeling head: maps hidden states to logits over vocabulary.
+
+    Args:
+        d_model: hidden size
+        vocab_size: vocabulary size
+        tie_embedding: optional nn.Embedding instance to tie weights with
+    """
+
+    def __init__(self, d_model: int, vocab_size: int, tie_embedding: Optional[nn.Embedding] = None):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.proj = nn.Linear(d_model, vocab_size, bias=True)
+
+        if tie_embedding is not None:
+            # Validate sizes
+            assert tie_embedding.num_embeddings == vocab_size, "vocab size mismatch for weight tying"
+            assert tie_embedding.embedding_dim == d_model, "embedding dim must match d_model for weight tying"
+            # Tie weights: point the projection weight to the embedding weight Tensor
+            # Remove the existing projection parameter in favor of the embedding weight
+            # This keeps the same Parameter object, so updates affect both modules.
+            self.proj.weight = tie_embedding.weight
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """
+        hidden_states: (batch, seq_len, d_model)
+        returns logits: (batch, seq_len, vocab_size)
+        """
+        return self.proj(hidden_states)
+
+
+class ProjectionHead(nn.Module):
+    """
+    Simple projection head for representation learning.
+
+    Args:
+        d_model: input dimension
+        proj_dim: output projection dimension
+        hidden_dim: intermediate dimension (optional)
+        dropout: dropout probability
+    """
+
+    def __init__(self, d_model: int, proj_dim: int = 128, hidden_dim: Optional[int] = None, dropout: float = 0.1):
+        super().__init__()
+        if hidden_dim is None:
+            hidden_dim = max(d_model, proj_dim)
+        self.net = nn.Sequential(
+            nn.Linear(d_model, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, proj_dim),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x: (batch, d_model) or (batch, seq_len, d_model) - both supported.
+        Returns:
+            If input is 3D: (batch, seq_len, proj_dim)
+            If input is 2D: (batch, proj_dim)
+        """
+        orig_dim = x.dim()
+        if orig_dim == 3:
+            B, T, D = x.shape
+            out = self.net(x.view(B * T, D))
+            return out.view(B, T, -1)
+        elif orig_dim == 2:
+            return self.net(x)
+        else:
+            raise ValueError("Input must be 2D or 3D tensor")

src/models/multitask.py

@@ -0,0 +1,198 @@
+"""
+Multitask model composition utilities.
+
+Provides:
+- MultiTaskModel: lightweight wrapper to compose an encoder and/or decoder with
+  multiple task heads (classification, token classification, LM head, etc.)
+- add_head / remove_head helpers
+- forward(task_name, ...) that routes inputs to the correct sub-modules
+- compute_loss helper that uses common losses and ignore_index support
+
+Design goals:
+- Keep composition simple and explicit (use named heads per task)
+- Support encoder-only tasks (classification, token classification) and
+  seq2seq tasks (encoder -> decoder -> LMHead)
+- Minimal dependencies on training loop; return logits and (optionally) loss
+"""
+from typing import Optional, Dict, Any, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import your components
+from .encoder import TransformerEncoder
+from .decoder import TransformerDecoder
+from .heads import ClassificationHead, TokenClassificationHead, LMHead
+
+
+class MultiTaskModel(nn.Module):
+    """
+    Compose encoder/decoder and task heads.
+
+    Usage patterns:
+    - Encoder-only classification:
+        mt = MultiTaskModel(encoder=enc)
+        mt.add_head("sentiment", ClassificationHead(...))
+        logits = mt.forward("sentiment", {"input_ids": src_ids})
+    - Seq2seq LM:
+        mt = MultiTaskModel(encoder=enc, decoder=dec)
+        mt.add_head("summarize", LMHead(...))
+        logits = mt.forward("summarize", {"src_ids": src_ids, "tgt_ids": tgt_ids})
+    """
+
+    def __init__(
+        self,
+        encoder: Optional[TransformerEncoder] = None,
+        decoder: Optional[TransformerDecoder] = None,
+    ):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.heads: Dict[str, nn.Module] = {}
+
+    def add_head(self, name: str, module: nn.Module) -> None:
+        """Register a head under a task name."""
+        if name in self.heads:
+            raise ValueError(f"Head '{name}' already exists")
+        self.heads[name] = module
+        self.add_module(f"head_{name}", module)
+
+    def remove_head(self, name: str) -> None:
+        """Remove a registered head."""
+        if name not in self.heads:
+            raise KeyError(name)
+        del self._modules[f"head_{name}"]
+        del self.heads[name]
+
+    def forward(
+        self,
+        task: str,
+        inputs: Dict[str, torch.Tensor],
+        return_loss: bool = False,
+        loss_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> Any:
+        """
+        Route inputs to appropriate model components and head.
+
+        Args:
+            task: registered head name
+            inputs: dictionary; common keys:
+                - For encoder tasks: "input_ids" or "embeddings" (B, S) or (B, S, d)
+                - For seq2seq: "src_ids" (B,S) or "src_embeddings", and "tgt_ids" (B,T) or "tgt_embeddings"
+                            when computing training loss, pass "labels" (B,T) for LM
+            return_loss: if True and labels provided, returns (loss, logits)
+            loss_kwargs: forwarded to compute_loss (e.g., ignore_index)
+
+        Returns:
+            logits (or (loss, logits) if return_loss True)
+        """
+        if task not in self.heads:
+            raise KeyError(f"Unknown task/head '{task}'")
+
+        head = self.heads[task]
+        loss_kwargs = loss_kwargs or {}
+
+        # Encoder-only heads expect encoder outputs
+        if isinstance(head, (ClassificationHead, TokenClassificationHead)):
+            if self.encoder is None:
+                raise RuntimeError("Encoder is required for encoder-side heads")
+            # accept either input_ids or embeddings
+            if "input_ids" in inputs:
+                enc_out = self.encoder(inputs["input_ids"])
+            elif "embeddings" in inputs:
+                enc_out = self.encoder(inputs["embeddings"])
+            else:
+                raise ValueError("inputs must contain 'input_ids' or 'embeddings' for encoder tasks")
+            logits = head(enc_out)
+
+            if return_loss:
+                labels = inputs.get("labels", None)
+                if labels is None:
+                    raise ValueError("return_loss=True requires 'labels' in inputs")
+                loss = self.compute_loss_for_head(head, logits, labels, **loss_kwargs)
+                return loss, logits
+            return logits
+
+        # LM/seq2seq head: run encoder -> decoder -> lm head
+        if isinstance(head, LMHead):
+            if self.encoder is None or self.decoder is None:
+                raise RuntimeError("Both encoder and decoder are required for LM-style heads")
+
+            # Build encoder memory
+            if "src_ids" in inputs:
+                memory = self.encoder(inputs["src_ids"])
+            elif "src_embeddings" in inputs:
+                memory = self.encoder(inputs["src_embeddings"])
+            else:
+                raise ValueError("inputs must contain 'src_ids' or 'src_embeddings' for seq2seq tasks")
+
+            # If training / teacher forcing: expect tgt_ids (shifted by caller) or embeddings
+            if "tgt_ids" in inputs:
+                decoder_inputs = inputs["tgt_ids"]
+            elif "tgt_embeddings" in inputs:
+                decoder_inputs = inputs["tgt_embeddings"]
+            else:
+                # For generation time you may call decoder.greedy_decode separately.
+                # Here we don't attempt to generate when labels not provided.
+                raise ValueError("Seq2seq tasks require 'tgt_ids' or 'tgt_embeddings' for training forward")
+
+            # Run decoder. Decoder returns logits shaped (B, T, vocab) in this codebase.
+            decoder_out = self.decoder(decoder_inputs, memory)
+
+            # If decoder already returned logits matching the head vocab size, use them directly.
+            # Otherwise, assume decoder returned hidden states and let the head project them.
+            if isinstance(decoder_out, torch.Tensor) and decoder_out.shape[-1] == head.vocab_size:
+                logits = decoder_out
+            else:
+                logits = head(decoder_out)
+
+            if return_loss:
+                labels = inputs.get("labels", None)
+                if labels is None:
+                    raise ValueError("return_loss=True requires 'labels' in inputs for seq2seq")
+                loss = self.compute_loss_for_head(head, logits, labels, **loss_kwargs)
+                return loss, logits
+            return logits
+
+        # Otherwise unsupported head type
+        raise RuntimeError(f"Unsupported head type: {type(head)}")
+
+    def compute_loss_for_head(
+        self,
+        head: nn.Module,
+        logits: torch.Tensor,
+        labels: torch.Tensor,
+        ignore_index: int = -100,
+    ) -> torch.Tensor:
+        """
+        Default loss dispatch:
+         - ClassificationHead: CrossEntropy on (B, num_labels)
+         - TokenClassificationHead: CrossEntropy per token (flattened)
+         - LMHead: CrossEntropy per token (flattened), ignore_index supported
+
+        Returns scalar loss.
+        """
+        if isinstance(head, ClassificationHead):
+            # logits: (B, num_labels) or (B, num_labels) direct
+            loss = F.cross_entropy(logits, labels.long())
+            return loss
+
+        if isinstance(head, TokenClassificationHead):
+            # logits: (B, T, C), labels: (B, T)
+            B, T, C = logits.shape
+            loss = F.cross_entropy(logits.view(B * T, C), labels.view(B * T).long(), ignore_index=ignore_index)
+            return loss
+
+        if isinstance(head, LMHead):
+            # logits: (B, T, V), labels: (B, T)
+            B, T, V = logits.shape
+            loss = F.cross_entropy(logits.view(B * T, V), labels.view(B * T).long(), ignore_index=ignore_index)
+            return loss
+
+        # Generic fall-back: try CrossEntropy on final dim
+        if logits.dim() == 2:
+            return F.cross_entropy(logits, labels.long())
+
+        # If we can't determine, raise
+        raise RuntimeError("Cannot compute loss for unknown head type")

tests/test_models/test_decoder.py

@@ -0,0 +1,152 @@
+import torch
+import pytest
+from src.models.decoder import (
+    create_causal_mask,
+    TransformerDecoderLayer,
+    TransformerDecoder,
+)
+
+
+def test_create_causal_mask_properties():
+    mask = create_causal_mask(5)
+    assert mask.shape == (5, 5)
+    # diagonal and below should be True
+    for i in range(5):
+        for j in range(5):
+            if j <= i:
+                assert mask[i, j].item() is True
+            else:
+                assert mask[i, j].item() is False
+
+
+def test_decoder_layer_shapes_and_grad():
+    torch.manual_seed(0)
+    d_model, num_heads, d_ff = 32, 4, 64
+    batch_size, tgt_len, src_len = 2, 6, 7
+
+    layer = TransformerDecoderLayer(d_model=d_model, num_heads=num_heads, d_ff=d_ff, dropout=0.0)
+    tgt = torch.randn(batch_size, tgt_len, d_model, requires_grad=True)
+    memory = torch.randn(batch_size, src_len, d_model)
+
+    # No masks
+    out, attn = layer(tgt, memory, tgt_mask=None, memory_mask=None)
+    assert out.shape == (batch_size, tgt_len, d_model)
+    assert isinstance(attn, dict)
+    assert "self" in attn and "cross" in attn
+    assert attn["self"].shape == (batch_size, num_heads, tgt_len, tgt_len)
+    assert attn["cross"].shape == (batch_size, num_heads, tgt_len, src_len)
+
+    # Backprop works
+    loss = out.sum()
+    loss.backward()
+    grads = [p.grad for p in layer.parameters() if p.requires_grad]
+    assert any(g is not None for g in grads)
+
+
+def test_decoder_layer_causal_mask_blocks_future():
+    torch.manual_seed(1)
+    d_model, num_heads, d_ff = 48, 6, 128
+    batch_size, tgt_len, src_len = 1, 5, 5
+
+    layer = TransformerDecoderLayer(d_model=d_model, num_heads=num_heads, d_ff=d_ff, dropout=0.0)
+    # create trivial increasing tgt embeddings so attention patterns are deterministic-ish
+    tgt = torch.randn(batch_size, tgt_len, d_model)
+    memory = torch.randn(batch_size, src_len, d_model)
+
+    causal = create_causal_mask(tgt_len, device=tgt.device)  # (T, T)
+    tgt_mask = causal.unsqueeze(0)  # (1, T, T) -> layer will handle unsqueeze to heads
+
+    out, attn = layer(tgt, memory, tgt_mask=tgt_mask, memory_mask=None)
+    self_attn = attn["self"].detach()
+    # Ensure upper triangle of attention weights is zero (no future attention)
+    # For each head and query i, keys j>i should be zero
+    B, H, Tq, Tk = self_attn.shape
+    for i in range(Tq):
+        for j in range(i + 1, Tk):
+            assert torch.allclose(self_attn[:, :, i, j], torch.zeros(B, H)), \
+                f"Found nonzero attention to future position {j} from query {i}"
+
+
+def test_decoder_stack_and_greedy_decode_shapes():
+    torch.manual_seed(2)
+    vocab_size = 30
+    d_model = 32
+    num_layers = 2
+    num_heads = 4
+    d_ff = 128
+    batch_size = 2
+    src_len = 7
+    max_tgt = 6
+
+    decoder = TransformerDecoder(
+        vocab_size=vocab_size,
+        d_model=d_model,
+        num_layers=num_layers,
+        num_heads=num_heads,
+        d_ff=d_ff,
+        dropout=0.0,
+        max_len=max_tgt,
+        pad_token_id=0,
+    )
+
+    # Random memory from encoder
+    memory = torch.randn(batch_size, src_len, d_model)
+
+    # Greedy decode: should produce (B, <= max_tgt)
+    generated = decoder.greedy_decode(memory, max_len=max_tgt, start_token_id=1, end_token_id=None)
+    assert generated.shape[0] == batch_size
+    assert generated.shape[1] <= max_tgt
+    assert (generated[:, 0] == 1).all()  # starts with start token
+
+    # Also test forward with embeddings and collect_attn
+    embeddings = torch.randn(batch_size, max_tgt, d_model)
+    logits, attn_list = decoder(embeddings, memory, collect_attn=True)
+    assert logits.shape == (batch_size, max_tgt, vocab_size)
+    assert isinstance(attn_list, list)
+    assert len(attn_list) == num_layers
+    for attn in attn_list:
+        assert "self" in attn and "cross" in attn
+
+
+def test_decoder_train_eval_dropout_behavior():
+    torch.manual_seed(3)
+    vocab_size = 40
+    d_model = 32
+    num_layers = 2
+    num_heads = 4
+    d_ff = 128
+    batch_size = 2
+    src_len = 6
+    tgt_len = 5
+
+    decoder = TransformerDecoder(
+        vocab_size=vocab_size,
+        d_model=d_model,
+        num_layers=num_layers,
+        num_heads=num_heads,
+        d_ff=d_ff,
+        dropout=0.4,
+        max_len=tgt_len,
+        pad_token_id=0,
+    )
+
+    # token ids with padding possible
+    input_ids = torch.randint(1, vocab_size, (batch_size, tgt_len), dtype=torch.long)
+    input_ids[0, -1] = 0
+
+    memory = torch.randn(batch_size, src_len, d_model)
+
+    decoder.train()
+    out1 = decoder(input_ids, memory)
+    out2 = decoder(input_ids, memory)
+    # With dropout in train mode, outputs should usually differ
+    assert not torch.allclose(out1, out2)
+
+    decoder.eval()
+    out3 = decoder(input_ids, memory)
+    out4 = decoder(input_ids, memory)
+    assert torch.allclose(out3, out4)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-q"])

tests/test_models/test_decoder_step.py

@@ -0,0 +1,98 @@
+import torch
+import pytest
+from typing import Any, Dict, cast
+from src.models.decoder import TransformerDecoder
+
+
+def test_step_equivalence_with_greedy_decode():
+    torch.manual_seed(7)
+    vocab_size = 25
+    d_model = 32
+    num_layers = 2
+    num_heads = 4
+    d_ff = 64
+    batch_size = 2
+    src_len = 6
+    max_tgt = 6
+
+    decoder = TransformerDecoder(
+        vocab_size=vocab_size,
+        d_model=d_model,
+        num_layers=num_layers,
+        num_heads=num_heads,
+        d_ff=d_ff,
+        dropout=0.0,
+        max_len=max_tgt,
+        pad_token_id=0,
+    )
+
+    memory = torch.randn(batch_size, src_len, d_model)
+
+    # 1) Get greedy sequence from naive greedy_decode
+    greedy = decoder.greedy_decode(memory, max_len=max_tgt, start_token_id=1, end_token_id=None)
+
+    # 2) Reproduce the same sequence with step() using cache
+    cache: Dict[str, Any] = {"past_length": 0}
+    generated = torch.full((batch_size, 1), 1, dtype=torch.long)
+    for _ in range(max_tgt - 1):
+        last_token = generated[:, -1:].to(memory.device)
+        logits, cache = decoder.step(cast(torch.LongTensor, last_token), memory, cache=cache)
+        next_token = logits.argmax(dim=-1, keepdim=True)
+        generated = torch.cat([generated, next_token], dim=1)
+
+    # Compare shapes & that sequences are identical
+    assert generated.shape == greedy.shape
+    assert torch.equal(generated, greedy)
+
+
+def test_step_cache_growth_and_shapes():
+    torch.manual_seed(9)
+    vocab_size = 20
+    d_model = 24
+    num_layers = 3
+    num_heads = 4
+    d_ff = 64
+    batch_size = 1
+    src_len = 5
+    steps = 4
+    max_tgt = 8
+
+    decoder = TransformerDecoder(
+        vocab_size=vocab_size,
+        d_model=d_model,
+        num_layers=num_layers,
+        num_heads=num_heads,
+        d_ff=d_ff,
+        dropout=0.0,
+        max_len=max_tgt,
+        pad_token_id=0,
+    )
+
+    memory = torch.randn(batch_size, src_len, d_model)
+
+    cache: Dict[str, Any] = {"past_length": 0}
+    last = torch.full((batch_size, 1), 1, dtype=torch.long)
+    for step_idx in range(steps):
+        logits, cache = decoder.step(cast(torch.LongTensor, last), memory, cache=cache)
+        # check updated past_length
+        assert cache["past_length"] == step_idx + 1
+        # check cached per-layer keys exist and have expected shape (B, H, seq_len, d_k)
+        for i in range(num_layers):
+            k = cache.get(f"self_k_{i}")
+            v = cache.get(f"self_v_{i}")
+            assert k is not None and v is not None
+            # seq_len should equal past_length
+            assert k.shape[2] == cache["past_length"]
+            # shapes match
+            assert k.shape[0] == batch_size
+            assert v.shape[0] == batch_size
+        # advance last token for next loop
+        last = logits.argmax(dim=-1, keepdim=True)
+
+    # Also ensure memory projections cached
+    for i in range(num_layers):
+        assert f"mem_k_{i}" in cache and f"mem_v_{i}" in cache
+        mem_k = cache[f"mem_k_{i}"]
+        mem_v = cache[f"mem_v_{i}"]
+        assert mem_k.shape[0] == batch_size
+        assert mem_k.shape[2] == src_len  # seq length of memory

tests/test_models/test_heads.py

@@ -0,0 +1,104 @@
+import torch
+import pytest
+import torch.nn as nn
+from src.models.heads import (
+    ClassificationHead,
+    TokenClassificationHead,
+    LMHead,
+    ProjectionHead,
+)
+
+
+def test_classification_head_shapes_and_dropout():
+    torch.manual_seed(0)
+    d_model = 64
+    num_labels = 5
+    batch_size = 3
+    seq_len = 10
+
+    head = ClassificationHead(d_model=d_model, num_labels=num_labels, pooler="mean", dropout=0.5)
+    head.train()
+    x = torch.randn(batch_size, seq_len, d_model)
+
+    out1 = head(x)
+    out2 = head(x)
+    # With dropout in train mode, outputs should usually differ
+    assert out1.shape == (batch_size, num_labels)
+    assert out2.shape == (batch_size, num_labels)
+    assert not torch.allclose(out1, out2)
+
+    head.eval()
+    out3 = head(x)
+    out4 = head(x)
+    assert torch.allclose(out3, out4), "Eval mode should be deterministic"
+
+
+def test_token_classification_head_shapes_and_grads():
+    torch.manual_seed(1)
+    d_model = 48
+    num_labels = 7
+    batch_size = 2
+    seq_len = 6
+
+    head = TokenClassificationHead(d_model=d_model, num_labels=num_labels, dropout=0.0)
+    x = torch.randn(batch_size, seq_len, d_model, requires_grad=True)
+    out = head(x)
+    assert out.shape == (batch_size, seq_len, num_labels)
+
+    loss = out.sum()
+    loss.backward()
+    grads = [p.grad for name, p in head.named_parameters() if p.requires_grad]
+    assert any(g is not None for g in grads)
+
+
+def test_lm_head_tie_weights_and_shapes():
+    torch.manual_seed(2)
+    vocab_size = 50
+    d_model = 32
+    batch_size = 2
+    seq_len = 4
+
+    embedding = nn.Embedding(vocab_size, d_model)
+    lm_tied = LMHead(d_model=d_model, vocab_size=vocab_size, tie_embedding=embedding)
+    lm_untied = LMHead(d_model=d_model, vocab_size=vocab_size, tie_embedding=None)
+
+    hidden = torch.randn(batch_size, seq_len, d_model)
+
+    # Shapes
+    logits_tied = lm_tied(hidden)
+    logits_untied = lm_untied(hidden)
+    assert logits_tied.shape == (batch_size, seq_len, vocab_size)
+    assert logits_untied.shape == (batch_size, seq_len, vocab_size)
+
+    # Weight tying: projection weight should be the same object as embedding.weight
+    assert lm_tied.proj.weight is embedding.weight
+
+    # Grad flows through tied weights
+    loss = logits_tied.sum()
+    loss.backward()
+    assert embedding.weight.grad is not None
+
+
+def test_projection_head_2d_and_3d_behavior_and_grad():
+    torch.manual_seed(3)
+    d_model = 40
+    proj_dim = 16
+    batch_size = 2
+    seq_len = 5
+
+    head = ProjectionHead(d_model=d_model, proj_dim=proj_dim, hidden_dim=64, dropout=0.0)
+    # 2D input
+    vec = torch.randn(batch_size, d_model, requires_grad=True)
+    out2 = head(vec)
+    assert out2.shape == (batch_size, proj_dim)
+
+    # 3D input
+    seq = torch.randn(batch_size, seq_len, d_model, requires_grad=True)
+    out3 = head(seq)
+    assert out3.shape == (batch_size, seq_len, proj_dim)
+
+    # Grad flow
+    loss = out3.sum()
+    loss.backward()
+    grads = [p.grad for p in head.parameters() if p.requires_grad]
+    assert any(g is not None for g in grads)

tests/test_models/test_multitask.py

@@ -0,0 +1,102 @@
+import torch
+import pytest
+from src.models.encoder import TransformerEncoder
+from src.models.decoder import TransformerDecoder
+from src.models.heads import ClassificationHead, LMHead, TokenClassificationHead
+from src.models.multitask import MultiTaskModel
+
+
+def test_multitask_encoder_classification_forward_and_loss():
+    torch.manual_seed(0)
+    vocab_size = 30
+    d_model = 32
+    num_layers = 2
+    num_heads = 4
+    d_ff = 64
+    batch_size = 3
+    seq_len = 8
+    num_labels = 5
+
+    enc = TransformerEncoder(vocab_size=vocab_size, d_model=d_model, num_layers=num_layers,
+                             num_heads=num_heads, d_ff=d_ff, dropout=0.0, max_len=seq_len, pad_token_id=0)
+
+    mt = MultiTaskModel(encoder=enc)
+    head = ClassificationHead(d_model=d_model, num_labels=num_labels, pooler="mean", dropout=0.0)
+    mt.add_head("sentiment", head)
+
+    input_ids = torch.randint(1, vocab_size, (batch_size, seq_len), dtype=torch.long)
+    labels = torch.randint(0, num_labels, (batch_size,), dtype=torch.long)
+
+    logits = mt.forward("sentiment", {"input_ids": input_ids})
+    assert logits.shape == (batch_size, num_labels)
+
+    loss, logits2 = mt.forward("sentiment", {"input_ids": input_ids, "labels": labels}, return_loss=True)
+    assert loss.item() >= 0
+    # grads
+    loss.backward()
+    grads = [p.grad for p in mt.parameters() if p.requires_grad]
+    assert any(g is not None for g in grads)
+
+
+def test_multitask_seq2seq_lm_forward_and_loss():
+    torch.manual_seed(1)
+    vocab_size = 40
+    d_model = 32
+    num_layers = 2
+    num_heads = 4
+    d_ff = 64
+    batch_size = 2
+    src_len = 7
+    tgt_len = 6
+
+    enc = TransformerEncoder(vocab_size=vocab_size, d_model=d_model, num_layers=num_layers,
+                             num_heads=num_heads, d_ff=d_ff, dropout=0.0, max_len=src_len, pad_token_id=0)
+    dec = TransformerDecoder(vocab_size=vocab_size, d_model=d_model, num_layers=num_layers,
+                             num_heads=num_heads, d_ff=d_ff, dropout=0.0, max_len=tgt_len, pad_token_id=0)
+    mt = MultiTaskModel(encoder=enc, decoder=dec)
+    lm_head = LMHead(d_model=d_model, vocab_size=vocab_size, tie_embedding=None)
+    mt.add_head("summarize", lm_head)
+
+    src_ids = torch.randint(1, vocab_size, (batch_size, src_len), dtype=torch.long)
+    # for training: provide decoder inputs (typically shifted right) and labels
+    tgt_ids = torch.randint(1, vocab_size, (batch_size, tgt_len), dtype=torch.long)
+    labels = tgt_ids.clone()
+
+    logits = mt.forward("summarize", {"src_ids": src_ids, "tgt_ids": tgt_ids})
+    assert logits.shape == (batch_size, tgt_len, vocab_size)
+
+    loss, logits2 = mt.forward("summarize", {"src_ids": src_ids, "tgt_ids": tgt_ids, "labels": labels}, return_loss=True)
+    assert loss.item() >= 0
+    loss.backward()
+    grads = [p.grad for p in mt.parameters() if p.requires_grad]
+    assert any(g is not None for g in grads)
+
+
+def test_token_classification_forward_and_loss():
+    torch.manual_seed(2)
+    vocab_size = 20
+    d_model = 24
+    num_layers = 2
+    num_heads = 4
+    d_ff = 64
+    batch_size = 2
+    seq_len = 5
+    num_labels = 7
+
+    enc = TransformerEncoder(vocab_size=vocab_size, d_model=d_model, num_layers=num_layers,
+                             num_heads=num_heads, d_ff=d_ff, dropout=0.0, max_len=seq_len, pad_token_id=0)
+    mt = MultiTaskModel(encoder=enc)
+    head = TokenClassificationHead(d_model=d_model, num_labels=num_labels, dropout=0.0)
+    mt.add_head("ner", head)
+
+    input_ids = torch.randint(1, vocab_size, (batch_size, seq_len), dtype=torch.long)
+    labels = torch.randint(0, num_labels, (batch_size, seq_len), dtype=torch.long)
+
+    logits = mt.forward("ner", {"input_ids": input_ids})
+    assert logits.shape == (batch_size, seq_len, num_labels)
+
+    loss, logits2 = mt.forward("ner", {"input_ids": input_ids, "labels": labels}, return_loss=True)
+    assert loss.item() >= 0
+    loss.backward()
+    grads = [p.grad for p in mt.parameters() if p.requires_grad]
+    assert any(g is not None for g in grads)