"""Query-based decoder for inverting mean-pooled DNA embeddings. Designed for the multi-length inversion-attack setting: a single fixed-size mean embedding must be decoded into a variable-length nucleotide sequence whose length itself is unknown. Architecture (mean mode only): 1. ``context_projection`` maps the mean embedding to ``n_context_tokens`` learnable summary tokens of size ``d_model``. These are the "memory" that cross-attention attends to. 2. ``length_head`` predicts a distribution over possible content lengths (0..seq_length-1) directly from the mean embedding. 3. The soft length distribution is mixed with ``length_embedding`` to produce a single conditioning vector that is added to every position query. This makes length prediction a first-class input to content generation, not a side branch: gradients flow from the sequence loss back through the length head, and the decoder is trained to produce EOS at the predicted position. 4. ``pos_queries`` are learnable per-output-slot embeddings. Combined with the length conditioning, they form the decoder inputs. 5. A stack of ``TransformerDecoderLayer`` blocks performs cross-attention (queries -> context) and self-attention (queries <-> queries). 6. ``output_projection`` produces per-position logits over the vocabulary (which includes the EOS symbol). Forward returns a tuple ``(seq_logits, length_logits)``. The training loop combines the sequence loss with a cross-entropy loss on the length head (weighted by ``aux_length_loss_weight``). At inference the length head's output is consumed only internally and is not exposed. """ from __future__ import annotations import logging from typing import List, Optional, Tuple import torch import torch.nn.functional as F from torch import nn class QueryDecoderReconstructor(nn.Module): def __init__( self, input_dim: int, mode: str, seq_length: int, output_dim: int, d_model: int = 256, nhead: int = 8, num_layers: int = 4, dim_feedforward: int = 1024, dropout: float = 0.1, n_context_tokens: int = 8, aux_length_loss_weight: float = 0.5, hidden_dims: Optional[List[int]] = None, ): super().__init__() assert input_dim > 0, "input_dim must be positive" assert mode == "mean", "QueryDecoderReconstructor supports mean mode only" assert seq_length > 0, "seq_length must be positive" assert output_dim > 0, "output_dim must be positive" assert d_model > 0, "d_model must be positive" assert nhead > 0 and d_model % nhead == 0, "nhead must divide d_model" assert num_layers > 0, "num_layers must be positive" assert dim_feedforward > 0, "dim_feedforward must be positive" assert 0.0 <= dropout <= 1.0, "dropout in [0, 1]" assert n_context_tokens > 0, "n_context_tokens must be positive" assert aux_length_loss_weight >= 0.0, "aux_length_loss_weight must be non-negative" self.input_dim = input_dim self.mode = mode self.seq_length = seq_length self.output_dim = output_dim self.d_model = d_model self.n_context_tokens = n_context_tokens self.aux_length_loss_weight = aux_length_loss_weight logger = logging.getLogger(__name__) logger.info( f"QueryDecoder: input_dim={input_dim} -> {n_context_tokens} context tokens " f"x d_model={d_model}, seq_length={seq_length} (incl. EOS slot), " f"output_dim={output_dim}, num_layers={num_layers}, nhead={nhead}, " f"aux_length_loss_weight={aux_length_loss_weight}" ) self.context_projection = nn.Linear(input_dim, n_context_tokens * d_model) self.length_head = nn.Sequential( nn.Linear(input_dim, d_model), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_model, seq_length), ) self.length_embedding = nn.Embedding(seq_length, d_model) self.pos_queries = nn.Embedding(seq_length, d_model) decoder_layer = nn.TransformerDecoderLayer( d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, dropout=dropout, activation="gelu", batch_first=True, norm_first=True, ) self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers) self.output_projection = nn.Linear(d_model, output_dim) nn.init.xavier_uniform_(self.output_projection.weight, gain=0.02) nn.init.zeros_(self.output_projection.bias) def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """Run the model. Parameters ---------- x : torch.Tensor Mean-pooled embeddings, shape (B, input_dim). Returns ------- seq_logits : torch.Tensor Per-position vocabulary logits, shape (B, seq_length, output_dim). length_logits : torch.Tensor Per-sample length logits, shape (B, seq_length). The class index equals the number of content tokens before the EOS slot. """ assert x.ndim == 2, f"Expected 2D mean embedding, got {x.ndim}D" B, D = x.shape assert D == self.input_dim, f"Expected input_dim={self.input_dim}, got {D}" context = self.context_projection(x).view(B, self.n_context_tokens, self.d_model) length_logits = self.length_head(x) length_probs = F.softmax(length_logits, dim=-1) length_emb = length_probs @ self.length_embedding.weight queries = self.pos_queries.weight.unsqueeze(0).expand(B, -1, -1) queries = queries + length_emb.unsqueeze(1) hidden = self.transformer_decoder(tgt=queries, memory=context) seq_logits = self.output_projection(hidden) return seq_logits, length_logits @staticmethod def length_targets_from_padded(batch_sequence: torch.Tensor) -> torch.Tensor: """Derive integer length targets from collated padded sequences. Each padded row is ``[content..., EOS, -100, -100, ...]``. The target we want the length head to predict is the count of content tokens before EOS, i.e. (non-pad count) - 1, clamped at 0 for the empty case. """ non_pad = (batch_sequence != -100).sum(dim=1) return (non_pad - 1).clamp(min=0)