P3712.第2题-大模型Attention模块开发

1000ms

Tried: 3066

Accepted: 1018

Difficulty: 4

所属公司 : 华为

算法与标签>模拟

题目内容

已知大模型常用的 Attention 模块定义如下：

$Y = \text{softmax}\left(\frac{QK^T}{\sqrt{h}}\right)V$

此处考虑二维情况，其中

$Q, K, V = XW_1, XW_2, XW_3 \in \mathbb{R}^{n \times h}, \quad X \in \mathbb{R}^{n \times m}, \quad W_1, W_2, W_3 \in \mathbb{R}^{m \times h}$

注意：

为简便起见，所有输入初始化为全 $1$ 矩阵，所有权重矩阵初始化为上三角全 $1$ 矩阵。
对任意矩阵 ( M ) 的 $softmax$ 计算简化为：

$\text{softmax}(M)_{ij} = \frac{M_{ij}}{M_i}, \quad M_i = \sum_j M_{ij}$

输入描述

输入为维度参数 $n, m$ 和 $h$ ，参数间使用空格隔开，均为小于 $100$ 的正整数

输出描述

输出为结果矩阵 $Y \in \mathbb{R}^{n \times h}$ 的所有元素之和，例如 $15$ ，输出在四舍五入后保留整数

样例1

输入

3 3 3

输出

说明

$X = \begin{pmatrix} 1 & 1 & 1 \\ 1 & 1 & 1 \\ 1 & 1 & 1 \end{pmatrix}, \quad W_1, W_2, W_3 = \begin{pmatrix} 1 & 1 & 1 \\ 0 & 1 & 1 \\ 0 & 0 & 1 \end{pmatrix}$

$Q, K, V = \begin{pmatrix} 1 & 2 & 3 \\ 1 & 2 & 3 \\ 1 & 2 & 3 \end{pmatrix}, \quad Y = \begin{pmatrix} 1 & 2 & 3 \\ 1 & 2 & 3 \\ 1 & 2 & 3 \end{pmatrix}$

输出为： $18$

样例2

输入

2 3 1

输出

说明

$X = \begin{pmatrix} 1 & 1 & 1 \\ 1 & 1 & 1 \end{pmatrix}, \quad W_1, W_2, W_3 = \begin{pmatrix} 1 \\ 0 \\ 0 \end{pmatrix}$

$Q, K, V = \begin{pmatrix} 1 \\ 1 \end{pmatrix}, \quad Y = \begin{pmatrix} 1 \\ 1 \end{pmatrix}$

输出为： $2$

提示

输入参数不包含 $0$ ，为正整数

输入

预期输出（选填）

解题思路

P3712.第2题-大模型Attention模块开发

题目内容

注意：

输入描述

输出描述

样例1

样例2

提示

Status

Development

Support

About