| """Losses and leakage-free gene-set augmentation.""" |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Sequence |
|
|
| import torch |
| from torch.nn import functional as F |
|
|
|
|
| def supervised_contrastive_loss( |
| embeddings: torch.Tensor, |
| group_ids: Sequence[str], |
| temperature: float = 0.07, |
| ) -> torch.Tensor: |
| if embeddings.ndim != 2 or embeddings.shape[0] != len(group_ids): |
| raise ValueError("Embeddings and group IDs have inconsistent batch dimensions") |
| if temperature <= 0: |
| raise ValueError("Temperature must be positive") |
| normalized = F.normalize(embeddings, dim=1) |
| logits = normalized @ normalized.T / temperature |
| identity = torch.eye(len(group_ids), dtype=torch.bool, device=embeddings.device) |
| groups = torch.tensor( |
| [[left == right for right in group_ids] for left in group_ids], |
| dtype=torch.bool, |
| device=embeddings.device, |
| ) |
| positives = groups & ~identity |
| positive_counts = positives.sum(dim=1) |
| if torch.any(positive_counts == 0): |
| raise ValueError("Every item must have at least one positive in its batch") |
| logits = logits.masked_fill(identity, float("-inf")) |
| log_probabilities = logits - torch.logsumexp(logits, dim=1, keepdim=True) |
| positive_log_probability = log_probabilities.masked_fill(~positives, 0.0).sum(dim=1) |
| positive_log_probability = positive_log_probability / positive_counts |
| return -positive_log_probability.mean() |
|
|
|
|
| def augment_gene_sets( |
| embeddings: torch.Tensor, |
| positions: torch.Tensor, |
| padding_mask: torch.Tensor, |
| gene_dropout: float, |
| position_jitter: float, |
| generator: torch.Generator | None = None, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| augmented_embeddings = embeddings.clone() |
| augmented_positions = positions.clone() |
| augmented_mask = padding_mask.clone() |
| for row in range(embeddings.shape[0]): |
| valid_indices = torch.nonzero(~padding_mask[row], as_tuple=False).flatten() |
| maximum_drop = max(0, len(valid_indices) - 2) |
| drop_count = min(maximum_drop, int(len(valid_indices) * gene_dropout)) |
| if drop_count: |
| order = torch.randperm(len(valid_indices), generator=generator, device="cpu") |
| dropped = valid_indices.cpu()[order[:drop_count]].to(embeddings.device) |
| augmented_mask[row, dropped] = True |
| augmented_embeddings[row, dropped] = 0 |
| noise = torch.empty(positions.shape[1], device=positions.device) |
| noise.uniform_(-position_jitter, position_jitter, generator=generator) |
| augmented_positions[row] = (augmented_positions[row] + noise).clamp(0.0, 1.0) |
| augmented_positions[row, augmented_mask[row]] = 0.0 |
| return augmented_embeddings, augmented_positions, augmented_mask |
|
|
|
|
| def masked_gene_loss( |
| predicted: torch.Tensor, |
| target: torch.Tensor, |
| masked_positions: torch.Tensor, |
| ) -> torch.Tensor: |
| if not masked_positions.any(): |
| raise ValueError("At least one gene must be masked") |
| return F.mse_loss(predicted[masked_positions], target[masked_positions]) |
|
|
