| """Set Transformer using only sequence embeddings and relative position.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import asdict, dataclass |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| @dataclass(frozen=True) |
| class ModelConfig: |
| architecture: str = "setnet" |
| esm_dimension: int = 1280 |
| pfam_dimension: int = 64 |
| pfam_vocab_size: int = 0 |
| hidden_dimension: int = 256 |
| output_dimension: int = 256 |
| attention_heads: int = 8 |
| inducing_points: int = 32 |
| feedforward_dimension: int = 512 |
| attention_blocks: int = 2 |
| position_bins: int = 64 |
| dropout: float = 0.1 |
|
|
| @classmethod |
| def from_dict(cls, values: dict[str, object]) -> "ModelConfig": |
| return cls(**{key: values[key] for key in asdict(cls()) if key in values}) |
|
|
|
|
| class GeneProjection(nn.Module): |
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| self.projection = nn.Linear(config.esm_dimension, config.hidden_dimension) |
| self.normalization = nn.LayerNorm(config.hidden_dimension) |
| self.position = nn.Embedding(config.position_bins, config.hidden_dimension) |
| self.dropout = nn.Dropout(config.dropout) |
| self.position_bins = config.position_bins |
|
|
| def forward(self, embeddings: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: |
| bins = (positions * (self.position_bins - 1)).long().clamp(0, self.position_bins - 1) |
| projected = F.gelu(self.normalization(self.projection(embeddings))) |
| return self.dropout(projected + self.position(bins)) |
|
|
|
|
| class PfamGeneProjection(nn.Module): |
| """Project ESM genes together with a BGC-level Pfam inventory embedding.""" |
|
|
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| if config.pfam_vocab_size < 3: |
| raise ValueError("Pfam vocabulary must contain padding, unknown, and one domain token") |
| self.pfam_embedding = nn.Embedding( |
| config.pfam_vocab_size, config.pfam_dimension, padding_idx=0 |
| ) |
| self.projection = nn.Linear( |
| config.esm_dimension + config.pfam_dimension, config.hidden_dimension |
| ) |
| self.normalization = nn.LayerNorm(config.hidden_dimension) |
| self.position = nn.Embedding(config.position_bins, config.hidden_dimension) |
| self.dropout = nn.Dropout(config.dropout) |
| self.position_bins = config.position_bins |
|
|
| def forward( |
| self, |
| embeddings: torch.Tensor, |
| positions: torch.Tensor, |
| pfam_tokens: torch.Tensor, |
| ) -> torch.Tensor: |
| bins = (positions * (self.position_bins - 1)).long().clamp(0, self.position_bins - 1) |
| token_mask = pfam_tokens.ne(0).unsqueeze(-1) |
| token_values = self.pfam_embedding(pfam_tokens).masked_fill(~token_mask, 0.0) |
| counts = token_mask.sum(dim=1).clamp_min(1) |
| pfam_summary = token_values.sum(dim=1) / counts |
| pfam_summary = pfam_summary.unsqueeze(1).expand(-1, embeddings.shape[1], -1) |
| merged = torch.cat((embeddings, pfam_summary), dim=-1) |
| projected = F.gelu(self.normalization(self.projection(merged))) |
| return self.dropout(projected + self.position(bins)) |
|
|
|
|
| class InducedSelfAttentionBlock(nn.Module): |
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| dimension = config.hidden_dimension |
| self.inducing = nn.Parameter(torch.empty(config.inducing_points, dimension)) |
| nn.init.xavier_uniform_(self.inducing) |
| self.inducing_norm = nn.LayerNorm(dimension) |
| self.input_norm = nn.LayerNorm(dimension) |
| self.summary_attention = nn.MultiheadAttention( |
| dimension, config.attention_heads, config.dropout, batch_first=True |
| ) |
| self.output_attention = nn.MultiheadAttention( |
| dimension, config.attention_heads, config.dropout, batch_first=True |
| ) |
| self.output_norm = nn.LayerNorm(dimension) |
| self.feedforward = nn.Sequential( |
| nn.Linear(dimension, config.feedforward_dimension), |
| nn.GELU(), |
| nn.Dropout(config.dropout), |
| nn.Linear(config.feedforward_dimension, dimension), |
| nn.Dropout(config.dropout), |
| ) |
|
|
| def forward(self, values: torch.Tensor, padding_mask: torch.Tensor | None) -> torch.Tensor: |
| batch_size = values.shape[0] |
| inducing = self.inducing.unsqueeze(0).expand(batch_size, -1, -1) |
| normalized_values = self.input_norm(values) |
| summary, _ = self.summary_attention( |
| inducing, normalized_values, normalized_values, key_padding_mask=padding_mask |
| ) |
| normalized_summary = self.inducing_norm(summary) |
| update, _ = self.output_attention( |
| normalized_values, normalized_summary, normalized_summary |
| ) |
| output = values + update |
| return output + self.feedforward(self.output_norm(output)) |
|
|
|
|
| class PoolingByAttention(nn.Module): |
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| self.seed = nn.Parameter(torch.empty(1, config.hidden_dimension)) |
| nn.init.xavier_uniform_(self.seed) |
| self.normalization = nn.LayerNorm(config.hidden_dimension) |
| self.attention = nn.MultiheadAttention( |
| config.hidden_dimension, config.attention_heads, config.dropout, batch_first=True |
| ) |
|
|
| def forward(self, values: torch.Tensor, padding_mask: torch.Tensor | None) -> torch.Tensor: |
| seed = self.seed.unsqueeze(0).expand(values.shape[0], -1, -1) |
| output, _ = self.attention( |
| seed, |
| self.normalization(values), |
| self.normalization(values), |
| key_padding_mask=padding_mask, |
| ) |
| return output.squeeze(1) |
|
|
|
|
| class LeakageFreeBGCSetNet(nn.Module): |
| """Map a variable-length BGC gene set to a normalized embedding.""" |
|
|
| input_names = ("gene_embeddings", "relative_positions", "padding_mask") |
|
|
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| self.config = config |
| self.gene_projection = GeneProjection(config) |
| self.blocks = nn.ModuleList( |
| InducedSelfAttentionBlock(config) for _ in range(config.attention_blocks) |
| ) |
| self.pooling = PoolingByAttention(config) |
| self.output = nn.Linear(config.hidden_dimension, config.output_dimension) |
|
|
| def encode_genes( |
| self, |
| gene_embeddings: torch.Tensor, |
| relative_positions: torch.Tensor, |
| padding_mask: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| values = self.gene_projection(gene_embeddings, relative_positions) |
| for block in self.blocks: |
| values = block(values, padding_mask) |
| return values |
|
|
| def forward( |
| self, |
| gene_embeddings: torch.Tensor, |
| relative_positions: torch.Tensor, |
| padding_mask: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| genes = self.encode_genes(gene_embeddings, relative_positions, padding_mask) |
| pooled = self.pooling(genes, padding_mask) |
| return F.normalize(self.output(pooled), p=2, dim=-1) |
|
|
|
|
| class MaskedGenePredictionHead(nn.Module): |
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| self.layers = nn.Sequential( |
| nn.Linear(config.hidden_dimension, config.feedforward_dimension), |
| nn.GELU(), |
| nn.Linear(config.feedforward_dimension, config.esm_dimension), |
| ) |
|
|
| def forward(self, contextual_embeddings: torch.Tensor) -> torch.Tensor: |
| return self.layers(contextual_embeddings) |
|
|
| class GatedDeepSets(nn.Module): |
| """Permutation-invariant learned pooling without gene-gene interactions.""" |
|
|
| input_names = ("gene_embeddings", "relative_positions", "padding_mask") |
|
|
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| self.config = config |
| self.gene_projection = GeneProjection(config) |
| self.gate = nn.Sequential( |
| nn.Linear(config.hidden_dimension, config.hidden_dimension // 2), |
| nn.GELU(), |
| nn.Linear(config.hidden_dimension // 2, 1), |
| ) |
| self.output = nn.Linear(config.hidden_dimension, config.output_dimension) |
|
|
| def encode_genes(self, gene_embeddings, relative_positions, padding_mask=None): |
| return self.gene_projection(gene_embeddings, relative_positions) |
|
|
| def forward(self, gene_embeddings, relative_positions, padding_mask=None): |
| genes = self.encode_genes(gene_embeddings, relative_positions, padding_mask) |
| logits = self.gate(genes).squeeze(-1) |
| if padding_mask is not None: |
| logits = logits.masked_fill(padding_mask, -torch.finfo(logits.dtype).max) |
| weights = torch.softmax(logits, dim=-1) |
| if padding_mask is not None: |
| weights = weights.masked_fill(padding_mask, 0.0) |
| pooled = (genes * weights.unsqueeze(-1)).sum(dim=1) |
| return F.normalize(self.output(pooled), p=2, dim=-1) |
|
|
|
|
| class RotarySelfAttentionBlock(nn.Module): |
| """Bidirectional self-attention with rotary position phases.""" |
|
|
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| d = config.hidden_dimension |
| self.heads = config.attention_heads |
| self.head_dim = d // self.heads |
| if self.head_dim % 2: |
| raise ValueError("Rotary attention requires an even per-head dimension") |
| self.norm = nn.LayerNorm(d) |
| self.qkv = nn.Linear(d, 3 * d) |
| self.output = nn.Linear(d, d) |
| self.dropout = nn.Dropout(config.dropout) |
| self.ffn_norm = nn.LayerNorm(d) |
| self.ffn = nn.Sequential( |
| nn.Linear(d, config.feedforward_dimension), nn.GELU(), |
| nn.Dropout(config.dropout), nn.Linear(config.feedforward_dimension, d), |
| nn.Dropout(config.dropout), |
| ) |
|
|
| def _rotate(self, values: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: |
| half = self.head_dim // 2 |
| inv = 1.0 / (10000.0 ** ( |
| torch.arange(half, device=values.device, dtype=values.dtype) / half |
| )) |
| angles = positions[:, None, :, None] * 128.0 * inv[None, None, None, :] |
| cos, sin = angles.cos(), angles.sin() |
| first, second = values[..., :half], values[..., half:] |
| return torch.cat((first * cos - second * sin, first * sin + second * cos), dim=-1) |
|
|
| def forward(self, values: torch.Tensor, positions: torch.Tensor, padding_mask=None) -> torch.Tensor: |
| normalized = self.norm(values) |
| batch, length, dimension = normalized.shape |
| qkv = self.qkv(normalized).view(batch, length, 3, self.heads, self.head_dim) |
| q, k, v = qkv.unbind(dim=2) |
| q, k = q.transpose(1, 2), k.transpose(1, 2) |
| v = v.transpose(1, 2) |
| q, k = self._rotate(q, positions), self._rotate(k, positions) |
| scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5) |
| if padding_mask is not None: |
| scores = scores.masked_fill( |
| padding_mask[:, None, None, :], -torch.finfo(scores.dtype).max |
| ) |
| attention = self.dropout(torch.softmax(scores, dim=-1)) |
| contextual = torch.matmul(attention, v).transpose(1, 2).reshape(batch, length, dimension) |
| output = values + self.output(contextual) |
| output = output + self.ffn(self.ffn_norm(output)) |
| if padding_mask is not None: |
| output = output.masked_fill(padding_mask.unsqueeze(-1), 0.0) |
| return output |
|
|
|
|
| class RoPETransformer(nn.Module): |
| """BGC-MAP-inspired bidirectional RoPE encoder with attentive pooling.""" |
|
|
| input_names = ("gene_embeddings", "relative_positions", "padding_mask") |
|
|
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| self.config = config |
| self.gene_projection = GeneProjection(config) |
| self.blocks = nn.ModuleList( |
| RotarySelfAttentionBlock(config) for _ in range(config.attention_blocks) |
| ) |
| self.pooling = PoolingByAttention(config) |
| self.output = nn.Linear(config.hidden_dimension, config.output_dimension) |
|
|
| def encode_genes(self, gene_embeddings, relative_positions, padding_mask=None): |
| values = self.gene_projection(gene_embeddings, relative_positions) |
| for block in self.blocks: |
| values = block(values, relative_positions, padding_mask) |
| return values |
|
|
| def forward(self, gene_embeddings, relative_positions, padding_mask=None): |
| genes = self.encode_genes(gene_embeddings, relative_positions, padding_mask) |
| return F.normalize(self.output(self.pooling(genes, padding_mask)), p=2, dim=-1) |
|
|
|
|
| class CrossAttentionPool(nn.Module): |
| """Learned context queries cross-attend to a self-attended BGC sequence.""" |
|
|
| input_names = ("gene_embeddings", "relative_positions", "padding_mask") |
|
|
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| self.config = config |
| d = config.hidden_dimension |
| self.gene_projection = GeneProjection(config) |
| layer = nn.TransformerEncoderLayer( |
| d, config.attention_heads, config.feedforward_dimension, config.dropout, |
| batch_first=True, norm_first=True, activation="gelu", |
| ) |
| self.encoder = nn.TransformerEncoder(layer, config.attention_blocks) |
| self.queries = nn.Parameter(torch.empty(2, d)) |
| nn.init.xavier_uniform_(self.queries) |
| self.cross = nn.MultiheadAttention(d, config.attention_heads, config.dropout, batch_first=True) |
| self.output = nn.Linear(d, config.output_dimension) |
|
|
| def encode_genes(self, gene_embeddings, relative_positions, padding_mask=None): |
| values = self.gene_projection(gene_embeddings, relative_positions) |
| return self.encoder(values, src_key_padding_mask=padding_mask) |
|
|
| def forward(self, gene_embeddings, relative_positions, padding_mask=None): |
| genes = self.encode_genes(gene_embeddings, relative_positions, padding_mask) |
| queries = self.queries.unsqueeze(0).expand(genes.shape[0], -1, -1) |
| attended, _ = self.cross(queries, genes, genes, key_padding_mask=padding_mask) |
| return F.normalize(self.output(attended.mean(dim=1)), p=2, dim=-1) |
|
|
|
|
| class LocalGlobalEncoder(nn.Module): |
| """PST-inspired local adjacency message passing followed by global pooling.""" |
|
|
| input_names = ("gene_embeddings", "relative_positions", "padding_mask") |
|
|
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| self.config = config |
| d = config.hidden_dimension |
| self.gene_projection = GeneProjection(config) |
| self.local = nn.Conv1d(d, d, kernel_size=3, padding=1, groups=1) |
| self.local_gate = nn.Sequential(nn.Linear(d, d), nn.Sigmoid()) |
| self.global_block = InducedSelfAttentionBlock(config) |
| self.pooling = PoolingByAttention(config) |
| self.output = nn.Linear(d, config.output_dimension) |
|
|
| def encode_genes(self, gene_embeddings, relative_positions, padding_mask=None): |
| values = self.gene_projection(gene_embeddings, relative_positions) |
| local = self.local(values.transpose(1, 2)).transpose(1, 2) |
| values = values + local * self.local_gate(values) |
| return self.global_block(values, padding_mask) |
|
|
| def forward(self, gene_embeddings, relative_positions, padding_mask=None): |
| genes = self.encode_genes(gene_embeddings, relative_positions, padding_mask) |
| return F.normalize(self.output(self.pooling(genes, padding_mask)), p=2, dim=-1) |
|
|
|
|
| class DilatedCNN(nn.Module): |
| """BiGCARP-inspired residual dilated convolutional encoder.""" |
|
|
| input_names = ("gene_embeddings", "relative_positions", "padding_mask") |
|
|
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| self.config = config |
| d = config.hidden_dimension |
| self.gene_projection = GeneProjection(config) |
| self.blocks = nn.ModuleList( |
| nn.Sequential( |
| nn.Conv1d(d, d, 3, padding=dilation, dilation=dilation), |
| nn.GELU(), nn.Dropout(config.dropout), nn.Conv1d(d, d, 1), |
| ) for dilation in (1, 2, 4, 8, 16) |
| ) |
| self.pooling = PoolingByAttention(config) |
| self.output = nn.Linear(d, config.output_dimension) |
|
|
| def encode_genes(self, gene_embeddings, relative_positions, padding_mask=None): |
| values = self.gene_projection(gene_embeddings, relative_positions) |
| for block in self.blocks: |
| updated = block(values.transpose(1, 2)).transpose(1, 2) |
| values = values + updated |
| if padding_mask is not None: |
| values = values.masked_fill(padding_mask.unsqueeze(-1), 0.0) |
| return values |
|
|
| def forward(self, gene_embeddings, relative_positions, padding_mask=None): |
| genes = self.encode_genes(gene_embeddings, relative_positions, padding_mask) |
| return F.normalize(self.output(self.pooling(genes, padding_mask)), p=2, dim=-1) |
|
|
|
|
| class HierarchicalLocalGlobal(nn.Module): |
| """Multi-scale local convolutions followed by a small global Transformer.""" |
|
|
| input_names = ("gene_embeddings", "relative_positions", "padding_mask") |
|
|
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| self.config = config |
| d = config.hidden_dimension |
| self.gene_projection = GeneProjection(config) |
| self.local3 = nn.Conv1d(d, d // 2, 3, padding=1) |
| self.local5 = nn.Conv1d(d, d // 2, 5, padding=2) |
| self.merge = nn.Linear(d, d) |
| layer = nn.TransformerEncoderLayer( |
| d, config.attention_heads, config.feedforward_dimension, config.dropout, |
| batch_first=True, norm_first=True, activation="gelu", |
| ) |
| self.global_encoder = nn.TransformerEncoder( |
| layer, max(1, config.attention_blocks // 2) |
| ) |
| self.pooling = PoolingByAttention(config) |
| self.output = nn.Linear(d, config.output_dimension) |
|
|
| def encode_genes(self, gene_embeddings, relative_positions, padding_mask=None): |
| values = self.gene_projection(gene_embeddings, relative_positions) |
| transposed = values.transpose(1, 2) |
| local = torch.cat((self.local3(transposed), self.local5(transposed)), dim=1).transpose(1, 2) |
| values = values + self.merge(local) |
| return self.global_encoder(values, src_key_padding_mask=padding_mask) |
|
|
| def forward(self, gene_embeddings, relative_positions, padding_mask=None): |
| genes = self.encode_genes(gene_embeddings, relative_positions, padding_mask) |
| return F.normalize(self.output(self.pooling(genes, padding_mask)), p=2, dim=-1) |
|
|
|
|
| class PfamAugmentedSetNet(nn.Module): |
| """SetNet whose gene projection is conditioned on the BGC Pfam inventory.""" |
|
|
| input_names = ("gene_embeddings", "relative_positions", "padding_mask", "pfam_tokens") |
| uses_pfam = True |
|
|
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| self.config = config |
| self.gene_projection = PfamGeneProjection(config) |
| self.blocks = nn.ModuleList( |
| InducedSelfAttentionBlock(config) for _ in range(config.attention_blocks) |
| ) |
| self.pooling = PoolingByAttention(config) |
| self.output = nn.Linear(config.hidden_dimension, config.output_dimension) |
|
|
| def encode_genes( |
| self, |
| gene_embeddings: torch.Tensor, |
| relative_positions: torch.Tensor, |
| padding_mask: torch.Tensor | None = None, |
| pfam_tokens: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| if pfam_tokens is None: |
| raise ValueError("Pfam-augmented SetNet requires pfam_tokens") |
| values = self.gene_projection(gene_embeddings, relative_positions, pfam_tokens) |
| for block in self.blocks: |
| values = block(values, padding_mask) |
| return values |
|
|
| def forward( |
| self, |
| gene_embeddings: torch.Tensor, |
| relative_positions: torch.Tensor, |
| padding_mask: torch.Tensor | None = None, |
| pfam_tokens: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| genes = self.encode_genes( |
| gene_embeddings, relative_positions, padding_mask, pfam_tokens |
| ) |
| pooled = self.pooling(genes, padding_mask) |
| return F.normalize(self.output(pooled), p=2, dim=-1) |
|
|
|
|
| class WeightedPfamJaccard(nn.Module): |
| """Learn nonnegative Pfam importance weights for differentiable set Jaccard.""" |
|
|
| input_names = ("pfam_tokens",) |
| uses_pfam = True |
|
|
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| if config.pfam_vocab_size < 3: |
| raise ValueError("Pfam vocabulary must contain padding, unknown, and one domain token") |
| initial = float(torch.log(torch.expm1(torch.tensor(1.0)))) |
| self.raw_weights = nn.Parameter( |
| torch.full((config.pfam_vocab_size,), initial, dtype=torch.float32) |
| ) |
| with torch.no_grad(): |
| self.raw_weights[0] = -20.0 |
|
|
| def domain_weights(self) -> torch.Tensor: |
| positive = F.softplus(self.raw_weights) |
| return torch.cat((positive[:1] * 0.0, positive[1:])) |
|
|
| def pairwise_jaccard(self, pfam_tokens: torch.Tensor) -> torch.Tensor: |
| if pfam_tokens.ndim != 2: |
| raise ValueError("Pfam tokens must have shape [batch, domains]") |
| batch_size = pfam_tokens.shape[0] |
| vocabulary = self.raw_weights.shape[0] |
| presence = torch.zeros( |
| batch_size, vocabulary, device=pfam_tokens.device, dtype=torch.float32 |
| ) |
| presence.scatter_(1, pfam_tokens.clamp_min(0), 1.0) |
| presence[:, 0] = 0.0 |
| weighted = presence * self.domain_weights().to(pfam_tokens.device) |
| totals = weighted.sum(dim=1) |
| intersection = weighted @ presence.T |
| union = totals[:, None] + totals[None, :] - intersection |
| return intersection / union.clamp_min(1e-8) |
|
|
| def forward(self, pfam_tokens: torch.Tensor) -> torch.Tensor: |
| return self.domain_weights() |
|
|
|
|
| ARCHITECTURES = { |
| "setnet": LeakageFreeBGCSetNet, |
| "weighted_pfam_jaccard": WeightedPfamJaccard, |
| "pfam_setnet": PfamAugmentedSetNet, |
| "gated_deepsets": GatedDeepSets, |
| "rope_transformer": RoPETransformer, |
| "cross_attention": CrossAttentionPool, |
| "local_global": LocalGlobalEncoder, |
| "dilated_cnn": DilatedCNN, |
| "hierarchical": HierarchicalLocalGlobal, |
| } |
|
|
|
|
| def build_model(config: ModelConfig) -> nn.Module: |
| try: |
| model_class = ARCHITECTURES[config.architecture] |
| except KeyError as exc: |
| raise ValueError(f"Unknown architecture: {config.architecture}") from exc |
| return model_class(config) |
|
|
