BART由Facebook在2019年提出,全称Bidirectional and Auto-Regressive Transformers,它把BERT式的双向编码器与GPT式的自回归解码器串联起来,用一套去噪预训练目标让模型学会从被破坏的文本中还原原文。这个设计让BART在文本摘要、机器翻译、文本生成等任务上表现出色。虽然深度学习领域几乎被Python垄断,但得益于TensorFlow.js等框架的成熟,Node.js同样可以承担模型训练与推理的工作,尤其适合前后端统一技术栈的团队。本文详细介绍如何用Node.js实现BART的去噪预训练流程。

一、BART架构与去噪预训练的原理
BART的整体结构是一个标准的Seq2Seq Transformer。输入文本经过词嵌入和位置编码后,进入若干层双向编码器,编码器输出的隐状态再作为解码器的交叉注意力来源,解码器以自回归方式逐词生成目标序列。与BERT只能做理解任务、GPT只能做生成任务不同,BART因为这种混合结构,既能理解又能生成。
去噪预训练的思路很直观:拿一段正常文本,人为地破坏它,比如遮盖一些词、打乱句子顺序、删除部分片段,然后把破坏后的文本喂给编码器,让解码器输出原始的干净文本。模型为了完成任务,必须学会语法结构、语义关联和上下文推理,这就是预训练信号的来源。可以说噪声设计的好坏直接决定了预训练质量。
BART原论文中比较了几种噪声方案:词遮盖借鉴了BERT的MLM,但效果一般;句子乱序和文档旋转让模型学习全局结构;删除词让模型判断缺失内容;效果最好的是文本填充,即用一个遮盖符替换连续的多个词,模型需要推断被遮盖部分有多长以及内容是什么,任务难度更高,学到的表示也更丰富。
二、在Node.js中准备数据与实现噪声注入
第一步是语料准备。假设我们有一个纯文本语料库,每行一个文档。在Node.js中可以很方便地用流式读取处理大文件,避免一次性把所有语料加载进内存。读取后需要做分词处理,可以直接按空格或标点做简单分词,也可以借助一些JavaScript分词库。
噪声注入是去噪预训练的灵魂,下面给出文本填充和句子乱序的Node.js实现。文本填充的实现思路是:随机选取若干起始位置,将连续的多个词替换为遮盖符<mask>;句子乱序则是按句子切分后打乱顺序。
// 文本填充:将连续token替换为遮盖符
function textInfilling(tokens, maskRatio = 0.3, lambda = 3) {
const output = [];
let i = 0;
const totalLen = tokens.length;
const maskBudget = Math.floor(totalLen * maskRatio);
let masked = 0;
while (i < tokens.length) {
// 按泊松分布决定遮盖长度,均值lambda
if (masked < maskBudget && Math.random() < 0.5) {
const spanLen = poissonSample(lambda);
// 遮盖跨度不超过剩余预算
const actualLen = Math.min(spanLen, maskBudget - masked);
output.push('<mask>');
i += actualLen;
masked += actualLen;
} else {
output.push(tokens[i]);
i++;
}
}
return output;
}
// 句子乱序:切句后随机打乱
function sentencePermutation(text) {
const sentences = text.split(/(?<=[。!?.!?])/);
for (let i = sentences.length - 1; i > 0; i--) {
const j = Math.floor(Math.random() * (i + 1));
[sentences[i], sentences[j]] = [sentences[j], sentences[i]];
}
return sentences.join('');
}
// 泊松分布采样
function poissonSample(lambda) {
let L = Math.exp(-lambda), k = 0, p = 1;
do {
k++;
p *= Math.random();
} while (p > L);
return k - 1;
}上面代码中,泊松采样的参数lambda控制遮盖跨度的平均长度,论文推荐值为3,即平均每次遮盖三个词。maskRatio控制总遮盖比例,通常在0.3到0.5之间取值。实际训练时建议将多种噪声混合使用,比如30%的样本用文本填充,30%用句子乱序,剩余用词删除,这样模型能同时学到局部和全局的语言知识。
三、用TensorFlow.js搭建编码器解码器并训练
有了带噪声的样本后,接下来搭建模型。TensorFlow.js在Node.js环境下通过@tensorflow/tfjs-node提供原生绑定,性能远高于纯CPU的JavaScript实现。首先安装依赖:
npm install @tensorflow/tfjs-node
模型搭建分三部分:共享的词嵌入层、编码器和解码器。为简化演示,下面用函数式API构建一个小型BART模型,编码器和解码器各两层,隐藏维度256,注意力头数8。
const tf = require('@tensorflow/tfjs-node');
const VOCAB_SIZE = 30000;
const D_MODEL = 256;
const NUM_HEADS = 8;
const FF_DIM = 1024;
const MAX_LEN = 128;
function buildBart() {
// 编码器输入:加噪后的token序列
const encInput = tf.input({ shape: [MAX_LEN] });
// 解码器输入:目标序列右移一位(teacher forcing)
const decInput = tf.input({ shape: [MAX_LEN] });
const embedding = tf.layers.embedding({
inputDim: VOCAB_SIZE, outputDim: D_MODEL
});
const encEmb = embedding.apply(encInput);
const decEmb = embedding.apply(decInput);
// 编码器:两层双向Transformer块
let x = encEmb;
for (let i = 0; i < 2; i++) {
const attn = tf.layers.multiHeadAttention({
headSize: D_MODEL, numHeads: NUM_HEADS
}).apply(x);
x = tf.layers.add().apply([x, attn]);
x = tf.layers.layerNormalization().apply(x);
let ff = tf.layers.dense({ units: FF_DIM, activation: 'gelu' }).apply(x);
ff = tf.layers.dense({ units: D_MODEL }).apply(ff);
x = tf.layers.add().apply([x, ff]);
x = tf.layers.layerNormalization().apply(x);
}
// 解码器:自回归注意力 + 交叉注意力
let y = decEmb;
for (let i = 0; i < 2; i++) {
const selfAttn = tf.layers.multiHeadAttention({
headSize: D_MODEL, numHeads: NUM_HEADS
}).apply(y);
y = tf.layers.add().apply([y, selfAttn]);
y = tf.layers.layerNormalization().apply(y);
const crossAttn = tf.layers.multiHeadAttention({
headSize: D_MODEL, numHeads: NUM_HEADS
}).apply([y, x]); // 交叉注意力:以编码器输出为上下文
y = tf.layers.add().apply([y, crossAttn]);
y = tf.layers.layerNormalization().apply(y);
}
const logits = tf.layers.dense({ units: VOCAB_SIZE }).apply(y);
const model = tf.model({ inputs: [encInput, decInput], outputs: logits });
model.compile({
optimizer: tf.train.adam(3e-4),
loss: 'sparseCategoricalCrossentropy'
});
return model;
}训练循环的核心是构造批次数据:每个批次包含加噪后的编码器输入和右移一位的解码器输入标签。解码器输入是目标序列前面拼一个起始符,而标签就是原始序列,这样模型在每个位置预测下一个词,这正是自回归训练的标准做法。
async function train(model, corpus, epochs = 10, batchSize = 32) {
for (let epoch = 0; epoch < epochs; epoch++) {
for (let step = 0; step + batchSize < corpus.length; step += batchSize) {
const batch = corpus.slice(step, step + batchSize);
const encInputs = [], decInputs = [], labels = [];
for (const text of batch) {
const clean = tokenize(text);
const noisy = applyNoise(clean); // 混合噪声
encInputs.push(pad(noisy));
// 解码器输入:起始符 + 干净序列去掉最后一个词
decInputs.push(pad([BOS_ID, ...clean.slice(0, -1)]));
labels.push(pad(clean));
}
const loss = await model.trainOnBatch([
tf.tensor2d(encInputs), tf.tensor2d(decInputs)
], tf.tensor2d(labels));
console.log(`epoch ${epoch} step ${step} loss ${loss.toFixed(4)}`);
}
}
}四、性能瓶颈与工程化建议
在Node.js中训练Transformer模型有几个现实问题需要面对。第一是计算性能,@tensorflow/tfjs-node虽然支持原生后端,但默认只使用CPU,如果服务器有GPU,需要安装@tensorflow/tfjs-node-gpu并配置CUDA环境,否则多层Transformer的训练速度会比Python慢一个量级。第二是内存管理,长序列大批次容易把内存打爆,建议用tf.tidy包裹中间张量操作,并在每步训练后调用tf.disposeVariables之外的清理逻辑及时释放不再使用的张量。
一个务实的工程策略是混合方案:预训练在Python侧完成,导出为ONNX格式后在Node.js中用onnxruntime-node做推理或微调,这样既保留了JavaScript技术栈的接口统一性,又获得了接近原生的训练效率。如果只是做小规模实验或教学演示,纯TensorFlow.js的方案完全够用;如果要在生产环境跑真实规模的BART,强烈建议采用这种混合架构。
另外,分词器建议直接复用SentencePiece训练好的词表模型,社区有JavaScript移植版本可以在Node.js中调用,避免自己维护一套大词表逻辑。训练过程中记得定期保存模型权重,model.save('file://./bart-checkpoint')即可把权重落到本地,断点续训时用tf.loadLayersModel恢复,这在长周期预训练任务中非常关键。