github
/
minimind
mirror of https://github.com/jingyaogong/minimind.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Feat: Introduce Gated Attention to improve model performance

pull/579/head
wanghao 2025-12-12 21:32:40 +08:00
parent 66263101b9
commit 05616a8bd1
1 changed files with 9 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
model/model_minimind.py
View File

@ -160,6 +160,7 @@ class Attention(nn.Module):
self.k_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
self.v_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
self.o_proj = nn.Linear(args.num_attention_heads * self.head_dim, args.hidden_size, bias=False) self.o_proj = nn.Linear(args.num_attention_heads * self.head_dim, args.hidden_size, bias=False)
self.gate_proj = nn.Linear(args.hidden_size, args.num_attention_heads, bias=False)
self.attn_dropout = nn.Dropout(args.dropout) self.attn_dropout = nn.Dropout(args.dropout)
self.resid_dropout = nn.Dropout(args.dropout) self.resid_dropout = nn.Dropout(args.dropout)
self.dropout = args.dropout self.dropout = args.dropout
@ -211,8 +212,15 @@ class Attention(nn.Module):
scores = self.attn_dropout(scores) scores = self.attn_dropout(scores)
output = scores @ xv output = scores @ xv
output = output.transpose(1, 2).reshape(bsz, seq_len, -1) output = output.transpose(1, 2).contiguous()
gate_logits = self.gate_proj(x)
gate_scores = torch.sigmoid(gate_logits)
output = output * gate_scores.unsqueeze(-1)
output = output.view(bsz, seq_len, -1)
output = self.resid_dropout(self.o_proj(output)) output = self.resid_dropout(self.o_proj(output))
return output, past_kv return output, past_kv