| from __future__ import annotations |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| class StochasticGate(nn.Module): |
| def __init__( |
| self, |
| input_dim: int, |
| sigma: float, |
| temperature: float, |
| top_k: int, |
| ) -> None: |
| super().__init__() |
| self.logits = nn.Parameter(torch.zeros(input_dim)) |
| self.sigma = sigma |
| self.temperature = temperature |
| self.top_k = top_k |
|
|
| def gate_probabilities(self) -> torch.Tensor: |
| return torch.sigmoid(self.logits / self.temperature) |
|
|
| def forward(self, x: torch.Tensor, hard_topk: bool = False) -> torch.Tensor: |
| probabilities = self.gate_probabilities() |
| gates = probabilities |
| if hard_topk and 0 < self.top_k < gates.numel(): |
| _, top_indices = torch.topk(probabilities, self.top_k) |
| mask = torch.zeros_like(gates) |
| mask[top_indices] = 1.0 |
| gates = mask |
| return x * gates |
|
|
|
|
| class ResidualMLPBlock(nn.Module): |
| def __init__(self, dim: int, dropout: float) -> None: |
| super().__init__() |
| self.block = nn.Sequential( |
| nn.LayerNorm(dim), |
| nn.Linear(dim, dim * 2), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(dim * 2, dim), |
| nn.Dropout(dropout), |
| ) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return x + self.block(x) |
|
|
|
|
| class KGStableGateDANN(nn.Module): |
| def __init__( |
| self, |
| input_dim: int, |
| n_classes: int, |
| n_domains: int, |
| hidden_dim: int, |
| projection_dim: int, |
| encoder_blocks: int, |
| dropout: float, |
| top_k: int, |
| gate_sigma: float, |
| gate_temperature: float, |
| ) -> None: |
| super().__init__() |
| self.gates = StochasticGate(input_dim, gate_sigma, gate_temperature, top_k) |
| self.input_projection = nn.Sequential( |
| nn.Linear(input_dim, hidden_dim), |
| nn.LayerNorm(hidden_dim), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| ) |
| self.encoder = nn.Sequential( |
| *[ResidualMLPBlock(hidden_dim, dropout) for _ in range(encoder_blocks)] |
| ) |
| self.projection_head = nn.Sequential( |
| nn.LayerNorm(hidden_dim), |
| nn.Linear(hidden_dim, projection_dim), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(projection_dim, projection_dim), |
| nn.LayerNorm(projection_dim), |
| ) |
| self.disease_classifier = nn.Sequential( |
| nn.Linear(projection_dim, max(2, projection_dim // 2)), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(max(2, projection_dim // 2), n_classes), |
| ) |
| self.domain_classifier = nn.Sequential( |
| nn.Linear(projection_dim, max(2, projection_dim // 2)), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(max(2, projection_dim // 2), n_domains), |
| ) |
| self.register_buffer("gate_prior_target", torch.zeros(input_dim)) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = self.gates(x, hard_topk=True) |
| hidden = self.input_projection(x) |
| hidden = self.encoder(hidden) |
| embedding = self.projection_head(hidden) |
| return self.disease_classifier(embedding) |
|
|
