### MatMul 函数实现 矩阵乘法（`matmul`）是深度学习中最基本的操作之一，在 PyTorch 中可以通过 `torch.matmul` 或者 `@` 运算符来实现。以下是手动实现的一个简单版本： ```python import torch def matmul(a, b): """ 手动实现矩阵乘法 a @ b """ assert a.size(-1) == b.size(-2), "Matrix dimensions must agree" result = torch.zeros(*a.shape[:-1], b.shape[-1]) for i in range(a.size(-2)): for j in range(b.size(-1)): result[..., i, j] = (a[..., i, :] * b[..., :, j]).sum(dim=-1) return result ``` 上述代码实现了两个张量之间的逐元素相乘并累加操作[^1]。 --- ### Softmax 函数实现 Softmax 是一种常用的激活函数，用于将输入转换为概率分布形式。其定义为： \[ \text{Softmax}(x)_i = \frac{\exp(x_i)}{\sum_j \exp(x_j)} \] 下面是基于 Python 和 PyTorch 的实现方法： ```python def softmax(x, dim=None): """ 手动实现 Softmax 函数 """ exp_x = torch.exp(x - torch.max(x, dim=dim, keepdim=True)[0]) # 防止溢出 sum_exp_x = torch.sum(exp_x, dim=dim, keepdim=True) return exp_x / sum_exp_x ``` 此实现通过减去最大值防止指数运算过程中可能出现的数值不稳定问题[^2]。 --- ### RMSNorm 函数实现 RMSNorm 是 LayerNorm 的简化版，它不涉及均值计算和偏置参数 \( \beta \)，仅依赖于输入特征平方的平均值进行归一化处理。具体公式如下： \[ y = \frac{x}{\sqrt{\operatorname{Mean}(x^2) + \epsilon}} * \gamma \] 以下是基于 PyTorch 的实现代码： ```python class RMSNorm(torch.nn.Module): def __init__(self, dim, eps=1e-6): super(RMSNorm, self).__init__() self.dim = dim self.eps = eps self.weight = torch.nn.Parameter(torch.ones(dim)) def forward(self, x): variance = torch.mean(x ** 2, dim=-1, keepdim=True) # 均方根部分 normalized_x = x * torch.rsqrt(variance + self.eps) # 归一化 return self.weight * normalized_x # 缩放权重 ``` 该实现遵循了原始论文中的描述，即去除均值计算过程，并利用均方根替代标准差计算[^3]。 --- ### 总结 以上分别展示了 `matmul`、`softmax` 和 `rmsnorm` 的实现方式及其背后的原理。这些基础组件构成了现代神经网络的核心模块，尤其在 Transformer 架构中具有重要作用。