| from __future__ import annotations |
|
|
| from dataclasses import dataclass, fields as dataclass_fields |
| from contextlib import ExitStack |
| from typing import Any, Dict, List, Optional, Tuple |
|
|
| import json |
| import math |
| import os |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch.utils.checkpoint import checkpoint |
|
|
| try: |
| from torch.distributed.fsdp import FullyShardedDataParallel as FSDP |
| except Exception: |
| FSDP = None |
|
|
|
|
| def _dist_is_initialized() -> bool: |
| return torch.distributed.is_available() and torch.distributed.is_initialized() |
|
|
|
|
| def _leaf_fsdp_modules(module: nn.Module) -> List[nn.Module]: |
| if FSDP is None: |
| return [] |
|
|
| fsdp_modules: List[nn.Module] = [] |
| for child in module.modules(): |
| if not isinstance(child, FSDP): |
| continue |
| has_fsdp_descendant = any(isinstance(grandchild, FSDP) for grandchild in child.children()) |
| if not has_fsdp_descendant: |
| fsdp_modules.append(child) |
| return fsdp_modules |
|
|
|
|
| def _json_dump(obj: Any, path: str) -> None: |
| os.makedirs(os.path.dirname(path), exist_ok=True) |
| with open(path, "w", encoding="utf-8") as f: |
| json.dump(obj, f, ensure_ascii=False, indent=2) |
|
|
|
|
| class RmsNorm(nn.Module): |
| def __init__(self, dim: int, eps: float = 1e-6): |
| super().__init__() |
| self.eps = float(eps) |
| self.scale = nn.Parameter(torch.ones(dim)) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| inv = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt() |
| return x * inv * self.scale |
|
|
|
|
| class SdpaSelfAttention(nn.Module): |
| def __init__(self, d_model: int, n_heads: int, dropout: float): |
| super().__init__() |
| if d_model % n_heads != 0: |
| raise ValueError("d_model must be divisible by n_heads") |
| self.d_model = int(d_model) |
| self.n_heads = int(n_heads) |
| self.head_dim = self.d_model // self.n_heads |
| self.dropout = float(dropout) |
|
|
| self.q_proj = nn.Linear(self.d_model, self.d_model, bias=False) |
| self.k_proj = nn.Linear(self.d_model, self.d_model, bias=False) |
| self.v_proj = nn.Linear(self.d_model, self.d_model, bias=False) |
| self.o_proj = nn.Linear(self.d_model, self.d_model, bias=False) |
|
|
| def forward(self, x: torch.Tensor, *, attn_mask_valid: torch.Tensor) -> torch.Tensor: |
| batch_size, seq_len, _ = x.shape |
| q = self.q_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2) |
| k = self.k_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2) |
| v = self.v_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2) |
|
|
| if attn_mask_valid.shape != (batch_size, seq_len): |
| raise ValueError( |
| "attn_mask_valid must have shape [batch_size, seq_len], " |
| f"got {tuple(attn_mask_valid.shape)} for x={tuple(x.shape)}" |
| ) |
| |
| |
| key_mask = attn_mask_valid.to(device=x.device, dtype=torch.bool)[:, None, None, :] |
| drop_p = self.dropout if self.training else 0.0 |
| out = F.scaled_dot_product_attention( |
| q, |
| k, |
| v, |
| attn_mask=key_mask, |
| dropout_p=drop_p, |
| is_causal=False, |
| ) |
| out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) |
| return self.o_proj(out) |
|
|
|
|
| class EncoderBlock(nn.Module): |
| def __init__(self, d_model: int, n_heads: int, dropout: float): |
| super().__init__() |
| self.norm1 = RmsNorm(d_model) |
| self.attn = SdpaSelfAttention(d_model, n_heads, dropout) |
| self.drop = nn.Dropout(dropout) |
| self.norm2 = RmsNorm(d_model) |
| self.ffn = nn.Sequential( |
| nn.Linear(d_model, 4 * d_model, bias=False), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(4 * d_model, d_model, bias=False), |
| ) |
|
|
| def forward(self, x: torch.Tensor, *, attn_mask_valid: torch.Tensor) -> torch.Tensor: |
| x = x * attn_mask_valid.unsqueeze(-1).to(dtype=x.dtype) |
| h = self.norm1(x) |
| x = x + self.drop(self.attn(h, attn_mask_valid=attn_mask_valid)) |
| x = x + self.drop(self.ffn(self.norm2(x))) |
| return x * attn_mask_valid.unsqueeze(-1).to(dtype=x.dtype) |
|
|
|
|
| @dataclass |
| class FoundationConfig: |
| d_model: int = 768 |
| n_heads: int = 12 |
| dropout: float = 0.0 |
| encoder_layers: int = 12 |
| gradient_checkpointing: bool = False |
|
|
| gene_vocab_size: int = 0 |
| bin_vocab_size: int = 0 |
| gene_pad_id: int = 0 |
| bin_pad_id: int = 0 |
| bin_mask_id: int = 1 |
| bin_offset: int = 3 |
| gene_cls_id: int = 1 |
| expr_embedding_mode: str = "bin" |
| expr_rank_mlp_hidden_dim: int = 256 |
|
|
| dino_out_dim: int = 4096 |
| dino_hidden_dim: int = 2048 |
| dino_bottleneck_dim: int = 256 |
| student_temperature: float = 0.10 |
| teacher_temperature: float = 0.04 |
| teacher_momentum: float = 0.996 |
| teacher_momentum_end: float = 1.0 |
| center_momentum: float = 0.9 |
| token_center_momentum: float = 0.9 |
| cross_modal_dino_weight: float = 0.25 |
| cross_modal_pairing_mode: str = "disabled" |
| token_loss_weight: float = 1.0 |
| token_unmasked_loss_weight: float = 0.0 |
| inference_use_teacher: bool = True |
| teacher_cls_decoder_loss_weight: float = 0.0 |
| teacher_cls_decoder_hidden_dim: int = 1024 |
| teacher_cls_decoder_condition_on_obs: bool = False |
| teacher_cls_decoder_rank_loss_alpha: float = 0.05 |
| teacher_cls_decoder_rank_num_pairs: int = 4096 |
| teacher_student_decoder_distill_weight: float = 0.0 |
| teacher_student_decoder_loss_mode: str = "distribution" |
| teacher_student_decoder_rank_alpha: float = 0.05 |
| teacher_student_decoder_raw_mmd_weight: float = 0.1 |
| teacher_student_decoder_raw_rank_weight: float = 0.01 |
| teacher_student_decoder_raw_mse_weight: float = 1.0 |
|
|
|
|
| class GeneBinTokenizer(nn.Module): |
| def __init__(self, cfg: FoundationConfig): |
| super().__init__() |
| self.gene_emb = nn.Embedding(cfg.gene_vocab_size, cfg.d_model, padding_idx=cfg.gene_pad_id) |
| self.bin_emb = nn.Embedding(cfg.bin_vocab_size, cfg.d_model, padding_idx=cfg.bin_pad_id) |
| self.expr_embedding_mode = str(getattr(cfg, "expr_embedding_mode", "bin")).strip().lower() |
| if self.expr_embedding_mode not in {"bin", "rank"}: |
| raise ValueError(f"Unsupported expr_embedding_mode={self.expr_embedding_mode!r}") |
| self.bin_mask_id = int(getattr(cfg, "bin_mask_id", 1)) |
| self.bin_offset = int(getattr(cfg, "bin_offset", 3)) |
| rank_hidden_dim = int(getattr(cfg, "expr_rank_mlp_hidden_dim", 256)) |
| if rank_hidden_dim <= 0: |
| raise ValueError("expr_rank_mlp_hidden_dim must be > 0") |
| self.rank_mlp: Optional[nn.Sequential] |
| self.rank_mask_emb: Optional[nn.Parameter] |
| if self.expr_embedding_mode == "rank": |
| self.rank_mlp = nn.Sequential( |
| nn.Linear(1, rank_hidden_dim, bias=True), |
| nn.SiLU(), |
| nn.Linear(rank_hidden_dim, cfg.d_model, bias=False), |
| ) |
| self.rank_mask_emb = nn.Parameter(torch.empty(cfg.d_model)) |
| nn.init.normal_(self.rank_mask_emb, mean=0.0, std=0.02) |
| else: |
| self.rank_mlp = None |
| self.rank_mask_emb = None |
|
|
| def _rank_expression_emb(self, expr_ids: torch.Tensor) -> torch.Tensor: |
| if self.rank_mlp is None: |
| raise RuntimeError("rank_mlp is not initialized for non-rank expr_embedding_mode") |
| out = self.gene_emb.weight.new_zeros((*expr_ids.shape, self.gene_emb.embedding_dim)) |
| positive_mask = expr_ids >= int(self.bin_offset) |
| if bool(positive_mask.any().item()): |
| rank_values = (expr_ids[positive_mask] - int(self.bin_offset) + 1).to(dtype=out.dtype).unsqueeze(-1) |
| rank_emb = self.rank_mlp(rank_values) |
| if rank_emb.dtype != out.dtype: |
| rank_emb = rank_emb.to(dtype=out.dtype) |
| out[positive_mask] = rank_emb |
| mask_mask = expr_ids == int(self.bin_mask_id) |
| if bool(mask_mask.any().item()): |
| if self.rank_mask_emb is None: |
| raise RuntimeError("rank_mask_emb is not initialized for rank expr_embedding_mode") |
| out[mask_mask] = self.rank_mask_emb.to(dtype=out.dtype) |
| return out |
|
|
| def forward( |
| self, |
| gene_ids: torch.Tensor, |
| bin_ids: torch.Tensor, |
| attn_mask_valid: torch.Tensor, |
| *, |
| use_bin_embeddings: bool, |
| ) -> torch.Tensor: |
| x = self.gene_emb(gene_ids) |
| if bool(use_bin_embeddings): |
| if self.expr_embedding_mode == "rank": |
| x = x + self._rank_expression_emb(bin_ids) |
| else: |
| x = x + self.bin_emb(bin_ids) |
| return x * attn_mask_valid.unsqueeze(-1).to(dtype=x.dtype) |
|
|
|
|
| class EncoderBackbone(nn.Module): |
| def __init__(self, cfg: FoundationConfig): |
| super().__init__() |
| self.gradient_checkpointing = bool(getattr(cfg, "gradient_checkpointing", False)) |
| self.tokenizer = GeneBinTokenizer(cfg) |
| self.encoder = nn.ModuleList( |
| [EncoderBlock(cfg.d_model, cfg.n_heads, cfg.dropout) for _ in range(cfg.encoder_layers)] |
| ) |
| self.encoder_norm = RmsNorm(cfg.d_model) |
|
|
| def forward( |
| self, |
| gene_ids: torch.Tensor, |
| bin_ids: torch.Tensor, |
| attention_mask: torch.Tensor, |
| *, |
| use_bin_embeddings: bool, |
| ) -> torch.Tensor: |
| x = self.tokenizer( |
| gene_ids, |
| bin_ids, |
| attention_mask, |
| use_bin_embeddings=use_bin_embeddings, |
| ) |
| for blk in self.encoder: |
| if self.gradient_checkpointing and self.training and torch.is_grad_enabled(): |
| def _block_forward(hidden: torch.Tensor, *, block: EncoderBlock = blk) -> torch.Tensor: |
| return block(hidden, attn_mask_valid=attention_mask) |
|
|
| x = checkpoint(_block_forward, x, use_reentrant=False) |
| else: |
| x = blk(x, attn_mask_valid=attention_mask) |
| x = self.encoder_norm(x) |
| return x * attention_mask.unsqueeze(-1).to(dtype=x.dtype) |
|
|
|
|
| class DinoHead(nn.Module): |
| def __init__(self, cfg: FoundationConfig): |
| super().__init__() |
| hidden_dim = int(cfg.dino_hidden_dim) |
| bottleneck_dim = int(cfg.dino_bottleneck_dim) |
| out_dim = int(cfg.dino_out_dim) |
| if hidden_dim <= 0 or bottleneck_dim <= 0 or out_dim <= 0: |
| raise ValueError("DINO head dimensions must all be > 0") |
| self.mlp = nn.Sequential( |
| nn.Linear(cfg.d_model, hidden_dim, bias=False), |
| nn.GELU(), |
| nn.Linear(hidden_dim, hidden_dim, bias=False), |
| nn.GELU(), |
| nn.Linear(hidden_dim, bottleneck_dim, bias=False), |
| ) |
| self.norm = nn.LayerNorm(bottleneck_dim) |
| self.prototypes = nn.Linear(bottleneck_dim, out_dim, bias=False) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = self.mlp(x) |
| x = self.norm(x) |
| x = F.normalize(x, dim=-1) |
| return self.prototypes(x) |
|
|
|
|
| class TeacherCLSExpressionDecoder(nn.Module): |
| def __init__(self, cfg: FoundationConfig): |
| super().__init__() |
| hidden_dim = int(getattr(cfg, "teacher_cls_decoder_hidden_dim", 1024)) |
| if hidden_dim <= 0: |
| raise ValueError("teacher_cls_decoder_hidden_dim must be > 0") |
| self.condition_on_obs = bool(getattr(cfg, "teacher_cls_decoder_condition_on_obs", False)) |
| self.cls_proj = nn.Sequential( |
| nn.Linear(cfg.d_model, hidden_dim, bias=False), |
| nn.LayerNorm(hidden_dim), |
| nn.SiLU(), |
| ) |
| self.obs_proj: Optional[nn.Sequential] |
| if self.condition_on_obs: |
| self.obs_proj = nn.Sequential( |
| nn.Linear(2, hidden_dim, bias=False), |
| nn.LayerNorm(hidden_dim), |
| nn.SiLU(), |
| ) |
| else: |
| self.obs_proj = None |
| self.gene_proj = nn.Sequential( |
| nn.Linear(cfg.d_model, hidden_dim, bias=False), |
| nn.LayerNorm(hidden_dim), |
| nn.SiLU(), |
| ) |
| self.out_proj = nn.Sequential( |
| nn.Linear(2 * hidden_dim, hidden_dim, bias=False), |
| nn.SiLU(), |
| nn.Linear(hidden_dim, 1, bias=True), |
| ) |
|
|
| def forward( |
| self, |
| cls_state: torch.Tensor, |
| gene_query_states: torch.Tensor, |
| obs_features: Optional[torch.Tensor] = None, |
| ) -> torch.Tensor: |
| cls_hidden = self.cls_proj(cls_state) |
| if self.condition_on_obs: |
| if obs_features is None: |
| raise ValueError("teacher_cls_decoder_condition_on_obs=True requires obs_features") |
| cls_hidden = cls_hidden + self.obs_proj( |
| obs_features.to(device=cls_hidden.device, dtype=cls_hidden.dtype) |
| ) |
| cls_hidden = cls_hidden.unsqueeze(1).expand(-1, gene_query_states.shape[1], -1) |
| gene_hidden = self.gene_proj(gene_query_states) |
| return self.out_proj(torch.cat([cls_hidden, gene_hidden], dim=-1)).squeeze(-1) |
|
|
|
|
| class CellContrastivePredictModel(nn.Module): |
| def __init__(self, cfg: FoundationConfig): |
| super().__init__() |
| self.cfg = cfg |
| if cfg.encoder_layers < 1: |
| raise ValueError("encoder_layers must be >= 1") |
| if cfg.n_heads < 1: |
| raise ValueError("n_heads must be >= 1") |
| if cfg.d_model <= 0: |
| raise ValueError("d_model must be > 0") |
| if cfg.gene_vocab_size <= cfg.gene_cls_id: |
| raise ValueError("gene_vocab_size too small for gene_cls_id") |
| if cfg.bin_vocab_size <= 0: |
| raise ValueError("bin_vocab_size must be > 0") |
| if not (0.0 < float(cfg.student_temperature)): |
| raise ValueError("student_temperature must be > 0") |
| if not (0.0 < float(cfg.teacher_temperature)): |
| raise ValueError("teacher_temperature must be > 0") |
| if not (0.0 <= float(cfg.token_loss_weight)): |
| raise ValueError("token_loss_weight must be >= 0") |
| if float(getattr(cfg, "token_unmasked_loss_weight", 0.0)) < 0.0: |
| raise ValueError("token_unmasked_loss_weight must be >= 0") |
| if not (0.0 <= float(getattr(cfg, "teacher_cls_decoder_loss_weight", 0.0))): |
| raise ValueError("teacher_cls_decoder_loss_weight must be >= 0") |
| if int(getattr(cfg, "teacher_cls_decoder_hidden_dim", 1024)) <= 0: |
| raise ValueError("teacher_cls_decoder_hidden_dim must be > 0") |
| if not isinstance(getattr(cfg, "teacher_cls_decoder_condition_on_obs", False), bool): |
| raise ValueError("teacher_cls_decoder_condition_on_obs must be a bool") |
| if float(getattr(cfg, "teacher_cls_decoder_rank_loss_alpha", 0.05)) < 0.0: |
| raise ValueError("teacher_cls_decoder_rank_loss_alpha must be >= 0") |
| if int(getattr(cfg, "teacher_cls_decoder_rank_num_pairs", 4096)) < 0: |
| raise ValueError("teacher_cls_decoder_rank_num_pairs must be >= 0") |
| if float(getattr(cfg, "teacher_student_decoder_distill_weight", 0.0)) < 0.0: |
| raise ValueError("teacher_student_decoder_distill_weight must be >= 0") |
| loss_mode = str(getattr(cfg, "teacher_student_decoder_loss_mode", "distribution")).strip().lower() |
| if loss_mode not in {"distribution", "mse"}: |
| raise ValueError("teacher_student_decoder_loss_mode must be one of {'distribution', 'mse'}") |
| cfg.teacher_student_decoder_loss_mode = loss_mode |
| if float(getattr(cfg, "teacher_student_decoder_rank_alpha", 0.05)) < 0.0: |
| raise ValueError("teacher_student_decoder_rank_alpha must be >= 0") |
| if float(getattr(cfg, "teacher_student_decoder_raw_mmd_weight", 0.1)) < 0.0: |
| raise ValueError("teacher_student_decoder_raw_mmd_weight must be >= 0") |
| if float(getattr(cfg, "teacher_student_decoder_raw_rank_weight", 0.01)) < 0.0: |
| raise ValueError("teacher_student_decoder_raw_rank_weight must be >= 0") |
| if float(getattr(cfg, "teacher_student_decoder_raw_mse_weight", 1.0)) < 0.0: |
| raise ValueError("teacher_student_decoder_raw_mse_weight must be >= 0") |
| if str(getattr(cfg, "expr_embedding_mode", "bin")).strip().lower() not in {"bin", "rank"}: |
| raise ValueError("expr_embedding_mode must be one of {'bin', 'rank'}") |
| if int(getattr(cfg, "expr_rank_mlp_hidden_dim", 256)) <= 0: |
| raise ValueError("expr_rank_mlp_hidden_dim must be > 0") |
| cross_modal_pairing_mode = str(getattr(cfg, "cross_modal_pairing_mode", "disabled")).strip().lower() |
| if cross_modal_pairing_mode not in {"disabled", "symmetric", "bin_to_gene_only"}: |
| raise ValueError("cross_modal_pairing_mode must be one of {'disabled', 'symmetric', 'bin_to_gene_only'}") |
| cfg.cross_modal_pairing_mode = cross_modal_pairing_mode |
|
|
| self.student_backbone = EncoderBackbone(cfg) |
| self.teacher_backbone = EncoderBackbone(cfg) |
| self.student_head = DinoHead(cfg) |
| self.teacher_head = DinoHead(cfg) |
| self.student_ibot_head = DinoHead(cfg) |
| self.teacher_ibot_head = DinoHead(cfg) |
| self.student_cls_expression_decoder: Optional[TeacherCLSExpressionDecoder] = None |
| self.teacher_cls_expression_decoder: Optional[TeacherCLSExpressionDecoder] = None |
| self._configure_teacher_cls_expression_decoder( |
| loss_weight=float(getattr(cfg, "teacher_cls_decoder_loss_weight", 0.0)), |
| hidden_dim=int(getattr(cfg, "teacher_cls_decoder_hidden_dim", 1024)), |
| ) |
|
|
| self.register_buffer("center", torch.zeros(1, int(cfg.dino_out_dim))) |
| self.register_buffer("token_center", torch.zeros(1, int(cfg.dino_out_dim))) |
| self.register_buffer("last_teacher_momentum", torch.tensor(float(cfg.teacher_momentum), dtype=torch.float32)) |
|
|
| self._sync_teacher_from_student() |
| self._freeze_teacher() |
|
|
| def _configure_teacher_cls_expression_decoder( |
| self, |
| *, |
| loss_weight: float, |
| hidden_dim: int, |
| condition_on_obs: Optional[bool] = None, |
| ) -> None: |
| current_hidden = int(getattr(self.cfg, "teacher_cls_decoder_hidden_dim", hidden_dim)) |
| current_condition = bool(getattr(self.cfg, "teacher_cls_decoder_condition_on_obs", False)) |
| next_condition = current_condition if condition_on_obs is None else bool(condition_on_obs) |
| distill_enabled = float(getattr(self.cfg, "teacher_student_decoder_distill_weight", 0.0)) > 0.0 |
| self.cfg.teacher_cls_decoder_loss_weight = float(loss_weight) |
| self.cfg.teacher_cls_decoder_hidden_dim = int(hidden_dim) |
| self.cfg.teacher_cls_decoder_condition_on_obs = bool(next_condition) |
| if float(loss_weight) <= 0.0: |
| self.student_cls_expression_decoder = None |
| self.teacher_cls_expression_decoder = None |
| return |
| if distill_enabled: |
| if ( |
| self.student_cls_expression_decoder is None |
| or int(current_hidden) != int(hidden_dim) |
| or current_condition != bool(next_condition) |
| ): |
| self.student_cls_expression_decoder = TeacherCLSExpressionDecoder(self.cfg) |
| if ( |
| self.teacher_cls_expression_decoder is None |
| or int(current_hidden) != int(hidden_dim) |
| or current_condition != bool(next_condition) |
| ): |
| self.teacher_cls_expression_decoder = TeacherCLSExpressionDecoder(self.cfg) |
| self.teacher_cls_expression_decoder.load_state_dict(self.student_cls_expression_decoder.state_dict()) |
| for param in self.teacher_cls_expression_decoder.parameters(): |
| param.requires_grad = False |
| return |
| self.student_cls_expression_decoder = None |
| if ( |
| self.teacher_cls_expression_decoder is None |
| or int(current_hidden) != int(hidden_dim) |
| or current_condition != bool(next_condition) |
| ): |
| self.teacher_cls_expression_decoder = TeacherCLSExpressionDecoder(self.cfg) |
|
|
| def configure_teacher_cls_expression_decoder( |
| self, |
| *, |
| loss_weight: float, |
| hidden_dim: int, |
| condition_on_obs: Optional[bool] = None, |
| ) -> None: |
| self._configure_teacher_cls_expression_decoder( |
| loss_weight=float(loss_weight), |
| hidden_dim=int(hidden_dim), |
| condition_on_obs=condition_on_obs, |
| ) |
|
|
| def _freeze_teacher(self) -> None: |
| for param in self.teacher_backbone.parameters(): |
| param.requires_grad = False |
| for param in self.teacher_head.parameters(): |
| param.requires_grad = False |
| for param in self.teacher_ibot_head.parameters(): |
| param.requires_grad = False |
| if ( |
| float(getattr(self.cfg, "teacher_student_decoder_distill_weight", 0.0)) > 0.0 |
| and self.teacher_cls_expression_decoder is not None |
| ): |
| for param in self.teacher_cls_expression_decoder.parameters(): |
| param.requires_grad = False |
|
|
| def _sync_teacher_from_student(self) -> None: |
| self.teacher_backbone.load_state_dict(self.student_backbone.state_dict()) |
| self.teacher_head.load_state_dict(self.student_head.state_dict()) |
| self.teacher_ibot_head.load_state_dict(self.student_ibot_head.state_dict()) |
| if self.student_cls_expression_decoder is not None and self.teacher_cls_expression_decoder is not None: |
| self.teacher_cls_expression_decoder.load_state_dict(self.student_cls_expression_decoder.state_dict()) |
|
|
| def _pool_cell_repr(self, token_states: torch.Tensor) -> torch.Tensor: |
| return token_states[:, 0, :] |
|
|
| def _pool_instance_repr(self, token_states: torch.Tensor) -> torch.Tensor: |
| return self._pool_cell_repr(token_states) |
|
|
| def _select_backbone(self, backbone: Optional[str] = None) -> EncoderBackbone: |
| chosen = backbone |
| if chosen is None: |
| chosen = "teacher" if bool(getattr(self.cfg, "inference_use_teacher", True)) else "student" |
| if chosen == "teacher": |
| return self.teacher_backbone |
| if chosen == "student": |
| return self.student_backbone |
| raise ValueError(f"Unknown backbone={backbone!r}") |
|
|
| def _encode_tokens( |
| self, |
| gene_ids: torch.Tensor, |
| bin_ids: torch.Tensor, |
| attention_mask: torch.Tensor, |
| *, |
| use_bin_embeddings: bool = True, |
| backbone: Optional[str] = None, |
| ) -> torch.Tensor: |
| module = self._select_backbone(backbone) |
| return module( |
| gene_ids, |
| bin_ids, |
| attention_mask, |
| use_bin_embeddings=use_bin_embeddings, |
| ) |
|
|
| def _encode( |
| self, |
| gene_ids: torch.Tensor, |
| bin_ids: torch.Tensor, |
| attention_mask: torch.Tensor, |
| *, |
| use_bin_embeddings: bool = True, |
| backbone: Optional[str] = None, |
| ) -> torch.Tensor: |
| return self._pool_cell_repr( |
| self._encode_tokens( |
| gene_ids, |
| bin_ids, |
| attention_mask, |
| use_bin_embeddings=use_bin_embeddings, |
| backbone=backbone, |
| ) |
| ) |
|
|
| def _encode_view_stack( |
| self, |
| gene_ids: torch.Tensor, |
| bin_ids: torch.Tensor, |
| attention_mask: torch.Tensor, |
| view_use_bin: torch.Tensor, |
| *, |
| backbone_name: str, |
| head: nn.Module, |
| capture_token_views: Optional[Tuple[int, ...]] = None, |
| ) -> Tuple[torch.Tensor, torch.Tensor, Dict[int, torch.Tensor]]: |
| num_views = int(gene_ids.shape[1]) |
| pooled_views = [] |
| logits_views = [] |
| token_view_set = set(capture_token_views or ()) |
| captured_token_states: Dict[int, torch.Tensor] = {} |
| for view_idx in range(num_views): |
| use_bin = bool(view_use_bin[view_idx].item()) |
| view_attention_mask = attention_mask[:, view_idx, :] |
| valid_positions = torch.nonzero(view_attention_mask, as_tuple=False) |
| if valid_positions.numel() == 0: |
| raise ValueError(f"view {view_idx} contains no valid tokens") |
| max_valid_len = int(valid_positions[:, -1].max().item()) + 1 |
| tokens = self._encode_tokens( |
| gene_ids[:, view_idx, :max_valid_len], |
| bin_ids[:, view_idx, :max_valid_len], |
| view_attention_mask[:, :max_valid_len], |
| use_bin_embeddings=use_bin, |
| backbone=backbone_name, |
| ) |
| pooled = self._pool_cell_repr(tokens) |
| pooled_views.append(pooled) |
| logits_views.append(head(pooled)) |
| if view_idx in token_view_set: |
| captured_token_states[view_idx] = tokens |
| return torch.stack(pooled_views, dim=1), torch.stack(logits_views, dim=1), captured_token_states |
|
|
| def _distributed_batch_mean(self, x: torch.Tensor) -> torch.Tensor: |
| mean = x.mean(dim=0, keepdim=True) |
| if _dist_is_initialized(): |
| mean = mean.clone() |
| torch.distributed.all_reduce(mean, op=torch.distributed.ReduceOp.SUM) |
| mean /= float(torch.distributed.get_world_size()) |
| return mean |
|
|
| def _distributed_optional_rank_mean( |
| self, |
| x: Optional[torch.Tensor], |
| *, |
| output_dim: int, |
| device: torch.device, |
| ) -> Optional[torch.Tensor]: |
| if x is None or int(x.numel()) == 0 or int(x.shape[0]) == 0: |
| mean = torch.zeros((1, int(output_dim)), device=device, dtype=torch.float32) |
| active = torch.zeros([], device=device, dtype=torch.float32) |
| else: |
| mean = x.reshape(-1, x.shape[-1]).detach().float().mean(dim=0, keepdim=True) |
| active = torch.ones([], device=mean.device, dtype=torch.float32) |
|
|
| if _dist_is_initialized(): |
| torch.distributed.all_reduce(mean, op=torch.distributed.ReduceOp.SUM) |
| torch.distributed.all_reduce(active, op=torch.distributed.ReduceOp.SUM) |
|
|
| if float(active.item()) <= 0.0: |
| return None |
| return mean / active.clamp_min(1.0) |
|
|
| def _update_center(self, teacher_logits: torch.Tensor) -> None: |
| batch_center = self._distributed_batch_mean(teacher_logits.reshape(-1, teacher_logits.shape[-1]).detach()) |
| batch_center = batch_center.to(device=self.center.device, dtype=self.center.dtype) |
| momentum = float(getattr(self.cfg, "center_momentum", 0.9)) |
| self.center.mul_(momentum).add_(batch_center, alpha=1.0 - momentum) |
|
|
| def _update_token_center(self, teacher_logits: Optional[torch.Tensor]) -> None: |
| batch_center = self._distributed_optional_rank_mean( |
| teacher_logits, |
| output_dim=int(self.token_center.shape[-1]), |
| device=self.token_center.device, |
| ) |
| if batch_center is None: |
| return |
| batch_center = batch_center.to(device=self.token_center.device, dtype=self.token_center.dtype) |
| momentum = float(getattr(self.cfg, "token_center_momentum", getattr(self.cfg, "center_momentum", 0.9))) |
| self.token_center.mul_(momentum).add_(batch_center, alpha=1.0 - momentum) |
|
|
| def _teacher_probs(self, logits: torch.Tensor) -> torch.Tensor: |
| temperature = float(getattr(self.cfg, "teacher_temperature", 0.04)) |
| logits_f = logits.float() |
| center_f = self.center.to(device=logits_f.device, dtype=logits_f.dtype) |
| return F.softmax((logits_f - center_f) / temperature, dim=-1) |
|
|
| def _teacher_token_probs(self, logits: torch.Tensor) -> torch.Tensor: |
| temperature = float(getattr(self.cfg, "teacher_temperature", 0.04)) |
| logits_f = logits.float() |
| center_f = self.token_center.to(device=logits_f.device, dtype=logits_f.dtype) |
| return F.softmax((logits_f - center_f) / temperature, dim=-1) |
|
|
| def _student_log_probs(self, logits: torch.Tensor) -> torch.Tensor: |
| temperature = float(getattr(self.cfg, "student_temperature", 0.10)) |
| return F.log_softmax(logits.float() / temperature, dim=-1) |
|
|
| def _cross_entropy_target(self, student_logits: torch.Tensor, teacher_probs: torch.Tensor) -> torch.Tensor: |
| student_log_probs = self._student_log_probs(student_logits) |
| return (-teacher_probs * student_log_probs).sum(dim=-1).mean() |
|
|
| @staticmethod |
| def _mean_pairwise_abs_1d(values: torch.Tensor) -> torch.Tensor: |
| n = int(values.numel()) |
| if n <= 1: |
| return values.new_zeros(()) |
| sorted_values = values.sort().values |
| coeff = torch.arange(n, device=values.device, dtype=values.dtype) * 2 - (n - 1) |
| return (2.0 * (coeff * sorted_values).sum()) / float(n * n) |
|
|
| @staticmethod |
| def _mean_cross_abs_1d(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| nx = int(x.numel()) |
| ny = int(y.numel()) |
| if nx == 0 or ny == 0: |
| return x.new_zeros(()) |
| y_sorted = y.sort().values |
| prefix = torch.cat([y_sorted.new_zeros((1,)), y_sorted.cumsum(dim=0)]) |
| right = torch.searchsorted(y_sorted, x, right=False) |
| left_sum = prefix[right] |
| total_sum = prefix[-1] |
| right_f = right.to(dtype=x.dtype) |
| ny_f = float(ny) |
| abs_sum = x * right_f - left_sum + (total_sum - left_sum) - x * (ny_f - right_f) |
| return abs_sum.mean() / ny_f |
|
|
| def _gene_mmd_loss(self, pred_expr: torch.Tensor, target_expr: torch.Tensor, target_mask: torch.Tensor) -> torch.Tensor: |
| mask = target_mask.to(dtype=torch.bool) |
| losses = [] |
| for batch_idx in range(int(pred_expr.shape[0])): |
| valid = mask[batch_idx] |
| if not bool(valid.any().item()): |
| continue |
| pred = pred_expr[batch_idx, valid].float() |
| target = target_expr[batch_idx, valid].to(device=pred.device, dtype=pred.dtype) |
| if int(pred.numel()) <= 1: |
| losses.append(F.smooth_l1_loss(pred, target, reduction="mean")) |
| continue |
| cross = self._mean_cross_abs_1d(pred, target) |
| pred_pair = self._mean_pairwise_abs_1d(pred) |
| target_pair = self._mean_pairwise_abs_1d(target) |
| losses.append((2.0 * cross - pred_pair - target_pair).clamp_min(0.0)) |
| if not losses: |
| return pred_expr.new_zeros(()) |
| return torch.stack(losses).mean() |
|
|
| def _rank_loss(self, pred_expr: torch.Tensor, target_expr: torch.Tensor, target_mask: torch.Tensor) -> torch.Tensor: |
| num_pairs = int(getattr(self.cfg, "teacher_cls_decoder_rank_num_pairs", 4096)) |
| if num_pairs <= 0: |
| return pred_expr.new_zeros(()) |
| mask = target_mask.to(dtype=torch.bool) |
| losses = [] |
| for batch_idx in range(int(pred_expr.shape[0])): |
| valid = mask[batch_idx] |
| if int(valid.sum().item()) <= 1: |
| continue |
| pred = pred_expr[batch_idx, valid].float() |
| target = target_expr[batch_idx, valid].to(device=pred.device, dtype=pred.dtype) |
| n = int(pred.numel()) |
| device = pred.device |
| expressed = torch.nonzero(target > 0, as_tuple=False).flatten() |
| zero = torch.nonzero(target <= 0, as_tuple=False).flatten() |
| pair_i = [] |
| pair_j = [] |
| ez_pairs = num_pairs // 2 if int(expressed.numel()) > 0 and int(zero.numel()) > 0 else 0 |
| if ez_pairs > 0: |
| pair_i.append(expressed[torch.randint(int(expressed.numel()), (ez_pairs,), device=device)]) |
| pair_j.append(zero[torch.randint(int(zero.numel()), (ez_pairs,), device=device)]) |
| random_pairs = num_pairs - ez_pairs |
| if random_pairs > 0: |
| pair_i.append(torch.randint(n, (random_pairs,), device=device)) |
| pair_j.append(torch.randint(n, (random_pairs,), device=device)) |
| i = torch.cat(pair_i, dim=0) |
| j = torch.cat(pair_j, dim=0) |
| true_diff = target[i] - target[j] |
| true_sign = torch.sign(true_diff) |
| valid_pair = true_sign != 0 |
| if not bool(valid_pair.any().item()): |
| continue |
| pred_diff = pred[i] - pred[j] |
| losses.append(F.softplus(-true_sign[valid_pair] * pred_diff[valid_pair]).mean()) |
| if not losses: |
| return pred_expr.new_zeros(()) |
| return torch.stack(losses).mean() |
|
|
| def _rank_loss_random_pairs( |
| self, |
| pred_expr: torch.Tensor, |
| target_expr: torch.Tensor, |
| target_mask: torch.Tensor, |
| ) -> torch.Tensor: |
| num_pairs = int(getattr(self.cfg, "teacher_cls_decoder_rank_num_pairs", 4096)) |
| if num_pairs <= 0: |
| return pred_expr.new_zeros(()) |
| mask = target_mask.to(dtype=torch.bool) |
| losses = [] |
| for batch_idx in range(int(pred_expr.shape[0])): |
| valid = mask[batch_idx] |
| if int(valid.sum().item()) <= 1: |
| continue |
| pred = pred_expr[batch_idx, valid].float() |
| target = target_expr[batch_idx, valid].to(device=pred.device, dtype=pred.dtype) |
| n = int(pred.numel()) |
| device = pred.device |
| i = torch.randint(n, (num_pairs,), device=device) |
| j = torch.randint(n, (num_pairs,), device=device) |
| true_sign = torch.sign(target[i] - target[j]) |
| valid_pair = true_sign != 0 |
| if not bool(valid_pair.any().item()): |
| continue |
| pred_diff = pred[i] - pred[j] |
| losses.append(F.softplus(-true_sign[valid_pair] * pred_diff[valid_pair]).mean()) |
| if not losses: |
| return pred_expr.new_zeros(()) |
| return torch.stack(losses).mean() |
|
|
| def _masked_mse_loss( |
| self, |
| pred_expr: torch.Tensor, |
| target_expr: torch.Tensor, |
| target_mask: torch.Tensor, |
| ) -> torch.Tensor: |
| mask = target_mask.to(device=pred_expr.device, dtype=torch.bool) |
| if not bool(mask.any().item()): |
| return pred_expr.new_zeros(()) |
| target = target_expr.to(device=pred_expr.device, dtype=pred_expr.dtype) |
| return F.mse_loss(pred_expr[mask].float(), target[mask].float(), reduction="mean") |
|
|
| def _teacher_cls_decoder_loss( |
| self, |
| *, |
| teacher_gene_ids: torch.Tensor, |
| teacher_bin_ids: torch.Tensor, |
| teacher_attention_mask: torch.Tensor, |
| teacher_view_use_bin: torch.Tensor, |
| target_gene_ids: torch.Tensor, |
| target_expr: torch.Tensor, |
| target_mask: torch.Tensor, |
| obs_features: Optional[torch.Tensor] = None, |
| ) -> torch.Tensor: |
| if self.teacher_cls_expression_decoder is None: |
| return torch.zeros([], device=teacher_gene_ids.device, dtype=torch.float32) |
| if target_gene_ids.numel() == 0 or not bool(target_mask.any().item()): |
| return torch.zeros([], device=teacher_gene_ids.device, dtype=torch.float32) |
|
|
| decoder_tokens = self._encode_tokens( |
| teacher_gene_ids[:, 0, :], |
| teacher_bin_ids[:, 0, :], |
| teacher_attention_mask[:, 0, :], |
| use_bin_embeddings=bool(teacher_view_use_bin[0].item()), |
| backbone="student", |
| ) |
| decoder_cls = self._pool_cell_repr(decoder_tokens) |
| gene_queries = self.student_backbone.tokenizer.gene_emb(target_gene_ids) |
| pred_expr = self.teacher_cls_expression_decoder(decoder_cls, gene_queries, obs_features=obs_features) |
|
|
| gene_mmd = self._gene_mmd_loss(pred_expr, target_expr, target_mask) |
| rank_alpha = float(getattr(self.cfg, "teacher_cls_decoder_rank_loss_alpha", 0.05)) |
| if rank_alpha <= 0.0: |
| return gene_mmd |
| return gene_mmd + rank_alpha * self._rank_loss(pred_expr, target_expr, target_mask) |
|
|
| def _teacher_student_cls_decoder_loss( |
| self, |
| *, |
| teacher_gene_ids: torch.Tensor, |
| teacher_bin_ids: torch.Tensor, |
| teacher_attention_mask: torch.Tensor, |
| teacher_view_use_bin: torch.Tensor, |
| student_gene_ids: torch.Tensor, |
| student_bin_ids: torch.Tensor, |
| student_attention_mask: torch.Tensor, |
| student_view_use_bin: torch.Tensor, |
| target_gene_ids: torch.Tensor, |
| target_expr: torch.Tensor, |
| target_mask: torch.Tensor, |
| obs_features: Optional[torch.Tensor] = None, |
| ) -> torch.Tensor: |
| if self.student_cls_expression_decoder is None or self.teacher_cls_expression_decoder is None: |
| return torch.zeros([], device=teacher_gene_ids.device, dtype=torch.float32) |
| if target_gene_ids.numel() == 0 or not bool(target_mask.any().item()): |
| return torch.zeros([], device=teacher_gene_ids.device, dtype=torch.float32) |
|
|
| student_tokens = self._encode_tokens( |
| student_gene_ids[:, 0, :], |
| student_bin_ids[:, 0, :], |
| student_attention_mask[:, 0, :], |
| use_bin_embeddings=bool(student_view_use_bin[0].item()), |
| backbone="student", |
| ) |
| student_cls = self._pool_cell_repr(student_tokens) |
| student_gene_queries = self.student_backbone.tokenizer.gene_emb(target_gene_ids) |
| student_scores = self.student_cls_expression_decoder( |
| student_cls, |
| student_gene_queries, |
| obs_features=obs_features, |
| ) |
|
|
| with torch.no_grad(): |
| teacher_tokens = self._encode_tokens( |
| teacher_gene_ids[:, 0, :], |
| teacher_bin_ids[:, 0, :], |
| teacher_attention_mask[:, 0, :], |
| use_bin_embeddings=bool(teacher_view_use_bin[0].item()), |
| backbone="teacher", |
| ) |
| teacher_cls = self._pool_cell_repr(teacher_tokens) |
| teacher_gene_queries = self.teacher_backbone.tokenizer.gene_emb(target_gene_ids) |
| teacher_scores = self.teacher_cls_expression_decoder( |
| teacher_cls, |
| teacher_gene_queries, |
| obs_features=obs_features, |
| ) |
|
|
| distill_weight = float(getattr(self.cfg, "teacher_student_decoder_distill_weight", 0.0)) |
| loss_mode = str(getattr(self.cfg, "teacher_student_decoder_loss_mode", "distribution")).strip().lower() |
| if loss_mode == "mse": |
| raw_mse_weight = float(getattr(self.cfg, "teacher_student_decoder_raw_mse_weight", 1.0)) |
| loss = student_scores.new_zeros(()) |
| if distill_weight > 0.0: |
| loss = loss + distill_weight * self._masked_mse_loss(student_scores, teacher_scores, target_mask) |
| if raw_mse_weight > 0.0: |
| loss = loss + raw_mse_weight * self._masked_mse_loss(student_scores, target_expr, target_mask) |
| return loss |
| if loss_mode != "distribution": |
| raise ValueError(f"Unsupported teacher_student_decoder_loss_mode={loss_mode!r}") |
|
|
| rank_alpha = float(getattr(self.cfg, "teacher_student_decoder_rank_alpha", 0.05)) |
| raw_mmd_weight = float(getattr(self.cfg, "teacher_student_decoder_raw_mmd_weight", 0.1)) |
| raw_rank_weight = float(getattr(self.cfg, "teacher_student_decoder_raw_rank_weight", 0.01)) |
|
|
| loss = student_scores.new_zeros(()) |
| if distill_weight > 0.0: |
| loss = loss + distill_weight * self._gene_mmd_loss(student_scores, teacher_scores, target_mask) |
| if rank_alpha > 0.0: |
| loss = loss + distill_weight * rank_alpha * self._rank_loss_random_pairs(student_scores, teacher_scores, target_mask) |
| if raw_mmd_weight > 0.0: |
| loss = loss + raw_mmd_weight * self._gene_mmd_loss(student_scores, target_expr, target_mask) |
| if raw_rank_weight > 0.0: |
| loss = loss + raw_rank_weight * self._rank_loss(student_scores, target_expr, target_mask) |
| return loss |
|
|
| def _allow_cross_mode_pair(self, *, teacher_is_bin: bool, student_is_bin: bool) -> bool: |
| mode = str(getattr(self.cfg, "cross_modal_pairing_mode", "bin_to_gene_only")) |
| if mode == "disabled": |
| return False |
| if mode == "symmetric": |
| return teacher_is_bin != student_is_bin |
| if mode == "bin_to_gene_only": |
| return bool(teacher_is_bin) and (not bool(student_is_bin)) |
| raise ValueError(f"Unknown cross_modal_pairing_mode={mode!r}") |
|
|
| def _pairwise_dino_loss( |
| self, |
| student_logits: torch.Tensor, |
| student_view_use_bin: torch.Tensor, |
| teacher_logits: torch.Tensor, |
| teacher_view_use_bin: torch.Tensor, |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| same_losses = [] |
| cross_losses = [] |
| cross_weight = float(getattr(self.cfg, "cross_modal_dino_weight", 0.0)) |
|
|
| num_teacher_views = int(teacher_logits.shape[1]) |
| num_student_views = int(student_logits.shape[1]) |
|
|
| for teacher_idx in range(num_teacher_views): |
| teacher_probs = self._teacher_probs(teacher_logits[:, teacher_idx, :]).detach() |
| teacher_is_bin = bool(teacher_view_use_bin[teacher_idx].item()) |
| for student_idx in range(num_student_views): |
| student_is_bin = bool(student_view_use_bin[student_idx].item()) |
| pair_loss = self._cross_entropy_target( |
| student_logits[:, student_idx, :], |
| teacher_probs, |
| ) |
| if student_is_bin == teacher_is_bin: |
| same_losses.append(pair_loss) |
| elif cross_weight > 0.0 and self._allow_cross_mode_pair( |
| teacher_is_bin=teacher_is_bin, |
| student_is_bin=student_is_bin, |
| ): |
| cross_losses.append(pair_loss) |
|
|
| device = student_logits.device |
| same_loss = ( |
| torch.stack(same_losses).mean() |
| if same_losses |
| else torch.zeros([], device=device, dtype=torch.float32) |
| ) |
| cross_loss = ( |
| torch.stack(cross_losses).mean() |
| if cross_losses |
| else torch.zeros([], device=device, dtype=torch.float32) |
| ) |
| total = same_loss + float(cross_weight) * cross_loss |
| return total, same_loss, cross_loss |
|
|
| def _matched_token_states_for_pair( |
| self, |
| *, |
| teacher_tokens: torch.Tensor, |
| teacher_gene_ids: torch.Tensor, |
| teacher_attention_mask: torch.Tensor, |
| student_tokens: torch.Tensor, |
| student_gene_ids: torch.Tensor, |
| student_attention_mask: torch.Tensor, |
| student_ibot_mask: torch.Tensor, |
| use_masked_tokens: bool = True, |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| teacher_states: List[torch.Tensor] = [] |
| student_states: List[torch.Tensor] = [] |
|
|
| batch_size = int(teacher_tokens.shape[0]) |
| hidden_dim = int(teacher_tokens.shape[-1]) |
| for batch_idx in range(batch_size): |
| teacher_valid = teacher_attention_mask[batch_idx].bool().clone() |
| student_attention_valid = student_attention_mask[batch_idx].bool() |
| student_mask = student_ibot_mask[batch_idx].bool() |
| if use_masked_tokens: |
| student_valid = (student_attention_valid & student_mask).clone() |
| else: |
| student_valid = (student_attention_valid & (~student_mask)).clone() |
|
|
| if teacher_valid.numel() <= 1 or student_valid.numel() <= 1: |
| continue |
|
|
| teacher_valid[0] = False |
| student_valid[0] = False |
| if not teacher_valid.any().item() or not student_valid.any().item(): |
| continue |
|
|
| teacher_positions = torch.nonzero(teacher_valid, as_tuple=False).squeeze(-1) |
| student_positions = torch.nonzero(student_valid, as_tuple=False).squeeze(-1) |
| teacher_genes = teacher_gene_ids[batch_idx, teacher_positions] |
| student_genes = student_gene_ids[batch_idx, student_positions] |
|
|
| matches = student_genes.unsqueeze(1) == teacher_genes.unsqueeze(0) |
| matched_student = matches.any(dim=1) |
| if not matched_student.any().item(): |
| continue |
|
|
| matched_student_positions = student_positions[matched_student] |
| matched_teacher_rel_pos = matches[matched_student].to(dtype=torch.int64).argmax(dim=1) |
| matched_teacher_positions = teacher_positions[matched_teacher_rel_pos] |
|
|
| teacher_states.append(teacher_tokens[batch_idx, matched_teacher_positions, :]) |
| student_states.append(student_tokens[batch_idx, matched_student_positions, :]) |
|
|
| if teacher_states: |
| return torch.cat(teacher_states, dim=0), torch.cat(student_states, dim=0) |
|
|
| empty = teacher_tokens.new_empty((0, hidden_dim)) |
| return empty, empty |
|
|
| def _pairwise_token_loss( |
| self, |
| *, |
| student_token_views: Dict[int, torch.Tensor], |
| student_gene_ids: torch.Tensor, |
| student_attention_mask: torch.Tensor, |
| student_ibot_mask: torch.Tensor, |
| student_view_use_bin: torch.Tensor, |
| student_view_is_global: torch.Tensor, |
| teacher_token_views: Dict[int, torch.Tensor], |
| teacher_gene_ids: torch.Tensor, |
| teacher_attention_mask: torch.Tensor, |
| teacher_view_use_bin: torch.Tensor, |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: |
| same_losses = [] |
| cross_losses = [] |
| unmasked_same_losses = [] |
| unmasked_cross_losses = [] |
| teacher_logits_for_center: List[torch.Tensor] = [] |
| cross_weight = float(getattr(self.cfg, "cross_modal_dino_weight", 0.0)) |
| unmasked_weight = float(getattr(self.cfg, "token_unmasked_loss_weight", 0.0)) |
|
|
| for teacher_idx, teacher_tokens in teacher_token_views.items(): |
| teacher_is_bin = bool(teacher_view_use_bin[teacher_idx].item()) |
| for student_idx, student_tokens in student_token_views.items(): |
| if not bool(student_view_is_global[student_idx].item()): |
| continue |
| if not student_ibot_mask[:, student_idx, :].any().item(): |
| continue |
|
|
| student_is_bin = bool(student_view_use_bin[student_idx].item()) |
| if student_is_bin == teacher_is_bin: |
| allow_pair = True |
| elif cross_weight > 0.0: |
| allow_pair = self._allow_cross_mode_pair( |
| teacher_is_bin=teacher_is_bin, |
| student_is_bin=student_is_bin, |
| ) |
| else: |
| allow_pair = False |
| if not allow_pair: |
| continue |
|
|
| teacher_states, student_states = self._matched_token_states_for_pair( |
| teacher_tokens=teacher_tokens, |
| teacher_gene_ids=teacher_gene_ids[:, teacher_idx, :], |
| teacher_attention_mask=teacher_attention_mask[:, teacher_idx, :], |
| student_tokens=student_tokens, |
| student_gene_ids=student_gene_ids[:, student_idx, :], |
| student_attention_mask=student_attention_mask[:, student_idx, :], |
| student_ibot_mask=student_ibot_mask[:, student_idx, :], |
| use_masked_tokens=True, |
| ) |
| if teacher_states.shape[0] > 0: |
| teacher_logits = self.teacher_ibot_head(teacher_states) |
| teacher_probs = self._teacher_token_probs(teacher_logits).detach() |
| student_logits = self.student_ibot_head(student_states) |
| pair_loss = self._cross_entropy_target(student_logits, teacher_probs) |
| teacher_logits_for_center.append(teacher_logits.detach()) |
|
|
| if student_is_bin == teacher_is_bin: |
| same_losses.append(pair_loss) |
| else: |
| cross_losses.append(pair_loss) |
|
|
| if unmasked_weight > 0.0: |
| unmasked_teacher_states, unmasked_student_states = self._matched_token_states_for_pair( |
| teacher_tokens=teacher_tokens, |
| teacher_gene_ids=teacher_gene_ids[:, teacher_idx, :], |
| teacher_attention_mask=teacher_attention_mask[:, teacher_idx, :], |
| student_tokens=student_tokens, |
| student_gene_ids=student_gene_ids[:, student_idx, :], |
| student_attention_mask=student_attention_mask[:, student_idx, :], |
| student_ibot_mask=student_ibot_mask[:, student_idx, :], |
| use_masked_tokens=False, |
| ) |
| if unmasked_teacher_states.shape[0] > 0: |
| unmasked_teacher_logits = self.teacher_ibot_head(unmasked_teacher_states) |
| unmasked_teacher_probs = self._teacher_token_probs(unmasked_teacher_logits).detach() |
| unmasked_student_logits = self.student_ibot_head(unmasked_student_states) |
| unmasked_pair_loss = self._cross_entropy_target( |
| unmasked_student_logits, |
| unmasked_teacher_probs, |
| ) |
| if student_is_bin == teacher_is_bin: |
| unmasked_same_losses.append(unmasked_pair_loss) |
| else: |
| unmasked_cross_losses.append(unmasked_pair_loss) |
|
|
| device = student_gene_ids.device |
| same_loss = ( |
| torch.stack(same_losses).mean() |
| if same_losses |
| else torch.zeros([], device=device, dtype=torch.float32) |
| ) |
| cross_loss = ( |
| torch.stack(cross_losses).mean() |
| if cross_losses |
| else torch.zeros([], device=device, dtype=torch.float32) |
| ) |
| unmasked_same_loss = ( |
| torch.stack(unmasked_same_losses).mean() |
| if unmasked_same_losses |
| else torch.zeros([], device=device, dtype=torch.float32) |
| ) |
| unmasked_cross_loss = ( |
| torch.stack(unmasked_cross_losses).mean() |
| if unmasked_cross_losses |
| else torch.zeros([], device=device, dtype=torch.float32) |
| ) |
| masked_total = same_loss + float(cross_weight) * cross_loss |
| unmasked_total = unmasked_same_loss + float(cross_weight) * unmasked_cross_loss |
| total = masked_total + unmasked_weight * unmasked_total |
| teacher_logits_center = torch.cat(teacher_logits_for_center, dim=0) if teacher_logits_for_center else None |
| self._last_token_unmasked_loss = unmasked_total.detach() |
| self._last_token_unmasked_same_mode_loss = unmasked_same_loss.detach() |
| self._last_token_unmasked_cross_mode_loss = unmasked_cross_loss.detach() |
| return total, same_loss, cross_loss, teacher_logits_center |
|
|
| def _teacher_momentum_value(self, global_step: int, max_steps: int) -> float: |
| start = float(getattr(self.cfg, "teacher_momentum", 0.996)) |
| end = float(getattr(self.cfg, "teacher_momentum_end", 1.0)) |
| if max_steps <= 1: |
| return start |
| progress = min(max(float(global_step), 0.0), float(max_steps - 1)) / float(max_steps - 1) |
| cosine = 0.5 * (1.0 + math.cos(math.pi * progress)) |
| return end - (end - start) * cosine |
|
|
| def _ema_update_module(self, student_module: nn.Module, teacher_module: nn.Module, momentum: float) -> None: |
| with ExitStack() as stack: |
| if FSDP is not None: |
| for wrapped in _leaf_fsdp_modules(student_module): |
| stack.enter_context(FSDP.summon_full_params(wrapped, recurse=False, writeback=False)) |
| for wrapped in _leaf_fsdp_modules(teacher_module): |
| stack.enter_context(FSDP.summon_full_params(wrapped, recurse=False, writeback=True)) |
|
|
| student_params = dict(student_module.named_parameters()) |
| teacher_params = dict(teacher_module.named_parameters()) |
| if student_params.keys() != teacher_params.keys(): |
| missing_student = sorted(set(teacher_params.keys()) - set(student_params.keys())) |
| missing_teacher = sorted(set(student_params.keys()) - set(teacher_params.keys())) |
| raise RuntimeError( |
| "Student/teacher parameter names do not match for EMA update: " |
| f"missing_student={missing_student[:5]}, missing_teacher={missing_teacher[:5]}" |
| ) |
|
|
| for name, teacher_param in teacher_params.items(): |
| student_param = student_params[name] |
| if teacher_param.data.shape != student_param.data.shape: |
| raise RuntimeError( |
| f"EMA shape mismatch for {name}: teacher={tuple(teacher_param.shape)} " |
| f"student={tuple(student_param.shape)}" |
| ) |
| teacher_param.data.mul_(momentum).add_(student_param.data, alpha=1.0 - momentum) |
|
|
| @torch.no_grad() |
| def update_teacher(self, *, global_step: int, max_steps: int) -> float: |
| momentum = self._teacher_momentum_value(global_step=global_step, max_steps=max_steps) |
| self.last_teacher_momentum.fill_(float(momentum)) |
|
|
| self._ema_update_module(self.student_backbone, self.teacher_backbone, momentum) |
| self._ema_update_module(self.student_head, self.teacher_head, momentum) |
| self._ema_update_module(self.student_ibot_head, self.teacher_ibot_head, momentum) |
| if self.student_cls_expression_decoder is not None and self.teacher_cls_expression_decoder is not None: |
| self._ema_update_module( |
| self.student_cls_expression_decoder, |
| self.teacher_cls_expression_decoder, |
| momentum, |
| ) |
| return momentum |
|
|
| def forward(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: |
| student_gene_ids = batch["student_gene_ids"] |
| student_bin_ids = batch["student_bin_ids"] |
| student_attention_mask = batch["student_attention_mask"] |
| student_ibot_mask = batch["student_ibot_mask"] |
| student_view_use_bin = batch["student_view_use_bin"] |
| student_view_is_global = batch["student_view_is_global"] |
|
|
| teacher_gene_ids = batch["teacher_gene_ids"] |
| teacher_bin_ids = batch["teacher_bin_ids"] |
| teacher_attention_mask = batch["teacher_attention_mask"] |
| teacher_view_use_bin = batch["teacher_view_use_bin"] |
|
|
| student_capture_views = tuple( |
| idx for idx in range(int(student_view_is_global.numel())) if bool(student_view_is_global[idx].item()) |
| ) |
| teacher_capture_views = tuple(range(int(teacher_view_use_bin.numel()))) |
|
|
| _, student_logits, student_token_views = self._encode_view_stack( |
| student_gene_ids, |
| student_bin_ids, |
| student_attention_mask, |
| student_view_use_bin, |
| backbone_name="student", |
| head=self.student_head, |
| capture_token_views=student_capture_views, |
| ) |
|
|
| with torch.no_grad(): |
| _, teacher_logits, teacher_token_views = self._encode_view_stack( |
| teacher_gene_ids, |
| teacher_bin_ids, |
| teacher_attention_mask, |
| teacher_view_use_bin, |
| backbone_name="teacher", |
| head=self.teacher_head, |
| capture_token_views=teacher_capture_views, |
| ) |
|
|
| loss_dino, loss_same_mode, loss_cross_mode = self._pairwise_dino_loss( |
| student_logits, |
| student_view_use_bin, |
| teacher_logits, |
| teacher_view_use_bin, |
| ) |
|
|
| token_loss_weight = float(getattr(self.cfg, "token_loss_weight", 1.0)) |
| if token_loss_weight > 0.0: |
| loss_token, loss_token_same_mode, loss_token_cross_mode, teacher_token_center_logits = self._pairwise_token_loss( |
| student_token_views=student_token_views, |
| student_gene_ids=student_gene_ids, |
| student_attention_mask=student_attention_mask, |
| student_ibot_mask=student_ibot_mask, |
| student_view_use_bin=student_view_use_bin, |
| student_view_is_global=student_view_is_global, |
| teacher_token_views=teacher_token_views, |
| teacher_gene_ids=teacher_gene_ids, |
| teacher_attention_mask=teacher_attention_mask, |
| teacher_view_use_bin=teacher_view_use_bin, |
| ) |
| else: |
| zero = torch.zeros([], device=student_logits.device, dtype=torch.float32) |
| loss_token = zero |
| loss_token_same_mode = zero |
| loss_token_cross_mode = zero |
| self._last_token_unmasked_loss = zero.detach() |
| self._last_token_unmasked_same_mode_loss = zero.detach() |
| self._last_token_unmasked_cross_mode_loss = zero.detach() |
| teacher_token_center_logits = None |
|
|
| decoder_loss_weight = float(getattr(self.cfg, "teacher_cls_decoder_loss_weight", 0.0)) |
| if decoder_loss_weight > 0.0: |
| if float(getattr(self.cfg, "teacher_student_decoder_distill_weight", 0.0)) > 0.0: |
| loss_decoder = self._teacher_student_cls_decoder_loss( |
| teacher_gene_ids=teacher_gene_ids, |
| teacher_bin_ids=teacher_bin_ids, |
| teacher_attention_mask=teacher_attention_mask, |
| teacher_view_use_bin=teacher_view_use_bin, |
| student_gene_ids=student_gene_ids, |
| student_bin_ids=student_bin_ids, |
| student_attention_mask=student_attention_mask, |
| student_view_use_bin=student_view_use_bin, |
| target_gene_ids=batch["teacher_cls_decoder_target_gene_ids"], |
| target_expr=batch["teacher_cls_decoder_target_expr"], |
| target_mask=batch["teacher_cls_decoder_target_mask"], |
| obs_features=batch.get("teacher_cls_decoder_obs_features"), |
| ) |
| else: |
| loss_decoder = self._teacher_cls_decoder_loss( |
| teacher_gene_ids=teacher_gene_ids, |
| teacher_bin_ids=teacher_bin_ids, |
| teacher_attention_mask=teacher_attention_mask, |
| teacher_view_use_bin=teacher_view_use_bin, |
| target_gene_ids=batch["teacher_cls_decoder_target_gene_ids"], |
| target_expr=batch["teacher_cls_decoder_target_expr"], |
| target_mask=batch["teacher_cls_decoder_target_mask"], |
| obs_features=batch.get("teacher_cls_decoder_obs_features"), |
| ) |
| else: |
| loss_decoder = torch.zeros([], device=student_logits.device, dtype=torch.float32) |
|
|
| total_loss = loss_dino + token_loss_weight * loss_token + decoder_loss_weight * loss_decoder |
|
|
| with torch.no_grad(): |
| self._update_center(teacher_logits) |
| self._update_token_center(teacher_token_center_logits) |
|
|
| return { |
| "loss": total_loss, |
| "loss_dino": loss_dino.detach(), |
| "loss_dino_same_mode": loss_same_mode.detach(), |
| "loss_dino_cross_mode": loss_cross_mode.detach(), |
| "loss_token": loss_token.detach(), |
| "loss_token_same_mode": loss_token_same_mode.detach(), |
| "loss_token_cross_mode": loss_token_cross_mode.detach(), |
| "loss_token_unmasked": getattr(self, "_last_token_unmasked_loss", loss_token.new_zeros(())).detach(), |
| "loss_token_unmasked_same_mode": getattr( |
| self, "_last_token_unmasked_same_mode_loss", loss_token.new_zeros(()) |
| ).detach(), |
| "loss_token_unmasked_cross_mode": getattr( |
| self, "_last_token_unmasked_cross_mode_loss", loss_token.new_zeros(()) |
| ).detach(), |
| "loss_decoder": loss_decoder.detach(), |
| "loss_total": total_loss.detach(), |
| "loss_contrastive": total_loss.detach(), |
| "teacher_momentum": self.last_teacher_momentum.detach(), |
| } |
|
|
| def save_pretrained( |
| self, |
| save_directory: str, |
| *, |
| save_safetensors: bool = True, |
| state_dict: Optional[Dict[str, torch.Tensor]] = None, |
| gene_vocab: Optional[Dict[str, Any]] = None, |
| meta_vocab: Optional[Dict[str, Any]] = None, |
| manifest: Optional[Any] = None, |
| extra_files: Optional[Dict[str, Any]] = None, |
| ) -> None: |
| os.makedirs(save_directory, exist_ok=True) |
|
|
| cfg_path = os.path.join(save_directory, "model_config.json") |
| _json_dump(self.cfg.__dict__, cfg_path) |
|
|
| if state_dict is None: |
| state_dict = self.state_dict() |
| cpu_state = {k: v.detach().cpu() for k, v in state_dict.items()} |
|
|
| if save_safetensors: |
| import safetensors.torch as st |
|
|
| st.save_file(cpu_state, os.path.join(save_directory, "model.safetensors")) |
| else: |
| torch.save(cpu_state, os.path.join(save_directory, "pytorch_model.bin")) |
|
|
| if gene_vocab is not None: |
| _json_dump(gene_vocab, os.path.join(save_directory, "gene_vocab.json")) |
| if meta_vocab is not None: |
| _json_dump(meta_vocab, os.path.join(save_directory, "meta_vocab.json")) |
| if manifest is not None: |
| _json_dump(manifest, os.path.join(save_directory, "manifest.json")) |
| if extra_files: |
| for name, obj in extra_files.items(): |
| _json_dump(obj, os.path.join(save_directory, name)) |
|
|
| @classmethod |
| def from_pretrained( |
| cls, |
| load_directory: str, |
| *, |
| device: Optional[str] = None, |
| dtype: Optional[torch.dtype] = None, |
| strict: bool = True, |
| load_assets: bool = True, |
| ) -> Tuple["CellContrastivePredictModel", Optional[Dict[str, Any]], Optional[Dict[str, Any]], Optional[Any]]: |
| cfg_path = os.path.join(load_directory, "model_config.json") |
| if not os.path.exists(cfg_path): |
| raise FileNotFoundError(f"Missing model_config.json in: {load_directory}") |
|
|
| with open(cfg_path, "r", encoding="utf-8") as f: |
| cfg_dict = json.load(f) |
|
|
| cfg_fields = {f.name for f in dataclass_fields(FoundationConfig)} |
| filtered_cfg = {k: v for k, v in cfg_dict.items() if k in cfg_fields} |
| cfg = FoundationConfig(**filtered_cfg) |
| model = cls(cfg) |
|
|
| st_path = os.path.join(load_directory, "model.safetensors") |
| pt_path = os.path.join(load_directory, "pytorch_model.bin") |
| if os.path.exists(st_path): |
| import safetensors.torch as st |
|
|
| state_dict = st.load_file(st_path, device="cpu") |
| elif os.path.exists(pt_path): |
| state_dict = torch.load(pt_path, map_location="cpu") |
| else: |
| raise FileNotFoundError(f"Missing model weights in: {load_directory}") |
|
|
| if "token_center" not in state_dict: |
| state_dict["token_center"] = torch.zeros_like(model.token_center) |
|
|
| for missing_prefix, source_prefix in ( |
| ("student_ibot_head.", "student_head."), |
| ("teacher_ibot_head.", "teacher_head."), |
| ): |
| if any(key.startswith(missing_prefix) for key in state_dict): |
| continue |
| for model_key in model.state_dict().keys(): |
| if not model_key.startswith(missing_prefix): |
| continue |
| source_key = source_prefix + model_key[len(missing_prefix) :] |
| if source_key in state_dict: |
| state_dict[model_key] = state_dict[source_key].clone() |
|
|
| model.load_state_dict(state_dict, strict=strict) |
| if dtype is not None: |
| model = model.to(dtype=dtype) |
| if device is not None: |
| model = model.to(device=device) |
|
|
| gene_vocab = None |
| meta_vocab = None |
| manifest = None |
| if load_assets: |
| gene_vocab_path = os.path.join(load_directory, "gene_vocab.json") |
| meta_vocab_path = os.path.join(load_directory, "meta_vocab.json") |
| manifest_path = os.path.join(load_directory, "manifest.json") |
| if os.path.exists(gene_vocab_path): |
| with open(gene_vocab_path, "r", encoding="utf-8") as f: |
| gene_vocab = json.load(f) |
| if os.path.exists(meta_vocab_path): |
| with open(meta_vocab_path, "r", encoding="utf-8") as f: |
| meta_vocab = json.load(f) |
| if os.path.exists(manifest_path): |
| with open(manifest_path, "r", encoding="utf-8") as f: |
| manifest = json.load(f) |
| return model, gene_vocab, meta_vocab, manifest |
|
|
|
|
| __all__ = ["CellContrastivePredictModel", "FoundationConfig"] |
|
|
