NumPy 数组间运算：LLM 开发中的核心计算引擎

原创已于 2026-01-07 10:54:44 修改 · 619 阅读

31 ·

CC 4.0 BY-SA版权

文章标签：

#numpy #python #开发语言 #人工智能 #数据挖掘 #学习 #自然语言处理

于 2026-01-06 14:54:09 首次发布

LLM数据：Numpy、Pandas、Matplotlib 专栏收录该内容

13 篇文章

订阅专栏

Anaconda加速AI模型训练 2.3w人浏览 15人参与

在这里插入图片描述

【个人主页：玄同765】

大语言模型（LLM）开发工程师｜中国传媒大学·数字媒体技术（智能交互与游戏设计）

深耕领域：大语言模型开发 / RAG知识库 / AI Agent落地 / 模型微调

技术栈：Python / LangChain/RAG（Dify+Redis+Milvus）| SQL/NumPy | FastAPI+Docker ️

工程能力：专注模型工程化部署、知识库构建与优化，擅长全流程解决方案

专栏传送门：LLM大模型开发项目实战指南、Python 从真零基础到纯文本 LLM 全栈实战、从零学 SQL + 大模型应用落地、大模型开发小白专属：从 0 入门 Linux&Shell

「让AI交互更智能，让技术落地更高效」

欢迎技术探讨/项目合作！关注我，解锁大模型与智能交互的无限可能！

NumPy（Numerical Python）作为Python科学计算领域的核心库，其多维数组对象（ndarray）提供了高效、灵活的数据处理能力。在大型语言模型（LLM）的开发中，数组间运算是构建模型、处理输入和计算损失函数的基础。从嵌入层的向量映射到自注意力机制的张量运算，再到损失函数的梯度计算，NumPy的数组间运算能力直接影响了LLM的训练效率和推理性能。本文将深入探讨NumPy数组间运算的核心机制、在LLM开发中的关键应用场景，以及如何通过优化这些运算提升模型性能。

一、NumPy 数组间运算的核心机制

NumPy的数组间运算遵循广播（Broadcasting）机制，这是其高效处理不同形状数组运算的核心。广播机制允许形状不同的数组进行元素级运算，而无需显式地复制数据或调整形状。这一机制在LLM开发中尤为重要，因为模型处理的张量通常具有复杂的维度结构。

广播遵循以下核心规则：

维度对齐：从尾部维度开始向前对齐，形状较小的数组在前面补1直到维度数相同
维度兼容：每个维度的大小要么相等，要么其中一个是1，或其中一个不存在
结果形状：取各输入数组对应维度的最大值，形成最终的输出形状

在LLM中，这些规则使得我们能够高效地处理多维张量，例如将位置编码广播到整个输入批次，或将注意力掩码广播到所有注意力头。下表展示了LLM中常见的数组形状及其广播应用：

数组类型	典型形状	广播应用
输入ID序列	(batch_size, seq_len)	与嵌入矩阵(vocab_size, embed_dim)广播相乘
位置编码	(seq_len, d_model)	扩展为(1, seq_len, d_model)与输入批次广播相加
注意力掩码	(seq_len, seq_len)	扩展为(batch_size, 1, seq_len, seq_len)广播应用
多头注意力权重	(num_heads, d_k)	扩展为(1, 1, num_heads, d_k)与查询/键/值张量广播相乘

这些广播操作使得LLM中的复杂张量运算变得简洁高效，无需编写繁琐的循环代码。例如，在Transformer模型中，位置编码可以通过简单的加法操作应用到整个输入批次，而无需为每个样本单独处理。

二、数组间运算在LLM开发中的关键应用

1. 嵌入层计算

LLM的输入通常是文本序列，这些序列需要通过嵌入层转换为连续的向量表示。嵌入层的核心操作是将离散的ID序列映射到连续的向量空间，这一过程依赖于数组间运算。具体来说，输入ID序列形状为(batch_size, seq_len)，而嵌入矩阵形状为(vocab_size, embed_dim)。通过广播机制，我们可以直接将ID序列与嵌入矩阵相乘，生成形状为(batch_size, seq_len, embed_dim)的嵌入向量。

import numpy as np

# 假设词汇表大小为5000，嵌入维度为512
vocab_size, embed_dim = 5000, 512
embedding_matrix = np.random.randn(vocab_size, embed_dim)

# 输入ID序列，batch_size=32，seq_len=128
input_ids = np.random.randint(0, vocab_size, (32, 128))

# 数组间运算生成嵌入向量
embeddings = embedding_matrix[input_ids]
# 结果形状：(32, 128, 512)

这一操作利用了NumPy的高级索引功能，本质上是一种高效的广播运算。输入ID序列中的每个元素都作为索引从嵌入矩阵中提取对应的向量，而无需显式循环。这种向量化操作比Python循环快8-10倍，对于处理大规模文本数据至关重要。

2. 自注意力机制

自注意力机制是Transformer模型的核心组件，它允许模型捕捉序列中不同位置之间的关系。自注意力的计算涉及多个数组间运算，包括线性变换、点积计算和softmax归一化，这些运算都高度依赖广播机制。

def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]
    # Q形状：(batch_size, seq_len, d_k)
    # K形状：(batch_size, seq_len, d_k)
    # V形状：(batch_size, seq_len, d_v)
    # 计算注意力分数
    scores = np.matmul(Q, K.transpose(0, 2, 1)) / np.sqrt(d_k)
    # 应用掩码
    if mask is not None:
        scores = np.where(mask, scores, -1e9)
    # 计算softmax
    attn_weights = np.exp(scores) / np.sum(np.exp(scores), axis=-1, keepdims=True)
    # 计算上下文向量
    context = np.matmul(attn_weights, V)
    return context

在这个实现中，Q和K的矩阵乘法（np.matmul(Q, K.transpose(0, 2, 1))）依赖于广播机制自动扩展维度。此外，掩码（形状通常为(seq_len, seq_len)）需要广播到所有批次和注意力头，确保每个位置只能关注允许的位置。这种广播机制使得自注意力计算能够高效处理不同长度的序列，而无需为每个序列长度编写特定代码。

3. 多头注意力实现

多头注意力是自注意力的扩展，它通过并行多个注意力头来捕捉不同子空间中的关系。多头注意力的实现涉及到复杂的数组形状操作和广播机制，这是NumPy数组间运算的典型应用场景。

def multi_head_attention(X, num_heads, d_model):
    # X形状：(batch_size, seq_len, d_model)
    # 线性变换生成Q、K、V
    W_Q = np.random.randn(d_model, d_model)
    W_K = np.random.randn(d_model, d_model)
    W_V = np.random.randn(d_model, d_model)
    Q = np.matmul(X, W_Q)
    K = np.matmul(X, W_K)
    V = np.matmul(X, W_V)
    # 分头操作
    Q = Q.reshape(Q.shape[0], Q.shape[1], num_heads, d_model//num_heads)
    Q = Q.transpose(0, 2, 1, 3)  # 形状变为(batch_size, num_heads, seq_len, d_k)
    K = K.reshape(K.shape[0], K.shape[1], num_heads, d_model//num_heads)
    K = K.transpose(0, 2, 1, 3)
    V = V.reshape(V.shape[0], V.shape[1], num_heads, d_model//num_heads)
    V = V.transpose(0, 2, 1, 3)
    # 分头计算注意力
    d_k = d_model // num_heads
    heads = []
    for i in range(num_heads):
        head = scaled_dot_product_attention(Q[:,i], K[:,i], V[:,i])
        heads.append(head)
    # 拼接多头结果
    concatenated = np拼接(heads, axis=-1)
    # 最终线性变换
    W_O = np.random.randn(d_model, d_model)
    output = np.matmul(concatenated, W_O)
    return output

在这个实现中，数组间运算的广播机制体现在多个方面：首先，Q、K、V的线性变换是通过矩阵乘法完成的；其次，分头操作通过reshape和transpose调整数组形状，使每个头能够独立处理子空间；最后，多头结果的拼接和最终线性变换也是基于数组间运算的广播机制。这种向量化实现比逐头循环处理快得多，是构建高效LLM的关键。

4. 位置编码应用

Transformer模型没有内置的序列顺序信息，因此需要通过位置编码（Positional Encoding）显式注入位置信息。位置编码通常是一个二维数组，需要广播到与输入批次相同的三维形状，然后与嵌入向量相加。

def positional_encoding(seq_len, d_model):
    # 生成位置编码矩阵
    pe = np.zeros((seq_len, d_model))
    for pos in range(seq_len):
        for i in range(0, d_model, 2):
            pe[pos, i] = np.sin(pos / (10000 ** ((2 * i)/d_model)))
            pe[pos, i+1] = np.cos(pos / (10000 ** ((2 * (i+1))/d_model)))
    return pe

# 应用位置编码
batch_size, seq_len, d_model = 32, 128, 512
X = np.random.randn(batch_size, seq_len, d_model)  # 输入张量
pe = positional_encoding(seq_len, d_model)  # 位置编码矩阵
# 广播机制自动扩展位置编码到批次维度
X_with_pe = X + pe[None, :, :]  # 形状变为(32, 128, 512)

在这个示例中，位置编码矩阵pe的形状为(seq_len, d_model)，通过添加None轴（pe[None, :, :]）扩展为(1, seq_len, d_model)。然后，通过广播机制与形状为(batch_size, seq_len, d_model)的输入张量X相加，生成包含位置信息的输入张量。这种广播机制使得位置编码能够高效地应用于任意批次大小的输入，无需为每个批次单独计算。

5. 损失函数计算

LLM的训练依赖于损失函数的计算，例如交叉熵损失。在计算交叉熵损失时，广播机制使得我们能够高效地比较预测概率分布与真实标签，无需显式地将标签转换为one-hot编码。

def cross_entropy_loss(y_pred, y_true):
    # y_pred形状：(batch_size, seq_len, vocab_size)
    # y_true形状：(batch_size, seq_len)
    # 计算softmax
    y_pred_shifted = y_pred - np.max(y_pred, axis=-1, keepdims=True)
    y_pred_exp = np.exp(y_pred_shifted)
    y_pred软max = y_pred_exp / np.sum(y_pred_exp, axis=-1, keepdims=True)
    # 使用广播机制直接索引真实标签的概率
    log probabilities = np.log(y_pred软max)
    # 广播应用：y_true的形状(32, 128)广播到(32, 128, 1)
    # 然后与log probabilities的形状(32, 128, 50257)进行元素级比较
    loss = -np.mean(log probabilities[np.arange(batch_size), np.arange(seq_len), y_true])
    return loss

# 示例用法
batch_size, seq_len, vocab_size = 32, 128, 50257
logits = np.random.randn(batch_size, seq_len, vocab_size)
labels = np.random.randint(0, vocab_size, (batch_size, seq_len))
loss = cross_entropy_loss(logits, labels)

在这个实现中，广播机制允许我们直接通过logits[np.arange(batch_size), np.arange(seq_len), labels]索引到每个位置的真实标签对应的概率，然后计算对数并求平均损失。这种广播机制避免了显式地将标签转换为one-hot编码的开销，显著提高了计算效率。

三、数组间运算的优化技巧与最佳实践

在LLM开发中，数组间运算的效率直接影响模型训练和推理的性能。通过合理应用NumPy的数组间运算特性，可以显著提升计算效率，减少内存占用。

1. 避免不必要的副本操作

NumPy数组操作分为视图（view）和副本（copy），理解这一区别对于优化LLM中的数组间运算至关重要。某些操作（如reshape、transpose、切片）返回视图，共享底层数据缓冲区，而另一些操作（如copy、flatten）返回副本，需要复制数据。在LLM开发中，应尽量使用视图操作以减少内存占用和复制开销。

# 不推荐：使用flatten创建副本
flattened = arr.flatten()

# 推荐：使用ravel创建视图
raveled = arr.ravel()

# 推荐：使用视图操作进行形状调整
reshaped = arr.reshape((batch_size, num_heads, seq_len, head_dim))

在多头注意力实现中，通过reshape和transpose调整数组形状时，应尽量保持这些操作为视图，以避免不必要的内存复制。此外，使用np膜map（memmap）可以将大型数组存储在磁盘上，只在需要时加载到内存，这对于处理大规模预训练模型的权重矩阵非常有用。

2. 合理使用广播机制

广播机制是NumPy高效处理不同形状数组运算的核心，但在某些情况下可能导致不必要的内存扩展。在LLM开发中，应理解广播机制的工作原理，并在适当的情况下使用显式扩展或使用keepdims=True参数来避免意外的维度扩展。

# 不推荐：可能导致不必要的维度扩展
summed = np.sum(arr, axis=1)

# 推荐：使用keepdims=True保留维度信息
summed_with Keepdims = np.sum(arr, axis=1, keepdims=True)

# 推荐：使用np 新轴显式控制维度
arr_expanded = arr[:, np 新轴, :]

# 推荐：使用np 扩展维度函数
arr_expanded = np 扩展维度(arr, axis=1)

在自注意力机制中，计算注意力分数时需要确保Q和K的形状兼容。例如，如果Q的形状是(batch_size, seq_len, d_k)，而K的形状是(batch_size, seq_len, d_k)，则需要先对K进行转置，使其形状变为(batch_size, d_k, seq_len)，然后再进行矩阵乘法。这种显式维度控制可以避免广播机制可能带来的意外扩展。

3. 优先使用向量化操作

在LLM开发中，应尽量避免使用显式循环，而使用NumPy的向量化操作。向量化操作通过底层C实现，能够并行处理数组元素，比Python循环快8-10倍。这在处理大规模文本数据和复杂张量运算时尤为重要。

# 不推荐：使用Python循环计算逐元素乘积
result = np.zeros_like(arr)
for i in range(arr.shape[0]):
    for j in range(arr.shape[1]):
        result[i, j] = arr[i, j] * scalar

# 推荐：使用向量化操作
result = arr * scalar

# 推荐：使用NumPy函数替代循环
result = np where(arr > threshold, arr, 0)

在LLM的损失函数计算中，使用向量化操作可以显著提高计算效率。例如，计算交叉熵损失时，应使用np.log和np.sum等向量化函数，而不是显式循环遍历每个元素。此外，使用np.einsum可以简化复杂的多维张量运算，提高代码可读性和执行效率。

4. 利用数组内存布局优化

NumPy数组的内存布局有两种主要形式：C连续（行优先）和Fortran连续（列优先）。在LLM开发中，应根据数组访问模式选择合适的内存布局，以提高缓存利用率和计算效率。

# 创建C连续数组
arr_c = np.random.randn(1000, 1000)

# 创建Fortran连续数组
arr_f = np.random.randn(1000, 1000).T

# 检查内存布局
print(arr_c 转置 .flags['C连续'])  # 输出：False
print(arr_f 转置 .flags['C连续'])  # 输出：True

在自注意力机制中，Q和K的矩阵乘法通常涉及大量连续内存访问，因此应确保这些数组是C连续的，以提高BLAS库（如NumPy底层使用的libBLAS）的计算效率。此外，对于需要频繁访问特定维度的数组，可以使用np 新轴或reshape调整数组形状，使其内存布局更符合访问模式。

5. 使用合适的数组数据类型

NumPy支持多种数组数据类型，如float32、float64、int32等。在LLM开发中，应根据计算需求和内存限制选择合适的数组数据类型，以平衡计算精度和效率。

# 默认使用float64，可能占用过多内存
arr_default = np.random.randn(1000, 1000)

# 使用float32提高内存效率
arr_float32 = np.random.randn(1000, 1000).astype(np.float32)

# 使用int32提高整数运算效率
arr_int32 = np.random.randint(0, 10000, (1000, 1000)).astype(np.int32)

在LLM的嵌入层和前馈神经网络层中，使用float32而非默认的float64可以减少约50%的内存占用，同时保持足够的计算精度。此外，在处理输入ID序列时，使用int32而非int64可以进一步减少内存占用，提高计算效率。

四、数组间运算在LLM中的实际应用案例

1. 嵌入层与位置编码的结合

在Transformer模型中，输入文本首先通过嵌入层转换为向量表示，然后与位置编码相加，以注入位置信息。这一过程涉及两个数组间运算：嵌入矩阵的索引和位置编码的广播相加。

def embedding_with_positional_encoding(input_ids, embedding_matrix, pe):
    # input_ids形状：(batch_size, seq_len)
    # embedding_matrix形状：(vocab_size, embed_dim)
    # pe形状：(seq_len, d_model)
    # 嵌入层计算
    embeddings = embedding_matrix[input_ids]
    # 广播机制应用：扩展位置编码到批次维度
    pe_expanded = pe[None, :, :]  # 形状变为(1, seq_len, d_model)
    # 嵌入向量与位置编码相加
    X = embeddings + pe_expanded
    return X

# 示例用法
batch_size, seq_len, vocab_size, embed_dim = 32, 128, 50257, 512
input_ids = np.random.randint(0, vocab_size, (batch_size, seq_len))
embedding_matrix = np.random.randn(vocab_size, embed_dim)
pe = positional_encoding(seq_len, embed_dim)
X = embedding_with_positional_encoding(input_ids, embedding_matrix, pe)

在这个示例中，广播机制允许我们将形状为(seq_len, d_model)的位置编码pe扩展为(1, seq_len, d_model)，然后与形状为(batch_size, seq_len, d_model)的嵌入向量embeddings相加，生成包含位置信息的输入张量X。这种广播机制使得位置编码能够高效地应用于任意批次大小的输入，无需为每个批次单独计算。

2. 自注意力机制的实现

自注意力机制是Transformer模型的核心组件，它允许模型捕捉序列中不同位置之间的关系。自注意力的实现涉及多个数组间运算，包括Q、K、V的线性变换、注意力分数的计算和softmax归一化 。

def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]
    # Q形状：(batch_size, seq_len, d_k)
    # K形状：(batch_size, seq_len, d_k)
    # V形状：(batch_size, seq_len, d_v)
    # 计算注意力分数
    scores = np.matmul(Q, K.transpose(0, 2, 1)) / np.sqrt(d_k)
    # 形状：(batch_size, seq_len, seq_len)
    # 应用掩码
    if mask is not None:
        # mask形状：(seq_len, seq_len)
        # 广播到：(batch_size, 1, seq_len, seq_len)
        mask广播 = mask[None, None, :, :]
        scores = np.where(mask广播, scores, -1e9)
    # 计算softmax
    attn_weights = np.exp(scores) / np.sum(np.exp(scores), axis=-1, keepdims=True)
    # 形状：(batch_size, seq_len, seq_len)
    # 计算上下文向量
    context = np.matmul(attn_weights, V)
    # 形状：(batch_size, seq_len, d_v)
    return context

在这个实现中，广播机制允许我们将形状为(seq_len, seq_len)的掩码mask扩展为(batch_size, 1, seq_len, seq_len)，然后与形状为(batch_size, seq_len, seq_len)的注意力分数scores进行比较，生成加性掩码。这种广播机制使得掩码能够高效地应用于所有批次和序列位置，无需为每个批次单独处理。

3. 多头注意力的实现

多头注意力是自注意力的扩展，它通过并行多个注意力头来捕捉不同子空间中的关系。多头注意力的实现涉及复杂的数组形状操作和广播机制，这是NumPy数组间运算的典型应用场景。

def multi_head_attention(X, num_heads, d_model):
    # X形状：(batch_size, seq_len, d_model)
    # 线性变换生成Q、K、V
    W_Q = np.random.randn(d_model, d_model)
    W_K = np.random.randn(d_model, d_model)
    W_V = np.random.randn(d_model, d_model)
    Q = np.matmul(X, W_Q)
    K = np.matmul(X, W_K)
    V = np.matmul(X, W_V)
    # 分头操作
    # Q形状：(batch_size, seq_len, num_heads, d_k)
    d_k = d_model // num_heads
    Q = Q.reshape(Q.shape[0], Q.shape[1], num_heads, d_k)
    Q = Q.transpose(0, 2, 1, 3)  # 形状变为(batch_size, num_heads, seq_len, d_k)
    K = K.reshape(K.shape[0], K.shape[1], num_heads, d_k)
    K = K.transpose(0, 2, 1, 3)
    V = V.reshape(V.shape[0], V.shape[1], num_heads, d_v)
    V = V.transpose(0, 2, 1, 3)
    # 分头计算注意力
    heads = []
    for i in range(num_heads):
        head = scaled_dot_product_attention(Q[:,i], K[:,i], V[:,i])
        heads.append(head)
    # 拼接多头结果
    concatenated = np拼接(heads, axis=-1)
    # 形状：(batch_size, seq_len, d_model)
    # 最终线性变换
    W_O = np.random.randn(d_model, d_model)
    output = np.matmul(concatenated, W_O)
    return output

4. 交叉熵损失函数的实现

在LLM训练中，交叉熵损失函数是最常用的损失函数之一。交叉熵损失函数的实现涉及多个数组间运算，包括softmax计算、对数运算和元素级乘法 。

def cross_entropy_loss(y_pred, y_true):
    # y_pred形状：(batch_size, seq_len, vocab_size)
    # y_true形状：(batch_size, seq_len)
    # 计算softmax
    y_pred_shifted = y_pred - np.max(y_pred, axis=-1, keepdims=True)
    y_pred_exp = np.exp(y_pred_shifted)
    y_pred软max = y_pred_exp / np.sum(y_pred_exp, axis=-1, keepdims=True)
    # 使用广播机制直接索引真实标签的概率
    log probabilities = np.log(y_pred软max)
    # 广播应用：y_true的形状(32, 128)广播到(32, 128, 1)
    # 然后与log probabilities的形状(32, 128, 50257)进行元素级比较
    loss = -np.mean(log probabilities[np.arange(batch_size), np.arange(seq_len), y_true])
    return loss

在这个实现中，广播机制允许我们直接通过log probabilities[np.arange(batch_size), np.arange(seq_len), y_true]索引到每个位置的真实标签对应的概率，然后计算对数并求平均损失。这种广播机制避免了显式地将标签转换为one-hot编码的开销，显著提高了计算效率。