如何在Node.js中实现Transformer Encoder编码器？完整代码详解

大多数Transformer教程都基于Python和PyTorch，但如果你是一名Node.js开发者，想在熟悉的JavaScript环境中理解并实现Transformer Encoder，可参考的资料就少了很多。其实编码器的核心只有三个组件：多头自注意力、前馈网络、残差连接加层归一化，配合位置编码就能完整跑通。本文将不依赖任何深度学习框架，用纯JavaScript实现一个可运行的Transformer Encoder，并逐个模块讲解原理和维度变化。

一、编码器的整体结构与数据流转

Transformer Encoder由N个相同的层堆叠而成，每一层包含两个子层：多头自注意力层和逐位置的前馈网络层。每个子层后面都跟着残差连接和层归一化，公式表达为LayerNorm(x + Sublayer(x))。残差连接的意义在于缓解深层网络的梯度消失问题，梯度可以通过捷径直接回传到底层；层归一化则对每个样本的每个位置独立做归一化，让数值分布更稳定。

在动手写代码前，先明确张量维度约定。假设输入是一个批次大小为B、序列长度为S、词嵌入维度为D的三维张量。自注意力层内部的Q、K、V矩阵会先把D投影到D，再拆分成H个头，每个头维度为dHead = D / H。注意力计算完成后各头拼接，再经过输出投影恢复到D。前馈层则先把D扩展到4倍的隐藏维度，经过ReLU激活后再压缩回D。整个过程中序列长度S和批次B始终不变，变的只是最后一维。

我们用最朴素的嵌套数组来表示张量，不引入额外的张量库。这样做虽然性能一般，但代码透明，非常适合理解算法本身。如果后续需要性能，可以把矩阵运算替换成math.js或自己写Flat64Array版本。

二、实现缩放点积注意力与多头机制

注意力机制的核心公式是softmax(QK^T / sqrt(dK)) * V。Q和K做点积得到序列中任意两个位置的相似度分数，除以sqrt(dK)是为了防止维度较大时点积值过大，导致softmax梯度趋近于零。softmax函数按最后一维做归一化，把分数转成权重分布，再用这个权重对V加权求和，得到融合了全局信息的输出。

多头机制的思路是把D维空间切分成H个独立的子空间，每个头在自己的子空间里独立做注意力计算，最后拼接结果。不同的头可以关注不同的语义模式，比如有的头关注相邻词，有的头关注句法结构。下面是核心实现代码：

// 矩阵转置
function transpose(m) {
  const rows = m.length, cols = m[0].length;
  const t = [];
  for (let j = 0; j < cols; j++) {
    t[j] = [];
    for (let i = 0; i < rows; i++) t[j][i] = m[i][j];
  }
  return t;
}

// 矩阵乘法
function matmul(a, b) {
  const n = a.length, m = b[0].length, k = b.length;
  const c = [];
  for (let i = 0; i < n; i++) {
    c[i] = [];
    for (let j = 0; j < m; j++) {
      let sum = 0;
      for (let p = 0; p < k; p++) sum += a[i][p] * b[p][j];
      c[i][j] = sum;
    }
  }
  return c;
}

// 按最后一维做softmax
function softmax(vec) {
  const max = Math.max(...vec);
  const exps = vec.map(v => Math.exp(v - max));
  const sum = exps.reduce((a, b) => a + b, 0);
  return exps.map(e => e / sum);
}

// 缩放点积注意力
function scaledDotProductAttention(Q, K, V) {
  const dK = K[0].length;
  const scores = matmul(Q, transpose(K)).map(
    row => row.map(v => v / Math.sqrt(dK))
  );
  const weights = scores.map(row => softmax(row));
  return matmul(weights, V);
}

// 多头注意力：输入x形状为 [S, D]
function multiHeadAttention(x, Wq, Wk, Wv, Wo, numHeads) {
  const S = x.length, D = x[0].length;
  const dHead = Math.floor(D / numHeads);

  const Q = matmul(x, Wq);
  const K = matmul(x, Wk);
  const V = matmul(x, Wv);

  const heads = [];
  for (let h = 0; h < numHeads; h++) {
    const start = h * dHead;
    // 切出当前头对应的子维度
    const Qh = Q.map(r => r.slice(start, start + dHead));
    const Kh = K.map(r => r.slice(start, start + dHead));
    const Vh = V.map(r => r.slice(start, start + dHead));
    heads.push(scaledDotProductAttention(Qh, Kh, Vh));
  }

  // 拼接所有头并做输出投影
  const concat = [];
  for (let i = 0; i < S; i++) {
    concat[i] = [].concat(...heads.map(h => h[i]));
  }
  return matmul(concat, Wo);
}

注意代码中softmax之前先减去了最大值，这是数值稳定的经典技巧。如果不减最大值，当分数较大时Math.exp可能返回Infinity，归一化结果就全变成NaN了。另外权重矩阵Wq、Wk、Wv的形状都是[D, D]，初始化时用小的随机数即可，实际训练中通常用Xavier初始化。

三、位置编码、前馈网络与层归一化

自注意力本身对位置是完全不敏感的，打乱输入顺序输出也一样，所以必须注入位置信息。原论文使用正弦余弦位置编码：偶数维度用sin，奇数维度用cos，不同维度对应不同波长的正弦波。这样设计的好处是位置p+k的编码可以表示成位置p编码的线性变换，模型容易学到相对位置关系。当然，你也可以用可学习的位置嵌入，效果同样不错。

前馈网络是逐位置独立的两层全连接，中间用ReLU激活。层归一化则对每个位置的D维向量求均值和方差，再做标准化并加上可学习的缩放和平移参数。这三个模块的实现如下：

// 正弦位置编码，shape为 [S, D]
function positionalEncoding(seqLen, dModel) {
  const pe = [];
  for (let pos = 0; pos < seqLen; pos++) {
    pe[pos] = [];
    for (let i = 0; i < dModel; i++) {
      if (i % 2 === 0) {
        pe[pos][i] = Math.sin(pos / Math.pow(10000, i / dModel));
      } else {
        pe[pos][i] = Math.cos(pos / Math.pow(10000, (i - 1) / dModel));
      }
    }
  }
  return pe;
}

// 前馈网络 [S, D] -> [S, D]
function feedForward(x, W1, b1, W2, b2) {
  return x.map(row => {
    const hidden = matmul([row], W1)[0].map((v, i) => Math.max(0, v + b1[i]));
    return matmul([hidden], W2)[0].map((v, i) => v + b2[i]);
  });
}

// 层归一化 [S, D]
function layerNorm(x, gamma, beta, eps = 1e-6) {
  return x.map(row => {
    const mean = row.reduce((a, b) => a + b, 0) / row.length;
    const variance = row.reduce((a, b) => a + (b - mean) ** 2, 0) / row.length;
    return row.map((v, i) =>
      gamma[i] * ((v - mean) / Math.sqrt(variance + eps)) + beta[i]
    );
  });
}

// 残差连接
function add(a, b) {
  return a.map((row, i) => row.map((v, j) => v + b[i][j]));
}

四、组装完整的TransformerEncoder

有了上面的组件，组装就很简单了。每个编码器层的流程是：输入先经过多头注意力，输出与输入做残差相加再归一化；结果送入前馈网络，再做一次残差相加和归一化。N个层依次串联，上一层的输出就是下一层的输入。完整示例如下：

// 随机初始化权重矩阵
function randomMatrix(rows, cols, scale = 0.1) {
  return Array.from({ length: rows }, () =>
    Array.from({ length: cols }, () => (Math.random() * 2 - 1) * scale)
  );
}

function createEncoderLayer(dModel, numHeads, dFF) {
  return {
    Wq: randomMatrix(dModel, dModel),
    Wk: randomMatrix(dModel, dModel),
    Wv: randomMatrix(dModel, dModel),
    Wo: randomMatrix(dModel, dModel),
    W1: randomMatrix(dModel, dFF),
    b1: new Array(dFF).fill(0),
    W2: randomMatrix(dFF, dModel),
    b2: new Array(dModel).fill(0),
    gamma1: new Array(dModel).fill(1),
    beta1: new Array(dModel).fill(0),
    gamma2: new Array(dModel).fill(1),
    beta2: new Array(dModel).fill(0),
    numHeads
  };
}

function encoderLayerForward(x, layer) {
  const attnOut = multiHeadAttention(x, layer.Wq, layer.Wk, layer.Wv, layer.Wo, layer.numHeads);
  const norm1 = layerNorm(add(x, attnOut), layer.gamma1, layer.beta1);
  const ffOut = feedForward(norm1, layer.W1, layer.b1, layer.W2, layer.b2);
  return layerNorm(add(norm1, ffOut), layer.gamma2, layer.beta2);
}

function transformerEncoderForward(x, layers) {
  let out = x;
  for (const layer of layers) {
    out = encoderLayerForward(out, layer);
  }
  return out;
}

// 使用示例
const dModel = 64, numHeads = 8, dFF = 256, numLayers = 3;
const seqLen = 10;

const layers = Array.from({ length: numLayers },
  () => createEncoderLayer(dModel, numHeads, dFF));

// 模拟词嵌入并加上位置编码
const embeddings = randomMatrix(seqLen, dModel, 0.5);
const input = add(embeddings, positionalEncoding(seqLen, dModel));

const output = transformerEncoderForward(input, layers);
console.log('输出维度:', output.length, 'x', output[0].length); // 10 x 64

运行后输出维度保持10 x 64，与输入一致，这正是编码器的特点：序列中每个位置都输出一个等长的向量，已经融合了全序列的上下文信息。由于权重是随机初始化的，输出没有实际语义，真实场景中需要加载预训练权重，把randomMatrix替换成从JSON或二进制文件读取的参数即可。

五、性能与工程化建议

纯JavaScript的嵌套数组实现在序列较长时会明显变慢，主要瓶颈在matmul的三重循环。优化方向有几个：一是把嵌套数组改为Float64Array扁平存储，减少指针追踪开销，通常能提速数倍；二是利用Node.js的worker_threads把多个头并行计算；三是如果只是做推理，可以考虑用onnxruntime-node加载现成的ONNX模型，直接获得完整训练好的编码器能力。

另一个容易忽略的点是因果掩码。本文实现的是编码器，不做掩码；如果你要改造成解码器或做自回归任务，需要在softmax之前把未来位置的分数设置为负无穷，防止信息泄露。此外，层归一化的位置也有Pre-LN和Post-LN两种方案，原论文是Post-LN（本文实现），而GPT等现代模型多用Pre-LN，即先归一化再进子层，训练时梯度更稳定，迁移到自己的项目时值得注意这个细节差异。

Node.jsTransformer编码器修改时间：2026-09-11 06:58:45