| import math |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| class PhaseFormerTransformerLayer(nn.Module): |
| """ |
| Transformer layer with phase-based temporal gating applied |
| to attention and feed-forward residual paths. |
| |
| Args: |
| d_model (int): Input/output dimension. |
| nhead (int): Number of attention heads. |
| dim_feedforward (int): FFN hidden layer size. |
| dropout (float): Dropout probability. |
| decay_rate (float): Decay coefficient lambda. |
| """ |
| def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, decay_rate=0.1): |
| super().__init__() |
| self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) |
| self.linear1 = nn.Linear(d_model, dim_feedforward) |
| self.dropout = nn.Dropout(dropout) |
| self.linear2 = nn.Linear(dim_feedforward, d_model) |
| self.norm1 = nn.LayerNorm(d_model) |
| self.norm2 = nn.LayerNorm(d_model) |
| self.decay_rate = decay_rate |
| self.phase_proj = nn.Linear(d_model, d_model) |
|
|
| def forward(self, src, t: float): |
| D_t = math.exp(-self.decay_rate * t) |
| phase = self.phase_proj(src) |
| g = D_t * torch.sin(phase) |
|
|
| attn_out, _ = self.self_attn(src, src, src) |
| src2 = self.norm1(src + g * attn_out) |
|
|
| ff = self.linear2(self.dropout(F.relu(self.linear1(src2)))) |
| return self.norm2(src2 + g * ff) |
|
|
