TensorFlow.js tf.multinomial结果不符：预期[75000,25000]却得[62083,37917]

问题根源 #

tf.multinomial（和tf.random.categorical功能完全一致）的参数不是直接传入概率值，而是传入对数概率（logits）——也就是未经过softmax归一化的权重。

你传入的[0.75, 0.25]被当作logits处理，内部会通过softmax转换为实际采样概率：

softmax([0.75, 0.25]) = [e^0.75/(e^0.75+e^0.25), e^0.25/(e^0.75+e^0.25)] ≈ [0.622, 0.378]

这个比例正好和你得到的[62083, 37917]匹配，这就是采样结果偏差的核心原因。

解决方法 #

有两种正确实现方式，同时建议优化采样效率：

1. 将概率转换为logits传入

对目标概率取自然对数，得到对应的logits后再传入函数：

const tf = require('@tensorflow/tfjs-node');

// 把概率转为对数概率（logits）
const logits = tf.tensor([Math.log(0.75), Math.log(0.25)], [2], 'float32');
// 一次性采样100000次，避免循环创建大量临时张量
const samples = tf.multinomial(logits, 100000).arraySync();

// 统计采样结果
const count = [0, 0];
samples.forEach(s => count[s]++);
console.log(count); // 结果会接近[75000, 25000]

2. 直接使用tf.random.categorical（推荐）

tf.random.categorical的文档对参数说明更清晰，同样接受logits，用法和multinomial一致：

const tf = require('@tensorflow/tfjs-node');

const logits = tf.tensor([Math.log(0.75), Math.log(0.25)], [2], 'float32');
const samples = tf.random.categorical(logits, 100000).arraySync();

const count = [0, 0];
samples.forEach(s => count[s]++);
console.log(count);

额外优化说明 #

你原来循环10万次每次采样1个结果的方式，会创建大量临时张量，既影响性能也可能导致内存占用过高。一次性采样10万次是更高效的做法。

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