Central Dogma Transformer

生物中心法则（Central Dogma of Molecular Biology）是分子生物学的核心理论，由弗朗西斯·克里克于1958年提出。它描述了遗传信息在生物大分子之间传递的基本方向与规律：遗传信息从DNA流向RNA（转录），再从RNA流向蛋白质（翻译）。这一过程是单向且不可逆的，即信息可以从核酸传递到核酸，或从核酸传递到蛋白质，但不能从蛋白质逆向传递回核酸。中心法则奠定了现代遗传学、基因工程和生物信息学的基础，是理解生命活动如生长、发育、遗传和疾病发生机制的关键框架。

架构示意图

模型参数定义

@dataclass
class CDT2StageVCEConfig:
    """CDT-III 2-Stage VCE Model Configuration"""
    # Input dimensions
    dna_dim: int = 3072
    dna_seq_len: int = 896
    n_genes: int = 2361      # RNA genes (2360 + GFI1B)
    n_proteins: int = 189    # ADT proteins (193 - 4 isotype controls)

    # Model dimensions
    hidden_dim: int = 512    # Same as CDT-II
    nhead: int = 8
    dropout: float = 0.3        # CDT-II components (preserved)
    protein_dropout: float = 0.5  # v2: stronger for protein path

    # Self-Attention layers
    dna_self_attn_layers: int = 2
    rna_self_attn_layers: int = 1
    protein_self_attn_layers: int = 1

模型组件

class RawExpressionEncoder(nn.Module):
    """Raw expression -> hidden_dim embeddings.
    Used for both RNA (2361 genes) and Protein (189 proteins).
    """
    def __init__(self, n_features: int, hidden_dim: int, dropout: float = 0.1):
        super().__init__()
        self.n_features = n_features
        self.hidden_dim = hidden_dim

        self.feature_embedding = nn.Embedding(n_features, hidden_dim)
        self.expr_projector = nn.Sequential(
            nn.Linear(1, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout)
        )
        self.combine = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.Dropout(dropout)
        )

    def forward(self, expression: torch.Tensor) -> torch.Tensor:
        batch_size = expression.size(0)
        device = expression.device

        feat_ids = torch.arange(self.n_features, device=device)
        feat_emb = self.feature_embedding(feat_ids)
        feat_emb = feat_emb.unsqueeze(0).expand(batch_size, -1, -1)

        expr_emb = self.expr_projector(expression.unsqueeze(-1))

        combined = torch.cat([feat_emb, expr_emb], dim=-1)
        return self.combine(combined)


class SequenceProjector(nn.Module):
    def __init__(self, input_dim: int, output_dim: int, dropout: float = 0.1):
        super().__init__()
        self.linear = nn.Linear(input_dim, output_dim)
        self.norm = nn.LayerNorm(output_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.dropout(self.norm(self.linear(x)))


class FlashSelfAttentionBlock(nn.Module):
    def __init__(self, d_model: int, nhead: int = 8, dropout: float = 0.1):
        super().__init__()
        self.d_model = d_model
        self.nhead = nhead
        self.head_dim = d_model // nhead
        self.dropout_p = dropout

        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)

        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_model * 4, d_model),
            nn.Dropout(dropout)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, return_attn: bool = False):
        batch_size, seq_len, _ = x.shape

        Q = self.q_proj(x).view(batch_size, seq_len, self.nhead, self.head_dim).transpose(1, 2)
        K = self.k_proj(x).view(batch_size, seq_len, self.nhead, self.head_dim).transpose(1, 2)
        V = self.v_proj(x).view(batch_size, seq_len, self.nhead, self.head_dim).transpose(1, 2)

        if return_attn:
            scale = self.head_dim ** -0.5
            attn_weights = torch.matmul(Q, K.transpose(-2, -1)) * scale
            attn_weights = F.softmax(attn_weights, dim=-1)
            attn_out = torch.matmul(attn_weights, V)
        else:
            attn_out = F.scaled_dot_product_attention(
                Q, K, V,
                dropout_p=self.dropout_p if self.training else 0.0
            )
            attn_weights = None

        attn_out = attn_out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        attn_out = self.out_proj(attn_out)

        x = self.norm1(x + self.dropout(attn_out))
        x = self.norm2(x + self.ffn(x))

        if return_attn:
            return x, attn_weights
        return x


class FlashCrossAttentionBlock(nn.Module):
    def __init__(self, d_model: int, nhead: int = 8, dropout: float = 0.1):
        super().__init__()
        self.d_model = d_model
        self.nhead = nhead
        self.head_dim = d_model // nhead
        self.dropout_p = dropout

        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)

        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_model * 4, d_model),
            nn.Dropout(dropout)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, query: torch.Tensor, key_value: torch.Tensor,
                return_attn: bool = False):
        batch_size, query_len, _ = query.shape
        key_len = key_value.shape[1]

        Q = self.q_proj(query).view(batch_size, query_len, self.nhead, self.head_dim).transpose(1, 2)
        K = self.k_proj(key_value).view(batch_size, key_len, self.nhead, self.head_dim).transpose(1, 2)
        V = self.v_proj(key_value).view(batch_size, key_len, self.nhead, self.head_dim).transpose(1, 2)

        if return_attn:
            scale = self.head_dim ** -0.5
            attn_weights = torch.matmul(Q, K.transpose(-2, -1)) * scale
            attn_weights = F.softmax(attn_weights, dim=-1)
            attn_out = torch.matmul(attn_weights, V)
        else:
            attn_out = F.scaled_dot_product_attention(
                Q, K, V,
                dropout_p=self.dropout_p if self.training else 0.0
            )
            attn_weights = None

        attn_out = attn_out.transpose(1, 2).contiguous().view(batch_size, query_len, self.d_model)
        attn_out = self.out_proj(attn_out)

        x = self.norm1(query + self.dropout(attn_out))
        x = self.norm2(x + self.ffn(x))

        if return_attn:
            return x, attn_weights
        return x

DNA-RNA

class VirtualCellEmbedderDNARNA(nn.Module):
    """VCE-T: DNA + RNA only (2 modalities) — CDT-II architecture"""

    def __init__(self, d_model: int, dropout: float = 0.1):
        super().__init__()
        self.d_model = d_model
        self.nhead = 4
        self.head_dim = d_model // self.nhead

        self.dna_query = nn.Parameter(torch.randn(1, 1, d_model))
        self.rna_query = nn.Parameter(torch.randn(1, 1, d_model))

        self.dna_q_proj = nn.Linear(d_model, d_model)
        self.dna_k_proj = nn.Linear(d_model, d_model)
        self.dna_v_proj = nn.Linear(d_model, d_model)
        self.dna_out_proj = nn.Linear(d_model, d_model)

        self.rna_q_proj = nn.Linear(d_model, d_model)
        self.rna_k_proj = nn.Linear(d_model, d_model)
        self.rna_v_proj = nn.Linear(d_model, d_model)
        self.rna_out_proj = nn.Linear(d_model, d_model)

        self.fusion = nn.Sequential(
            nn.Linear(d_model * 2, d_model * 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_model * 2, d_model),
            nn.LayerNorm(d_model)
        )

    def _attention_pool(self, query, key_value, q_proj, k_proj, v_proj, out_proj):
        batch_size = key_value.size(0)
        seq_len = key_value.size(1)
        query = query.expand(batch_size, -1, -1)

        Q = q_proj(query).view(batch_size, 1, self.nhead, self.head_dim).transpose(1, 2)
        K = k_proj(key_value).view(batch_size, seq_len, self.nhead, self.head_dim).transpose(1, 2)
        V = v_proj(key_value).view(batch_size, seq_len, self.nhead, self.head_dim).transpose(1, 2)

        attn_out = F.scaled_dot_product_attention(Q, K, V)
        attn_out = attn_out.transpose(1, 2).contiguous().view(batch_size, 1, self.d_model)
        return out_proj(attn_out).squeeze(1)

    def forward(self, dna_encoded, rna_encoded):
        dna_pooled = self._attention_pool(
            self.dna_query, dna_encoded,
            self.dna_q_proj, self.dna_k_proj, self.dna_v_proj, self.dna_out_proj
        )
        rna_pooled = self._attention_pool(
            self.rna_query, rna_encoded,
            self.rna_q_proj, self.rna_k_proj, self.rna_v_proj, self.rna_out_proj
        )

        concat = torch.cat([dna_pooled, rna_pooled], dim=-1)  # [B, d*2]
        return self.fusion(concat)  # [B, d]

RNA-Protein

class VirtualCellEmbedderProtein(nn.Module):
    """VCE-P: Protein modality embedder (Stage 2 of 2-Stage VCE)

    Input:
        cell_emb_rna: [B, d_model] from VCE-T
        protein_encoded: [B, n_proteins, d_model] after self-attn + cross-attn

    Output:
        cell_emb_protein: [B, d_model]
    """

    def __init__(self, d_model: int, dropout: float = 0.1):
        super().__init__()
        self.d_model = d_model
        self.nhead = 4
        self.head_dim = d_model // self.nhead

        # Protein attention pooling
        self.protein_query = nn.Parameter(torch.randn(1, 1, d_model))
        self.protein_q_proj = nn.Linear(d_model, d_model)
        self.protein_k_proj = nn.Linear(d_model, d_model)
        self.protein_v_proj = nn.Linear(d_model, d_model)
        self.protein_out_proj = nn.Linear(d_model, d_model)

        # Fusion: cell_emb_rna [512] + protein_pooled [512] → [512]
        # Same structure as VCE-T fusion (d*2 → d*2 → d)
        self.fusion = nn.Sequential(
            nn.Linear(d_model * 2, d_model * 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_model * 2, d_model),
            nn.LayerNorm(d_model)
        )

    def _attention_pool(self, query, key_value, q_proj, k_proj, v_proj, out_proj):
        batch_size = key_value.size(0)
        seq_len = key_value.size(1)
        query = query.expand(batch_size, -1, -1)

        Q = q_proj(query).view(batch_size, 1, self.nhead, self.head_dim).transpose(1, 2)
        K = k_proj(key_value).view(batch_size, seq_len, self.nhead, self.head_dim).transpose(1, 2)
        V = v_proj(key_value).view(batch_size, seq_len, self.nhead, self.head_dim).transpose(1, 2)

        attn_out = F.scaled_dot_product_attention(Q, K, V)
        attn_out = attn_out.transpose(1, 2).contiguous().view(batch_size, 1, self.d_model)
        return out_proj(attn_out).squeeze(1)

    def forward(self, cell_emb_rna, protein_encoded):
        # Attention-pool protein sequence → single vector
        protein_pooled = self._attention_pool(
            self.protein_query, protein_encoded,
            self.protein_q_proj, self.protein_k_proj,
            self.protein_v_proj, self.protein_out_proj
        )  # [B, d_model]

        # Fuse transcriptome context + proteome summary
        concat = torch.cat([cell_emb_rna, protein_pooled], dim=-1)  # [B, d*2]
        return self.fusion(concat)  # [B, d]

两阶段模型

class CDTTrimodal2StageModel(nn.Module):
    """
    CDT-III: 2-Stage VCE Trimodal Central Dogma Transformer

    Key difference from 1-stage VCE:
    - VCE-T (DNA+RNA) is IDENTICAL to CDT-II → fusion weights transfer 100%
    - VCE-P (RNA→Protein) is NEW → stacked on top of VCE-T output
    - RNA prediction uses cell_emb_rna (from VCE-T)
    - Protein prediction uses cell_emb_protein (from VCE-P)

    Architecture:
        DNA [896, 3072] → Projector → Self-Attn(×2)
        RNA [2361] → RawEncoder → Self-Attn(×1)
        DNA→RNA Cross-Attn

        VCE-T(DNA, RNA) → cell_emb_rna [512]     ← CDT-II (all weights transferred)
            ├── RNA TaskHead → [2361] log2FC

        Protein [189] → RawEncoder → Self-Attn(×1)
        RNA→Protein Cross-Attn

        VCE-P(cell_emb_rna, Protein) → cell_emb_protein [512]  ← NEW
            └── Protein TaskHead → [189] DSB effect
    """

    def __init__(self, config: Optional[CDT2StageVCEConfig] = None):
        super().__init__()
        if config is None:
            config = CDT2StageVCEConfig()
        self.config = config

        # === CDT-II components (ALL transferable) ===
        # DNA
        self.dna_projector = SequenceProjector(
            config.dna_dim, config.hidden_dim, config.dropout)
        self.dna_self_attn_layers = nn.ModuleList([
            FlashSelfAttentionBlock(config.hidden_dim, config.nhead, config.dropout)
            for _ in range(config.dna_self_attn_layers)
        ])
        # RNA
        self.rna_encoder = RawExpressionEncoder(
            config.n_genes, config.hidden_dim, config.dropout)
        self.rna_self_attn_layers = nn.ModuleList([
            FlashSelfAttentionBlock(config.hidden_dim, config.nhead, config.dropout)
            for _ in range(config.rna_self_attn_layers)
        ])
        # DNA→RNA Cross-Attention
        self.dna_to_rna = FlashCrossAttentionBlock(
            config.hidden_dim, config.nhead, config.dropout)
        # VCE-T (CDT-II VCE — fusion d*2→d, fully transferable)
        self.vce_t = VirtualCellEmbedderDNARNA(
            config.hidden_dim, config.dropout)
        # RNA Task Head
        self.rna_task_layer = nn.Sequential(
            nn.Linear(config.hidden_dim, config.hidden_dim),
            nn.GELU(),
            nn.Dropout(config.dropout),
            nn.Linear(config.hidden_dim, config.n_genes)
        )

        # === CDT-III NEW components ===
        # Protein Encoder
        self.protein_encoder = RawExpressionEncoder(
            config.n_proteins, config.hidden_dim, config.protein_dropout)
        # Protein Self-Attention
        self.protein_self_attn_layers = nn.ModuleList([
            FlashSelfAttentionBlock(config.hidden_dim, config.nhead, config.protein_dropout)
            for _ in range(config.protein_self_attn_layers)
        ])
        # RNA→Protein Cross-Attention
        self.rna_to_protein = FlashCrossAttentionBlock(
            config.hidden_dim, config.nhead, config.protein_dropout)
        # VCE-P (NEW — takes cell_emb_rna + protein)
        self.vce_p = VirtualCellEmbedderProtein(
            config.hidden_dim, config.protein_dropout)
        # Protein Task Head
        self.protein_task_layer = nn.Sequential(
            nn.Linear(config.hidden_dim, config.hidden_dim),
            nn.GELU(),
            nn.Dropout(config.protein_dropout),
            nn.Linear(config.hidden_dim, config.n_proteins)
        )

    def forward(self, dna_emb, rna_expr, protein_expr,
                return_attention: bool = False):
        """
        Args:
            dna_emb: [batch, 896, 3072]
            rna_expr: [batch, n_genes]
            protein_expr: [batch, n_proteins]

        Returns:
            rna_pred: [batch, n_genes]
            protein_pred: [batch, n_proteins]
        """
        attn_maps = {} if return_attention else None

        # === Encode ===
        dna = self.dna_projector(dna_emb)            # [B, 896, 512]
        rna = self.rna_encoder(rna_expr)             # [B, 2361, 512]
        protein = self.protein_encoder(protein_expr)  # [B, 189, 512]

        # === DNA Self-Attention ===
        for i, layer in enumerate(self.dna_self_attn_layers):
            if return_attention:
                dna, attn_w = layer(dna, return_attn=True)
                attn_maps[f'dna_self_attn_{i}'] = attn_w.detach().cpu()
            else:
                dna = layer(dna)

        # === RNA Self-Attention ===
        for i, layer in enumerate(self.rna_self_attn_layers):
            if return_attention:
                rna, attn_w = layer(rna, return_attn=True)
                attn_maps[f'rna_self_attn_{i}'] = attn_w.detach().cpu()
            else:
                rna = layer(rna)

        # === Protein Self-Attention ===
        for i, layer in enumerate(self.protein_self_attn_layers):
            if return_attention:
                protein, attn_w = layer(protein, return_attn=True)
                attn_maps[f'protein_self_attn_{i}'] = attn_w.detach().cpu()
            else:
                protein = layer(protein)

        # === Cross-Attention: DNA → RNA ===
        if return_attention:
            rna, attn_w = self.dna_to_rna(query=rna, key_value=dna, return_attn=True)
            attn_maps['dna_to_rna_cross'] = attn_w.detach().cpu()
        else:
            rna = self.dna_to_rna(query=rna, key_value=dna)

        # === Cross-Attention: RNA → Protein ===
        if return_attention:
            protein, attn_w = self.rna_to_protein(
                query=protein, key_value=rna, return_attn=True)
            attn_maps['rna_to_protein_cross'] = attn_w.detach().cpu()
        else:
            protein = self.rna_to_protein(query=protein, key_value=rna)

        # === Stage 1: VCE-T (DNA + RNA) → cell_emb_rna ===
        cell_emb_rna = self.vce_t(dna, rna)  # [B, 512]

        # === RNA Prediction (from VCE-T output) ===
        rna_pred = self.rna_task_layer(cell_emb_rna)  # [B, 2361]

        # === Stage 2: VCE-P (cell_emb_rna + Protein) → cell_emb_protein ===
        cell_emb_protein = self.vce_p(cell_emb_rna, protein)  # [B, 512]

        # === Protein Prediction (from VCE-P output) ===
        protein_pred = self.protein_task_layer(cell_emb_protein)  # [B, 189]

        if return_attention:
            return rna_pred, protein_pred, attn_maps
        return rna_pred, protein_pred

    # --- Freeze/Unfreeze methods for 2-phase training ---

    def get_cdt2_param_names(self):
        """CDT-II component prefixes (all transferable)."""
        return [
            'dna_projector', 'dna_self_attn_layers',
            'rna_encoder', 'rna_self_attn_layers',
            'dna_to_rna', 'vce_t', 'rna_task_layer'
        ]

    def get_cdt3_param_names(self):
        """CDT-III NEW component prefixes."""
        return [
            'protein_encoder', 'protein_self_attn_layers',
            'rna_to_protein', 'vce_p', 'protein_task_layer'
        ]

    def freeze_cdt2(self):
        """Phase 1: Freeze ALL CDT-II components (including VCE-T + rna_task_layer)."""
        cdt2_prefixes = self.get_cdt2_param_names()
        frozen_count = 0
        for name, param in self.named_parameters():
            if any(name.startswith(prefix) for prefix in cdt2_prefixes):
                param.requires_grad = False
                frozen_count += 1
        print(f"Frozen {frozen_count} CDT-II parameters (incl. VCE-T fusion)")
        print(f"Trainable: {self.get_num_params(trainable_only=True):,} params")

    def unfreeze_all(self):
        """Phase 2: Unfreeze all parameters for joint fine-tuning."""
        for param in self.parameters():
            param.requires_grad = True
        print(f"Unfrozen all {sum(1 for p in self.parameters())} parameters")

    def get_num_params(self, trainable_only: bool = False):
        if trainable_only:
            return sum(p.numel() for p in self.parameters() if p.requires_grad)
        return sum(p.numel() for p in self.parameters())

完整流程

请参考nobusama/CDT3，Central Dogma Transformer III: Interpretable AI Across DNA, RNA, and Protein