TensorFlow.js多头注意力反向传播报错：梯度与输入形状不一致

尝试使用tensorflowjs实现多头注意力机制时，训练阶段持续报错“梯度形状与输入形状不一致”。前向传播正常，手动核对张量形状也无问题，已尝试清理中间张量、明确指定更新权重，但报错未解决。

可复现代码：

const EMBEDDING_DIMENSION = 16;
const NUMBER_OF_HEADS = 2;
const BATCH_SIZE = 8;
const SEQUENCES_LENGTH = 4;
const EPOCHS = 10;

tf.setBackend("cpu");
tf.enableDebugMode();

function dotProductAttention(query, key, value) {
  return tf.tidy(() => {
    const matMul = query.matMul(key, false, true);
    const dk = tf.scalar(query.shape[query.shape.length - 1], 'float32');
    const scaled = matMul.div(tf.sqrt(dk));
    const attentionWeights = scaled.softmax(-1);
    return attentionWeights.matMul(value);
  })
}

function muliHeadAttention(input, numberOfHeads, qkvWeights, outputWeights) {
  return tf.tidy(() => {
    const batchSize = input.shape[0];
    const seuquencesLength = input.shape[1];
    const embeddingDimension = qkvWeights.shape[0];
    const headDimension = embeddingDimension / numberOfHeads;

    // calculate all query, key, and value by doing one big mat multipication.
    const QKV = input.matMul(qkvWeights);

    // reshap the QKV to include number of heads
    const QKVReshaped = QKV
      .reshape([batchSize, seuquencesLength, numberOfHeads, 3 * headDimension])
      .transpose([0, 2, 1, 3]);

    // split the mega QKV into 3 chunk each query, key, and value matrix
    const [query, key, value] = QKVReshaped.split(3, -1);

    // calculate the attention of query, key, and value
    const attention = dotProductAttention(query, key, value);

    // concatenate the attention values
    const concatenated = attention
      .transpose([0, 2, 1, 3])
      .reshape([batchSize, seuquencesLength, embeddingDimension]);

    // matrix multiply with output weights
    const output = concatenated.matMul(outputWeights);

    // finally, return it [batchSize, seuquencesLength, embeddingDimension]
    return output;
  })
}

function initializeQKVWeights(embeddingDimension) {
  return tf.variable(tf.randomNormal([embeddingDimension, 3 * embeddingDimension]));
}

function initializeOutputWeights(embeddingDimension) {
  return tf.variable(tf.randomNormal([embeddingDimension, embeddingDimension]));
}

// query, key, and value weights combined for better performance.
const QKVWeights = initializeQKVWeights(EMBEDDING_DIMENSION);

// final weight of muti head attention 
const outputWeights = initializeOutputWeights(EMBEDDING_DIMENSION);

// the predict function
const f = (x) => muliHeadAttention(x, NUMBER_OF_HEADS, QKVWeights, outputWeights);

// mean squared error
const loss = (predict, real) => predict.sub(real).square().mean();

// the optimizer
const optimizer = tf.train.sgd(0.01);

// temporary generated train datas
const xTrain = tf.randomNormal([BATCH_SIZE, SEQUENCES_LENGTH, EMBEDDING_DIMENSION]);
const yTrain = tf.randomNormal([BATCH_SIZE, SEQUENCES_LENGTH, EMBEDDING_DIMENSION]);

// training the model
for (let i = 0; i < EPOCHS; i++) {
  tf.tidy(() => {
    const lossValue = optimizer.minimize(() => loss(f(xTrain), yTrain), [QKVWeights, outputWeights]);

    // Logging loss every epoch
    lossValue.data().then(lossValueArray => {
      console.log(`Epoch ${i + 1}, Loss: ${lossValueArray[0]}`);
    });
  });
}

f(xTrain).array().then(pred => {
  console.log("Predictions:");
  console.log(pred);
  xTrain.dispose();
});

// Disposing tensors
yTrain.dispose();;

核心问题：tf.tidy在反向传播中的误用 #

在muliHeadAttention和dotProductAttention中用tf.tidy包裹了整个运算逻辑，但反向传播需要依赖前向过程生成的中间张量计算梯度，tf.tidy会自动清理这些中间张量，导致反向传播时缺失关键计算依赖，进而引发形状不匹配的报错。前向传播正常是因为无需依赖这些中间张量，反向传播阶段链路断裂才抛出异常。

修复步骤 #

1. 移除核心运算函数中的tf.tidy
dotProductAttention和muliHeadAttention是模型核心逻辑，前向生成的中间张量是反向传播的必要依赖，不能被tf.tidy清理。

2. 调整训练循环中的tf.tidy范围
   仅在训练循环中清理每次迭代生成的临时张量（如loss相关的展示用张量），避免影响模型运算的核心依赖。

修正后的关键代码片段 #

// 移除tf.tidy包裹
function dotProductAttention(query, key, value) {
  const matMul = query.matMul(key, false, true);
  const dk = tf.scalar(query.shape[query.shape.length - 1], 'float32');
  const scaled = matMul.div(tf.sqrt(dk));
  const attentionWeights = scaled.softmax(-1);
  return attentionWeights.matMul(value);
}

// 移除tf.tidy包裹
function muliHeadAttention(input, numberOfHeads, qkvWeights, outputWeights) {
  const batchSize = input.shape[0];
  const seuquencesLength = input.shape[1];
  const embeddingDimension = qkvWeights.shape[0];
  const headDimension = embeddingDimension / numberOfHeads;

  const QKV = input.matMul(qkvWeights);
  const QKVReshaped = QKV
    .reshape([batchSize, seuquencesLength, numberOfHeads, 3 * headDimension])
    .transpose([0, 2, 1, 3]);
  const [query, key, value] = QKVReshaped.split(3, -1);

  const attention = dotProductAttention(query, key, value);
  const concatenated = attention
    .transpose([0, 2, 1, 3])
    .reshape([batchSize, seuquencesLength, embeddingDimension]);
  const output = concatenated.matMul(outputWeights);

  return output;
}

// 训练循环调整tf.tidy范围
for (let i = 0; i < EPOCHS; i++) {
  optimizer.minimize(() => {
    const pred = f(xTrain);
    const lossVal = loss(pred, yTrain);
    // 仅在loss计算后清理临时张量
    tf.tidy(() => {
      lossVal.data().then(lossValueArray => {
        console.log(`Epoch ${i + 1}, Loss: ${lossValueArray[0]}`);
      });
    });
    return lossVal;
  }, [QKVWeights, outputWeights]);
}

额外优化建议 #

• 避免在训练循环中频繁调用data().then()，改用await同步获取loss值，简化代码同时避免异步导致的张量生命周期问题。
• 训练结束后手动清理所有不再使用的变量和张量，避免内存泄漏。

内容的提问来源于stack exchange，提问作者MrGeniusProgrammer