nanoGPT 逐行讲解——CausalSelfAttention-北京尧图网络

CausalSelfAttention第29-76行—— 最重要class CausalSelfAttention(nn.Module):def __init__(self, config):super().__init__()assert config.n_embd % config.n_head 0# QKV 投影一个线性层同时计算 Q, K, Vself.c_attn nn.Linear(config.n_embd, 3 * config.n_embd, biasconfig.bias)# 输出投影self.c_proj nn.Linear(config.n_embd, config.n_embd, biasconfig.bias)# 正则化self.attn_dropout nn.Dropout(config.dropout)self.resid_dropout nn.Dropout(config.dropout)self.n_head config.n_headself.n_embd config.n_embdself.dropout config.dropout# Flash Attention 检查PyTorch 2.0self.flash hasattr(torch.nn.functional, scaled_dot_product_attention)if not self.flash:# 因果掩码确保只能看到左边的tokenself.register_buffer(bias, torch.tril(torch.ones(config.block_size, config.block_size)).view(1, 1, config.block_size, config.block_size))def forward(self, x):B, T, C x.size() # Batch, Time, Channel# 1. 投影得到 Q, K, Vq, k, v self.c_attn(x).split(self.n_embd, dim2)# 2. 多头重塑[B, T, C] → [B, heads, T, head_dim]k k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)q q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)v v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)# 3. 计算注意力if self.flash:# 高效实现PyTorch 2.0y torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_maskNone,dropout_pself.dropout if self.training else 0,is_causalTrue)else:# 手动实现att (q k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))att att.masked_fill(self.bias[:,:,:T,:T] 0, float(-inf))att F.softmax(att, dim-1)att self.attn_dropout(att)y att v# 4. 合并多头输出y y.transpose(1, 2).contiguous().view(B, T, C)# 5. 输出投影y self.resid_dropout(self.c_proj(y))return y数据流详解输入: [B, T, C] (Batch , Sequence, Embedding Dim)│▼┌─────────────────────────────────────┐│ c_attn: Linear(C → 3*C) ││ 一次投影同时得到 Q, K, V │└─────────────────────────────────────┘│▼┌─────────────────────────────────────┐│ 重塑为多头 ││ [B, T, 3*C] → [B, heads, T, hs] ││ 其中 hs C / heads │└─────────────────────────────────────┘│▼┌─────────────────────────────────────┐│ Attention(Q, K, V) ││ softmax(QK^T / √hs) V │└─────────────────────────────────────┘│▼┌─────────────────────────────────────┐│ 合并多头输出投影 ││ [B, heads, T, hs] → [B, T, C] │└─────────────────────────────────────┘为什么叫Causal因果因为掩码矩阵 self.bias 是下三角矩阵[[1, 0, 0, 0],[1, 1, 0, 0],[1, 1, 1, 0],[1, 1, 1, 1]]这确保位置 i 只能看到位置 0~i看不到未来。

nanoGPT 逐行讲解——CausalSelfAttention

相关推荐

MLP / Feed-Forward Network

Cookie-Session、JWT与OAuth 2.0：主流登录凭证方案深度对比与选型指南

Transformer底层原理（二）Self-Attention 的数学本质

更多资讯

记忆机制：对话内与跨对话的记忆

激光共聚焦显微镜：从原理到应用的全面解析

一文读懂Prometheus+Grafana监控体系（进阶篇）

关于ELK日志分析平台搭建（进阶篇），你需要知道的一切

Pinyon Shift 原生渲染器深度解析：从PM4命令流到EDRAM、Resolve与纹理解码的完整指南

英语副词详解

今日更新

GUI自动化测试中的ScreenLayout：从硬编码坐标到声明式画面地图

LLC谐振变换器12V25A 300W开关电源设计与调试全攻略

ESP32多应用Flash数据隔离：分区表与NVS namespace实战

本周精选

面向复杂光照与角度的宠物面部识别鲁棒性增强策略

基于DRV8818与STM32的双极步进电机驱动方案实战解析

DRV8818PWPR+PIC18F45K80双极步进电机控制方案实战

本月热门

面向复杂光照与角度的宠物面部识别鲁棒性增强策略

基于DRV8818与STM32的双极步进电机驱动方案实战解析

DRV8818PWPR+PIC18F45K80双极步进电机控制方案实战

读完文章，想聊聊您的网站？