import torch from torch import nn import math from transformers.models.whisper.modeling_whisper import WhisperAttention from transformers.activations import ACT2FN class CustomLinear(nn.Linear): def __init__(self, *args, init_eye_val=0.0, fddt_init=None, init_fun=None, **kwargs): super().__init__(*args, **kwargs) self.init_eye_val = init_eye_val self.fddt_init = fddt_init self.init_fun = init_fun self.reset_parameters() # Ensure consistent init on creation def reset_parameters(self) -> None: with torch.no_grad(): # Apply custom init function if provided if hasattr(self,"init_fun") and self.init_fun is not None: self.init_fun(self) return # Default initialization nn.init.xavier_uniform_(self.weight) if self.bias is not None: nn.init.zeros_(self.bias) if hasattr(self, "fddt_init"): # FDDT-specific inits if self.fddt_init == 'non-disturbing': # Make weight an identity matrix (if possible) if self.weight.shape[0] == self.weight.shape[1]: self.weight.copy_(torch.eye(self.weight.shape[0], device=self.weight.device)) else: # Not square — fill first min(n, m) diagonals eye = torch.zeros_like(self.weight) n = min(self.weight.shape) eye[:n, :n] = torch.eye(n, device=self.weight.device) self.weight.copy_(eye) elif self.fddt_init == 'suppressive': if self.weight.shape[0] == self.weight.shape[1]: self.weight.copy_(self.init_eye_val * torch.eye(self.weight.shape[0], device=self.weight.device)) else: eye = torch.zeros_like(self.weight) n = min(self.weight.shape) eye[:n, :n] = self.init_eye_val * torch.eye(n, device=self.weight.device) self.weight.copy_(eye) class CustomDiagonalLinear(nn.Module): def __init__(self, d_model, bias=True, init_eye_val=0.0, fddt_init=None): super().__init__() self.init_eye_val = init_eye_val self.weight = nn.Parameter(torch.full((d_model,), init_eye_val)) self.bias = nn.Parameter(torch.zeros(d_model)) if bias else None self.fddt_init = fddt_init self.reset_parameters() def reset_parameters(self): with torch.no_grad(): # random init fan = self.weight.size(0) bound = math.sqrt(3.0 / fan) self.weight.uniform_(-bound, bound) if self.bias is not None: self.bias.zero_() # custom modes if self.fddt_init == 'non-disturbing': self.weight.fill_(1.0) elif self.fddt_init == 'suppressive': self.weight.fill_(self.init_eye_val) def forward(self, input): out = input * self.weight if self.bias is not None: out += self.bias return out class Gate(nn.Module): def __init__(self, items, init_val=0.0): super().__init__() self.init_val = init_val self.gate = nn.Parameter(torch.full((items,), init_val)) self.reset_parameters() def forward(self, orig_seq, new_seq): gate_act = torch.nn.functional.tanh(self.gate) output = orig_seq + gate_act * new_seq return output def reset_parameters(self): with torch.no_grad(): self.gate.fill_(self.init_val) def propagate_first_half_embeds_init(module): # Zero out all weights initially # module.weight.data.zero_() torch.nn.init.xavier_uniform_(module.weight, gain=1e-1) # Create identity mapping for first half of input (cross_attn_output) # Input: [cross_attn_output, q_orig] -> map cross_attn_output to first embed_dim outputs module.weight.data[:module.weight.shape[1] // 2, :module.weight.shape[1] // 2] += torch.eye( module.weight.shape[1] // 2) # Zero bias module.bias.data.zero_() def propage_first_embeds_to_match_output_dim_init(module): # module.weight.data.zero_() torch.nn.init.xavier_uniform_(module.weight, gain=1e-1) # Create identity mapping from first embed_dim inputs to output module.weight.data[:, :module.weight.shape[0]] += torch.eye(module.weight.shape[0]) # Zero bias for second linear module.bias.data.zero_() # Cross attention block that can easily learn to ignore cross attention initially class CrossAttentionEnrollBlock(nn.Module): def __init__(self, config): super().__init__() self.embed_dim = config.d_model self.ffn_dim = config.encoder_ffn_dim self.cross_attn = WhisperAttention( embed_dim=self.embed_dim, num_heads=config.encoder_attention_heads, dropout=config.attention_dropout, config=config, ) # Layer normalization (pre-norm style) # self.norm_attn = nn.LayerNorm(self.embed_dim, eps=layer_norm_eps) self.cross_gate = Gate(1,init_val=.0) # Feed-forward network that maps concat space back to single channel self.ffn = nn.Sequential( CustomLinear(self.embed_dim * 2, self.ffn_dim, init_fun=propagate_first_half_embeds_init), ACT2FN[config.activation_function], nn.Dropout(config.dropout if hasattr(config, 'dropout') else 0.1), CustomLinear(self.ffn_dim, self.embed_dim, init_fun=propage_first_embeds_to_match_output_dim_init), nn.Dropout(config.dropout if hasattr(config, 'dropout') else 0.1) ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: """ Args: hidden_states: (B, 2, T, F) - batch, channels, time, features Returns: Updated hidden states of same shape """ q = hidden_states[:, 0] # (B, T, F) kv = hidden_states[:, 1] # (B, T, F) # Cross-attention attn_output = self.cross_attn( hidden_states=q, key_value_states=kv, output_attentions=False )[0] # Concatenate attention output with original normalized query q_concat = torch.cat([attn_output, q], dim=-1) # (B, T, 2*F) # Feed-forward processing (no normalization to preserve initialization) updated_q = self.ffn(q_concat) # (B, T, F) q_out = self.cross_gate(q, updated_q) # Return stacked result (only query channel is updated) return torch.stack([q_out, kv], dim=1) class SpeakerCommunicationBlock(nn.Module): def __init__(self, config): super().__init__() self.streams = 2 self.config = config self.cae = CrossAttentionEnrollBlock(config) def forward(self, x): # x: (B, T, F) B, T, F = x.shape S = self.streams # Reshape to (B//S, S, T, F) x_reshaped = x.view(B//S, S, T, F) # Call the selected method out = self.cae(x_reshaped) # Reshape back (B, T, F) out_merged = out.view(B, T, F) return out_merged if __name__ == "__main__": model1 = CustomLinear(16 * 2, 64, init_fun=propagate_first_half_embeds_init) model2 = CustomLinear(64, 16, init_fun=propage_first_embeds_to_match_output_dim_init) input1 = torch.ones(16, 16) input2 = torch.zeros(16, 16) input = torch.concat((input1, input2), dim=-1) output = model2(model1(input)) print(f"Mean err: {(input1-output).mean()}") model_1 = CustomDiagonalLinear(4, bias=False, fddt_init='suppressive', init_eye_val=0.1) model_2 = CustomDiagonalLinear(4, bias=False, fddt_init='suppressive', init_eye_val=0.1) model_3 = CustomDiagonalLinear(4, bias=False, fddt_init='suppressive', init_eye_val=0.1) model_4 = CustomDiagonalLinear(4, bias=False, fddt_init='suppressive', init_eye_val=0.1) model = nn.Sequential(model_1, model_2, model_3, model_4) opt = torch.optim.Adam(model.parameters(), lr=0.01) model_1.reset_parameters() x = torch.ones(2, 4) y = torch.ones(2, 4) for i in range(100): opt.zero_grad() loss = ((model(x) - y) ** 2).mean() loss.backward() opt.step() print(f"Step {i}: mean weight {model_1.weight.mean().item():.4f}")