File size: 1,733 Bytes
2db32a1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | # coding=utf-8
"""Dual-Stream Feed-Forward (DSFF).
Implements Section IX of the Wiola paper. Two parallel *dense* streams of
different widths and activations are fused by a learned per-dimension gate:
Stream A (narrow, SwiGLU): a = D_A( SiLU(G_A x) * (U_A x) )
Stream B (wide, GELU): b = D_B( GELU(U_B x) )
gate: alpha = sigmoid(W_f [a; b]) in (0,1)^d
output: alpha * a + (1 - alpha) * b
Setting ``W_f = 0`` yields ``alpha = 0.5``, reducing DSFF to a simple ensemble
average; DSFF therefore strictly generalises a two-stream ensemble.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
class DualStreamFeedForward(nn.Module):
def __init__(self, hidden_size: int, narrow_size: int, wide_size: int):
super().__init__()
# Stream A: narrow SwiGLU.
self.gate_a = nn.Linear(hidden_size, narrow_size, bias=False) # G_A
self.up_a = nn.Linear(hidden_size, narrow_size, bias=False) # U_A
self.down_a = nn.Linear(narrow_size, hidden_size, bias=False) # D_A
# Stream B: wide GELU.
self.up_b = nn.Linear(hidden_size, wide_size, bias=False) # U_B
self.down_b = nn.Linear(wide_size, hidden_size, bias=False) # D_B
# Per-dimension fusion gate from concatenated stream outputs.
self.fusion = nn.Linear(2 * hidden_size, hidden_size, bias=False) # W_f
def forward(self, x: torch.Tensor) -> torch.Tensor:
a = self.down_a(F.silu(self.gate_a(x)) * self.up_a(x))
b = self.down_b(F.gelu(self.up_b(x)))
alpha = torch.sigmoid(self.fusion(torch.cat((a, b), dim=-1)))
return alpha * a + (1.0 - alpha) * b
|