为何PyTorch EmbeddingBag的"max"模式不支持传入per_sample_weights？

实现层的硬限制 #

• 翻PyTorch核心的EmbeddingBag.forward校验逻辑就能看到，只要传入per_sample_weights，第一步就会做分支判断：仅当mode == "sum"时才会继续执行后续流程，一旦是max模式直接抛出异常说明加权参数仅对sum模式生效，从函数入口就堵上了兼容的可能。
• 原生max模式的计算流是：先按索引批量查嵌入向量，再在每个样本的变长序列维度上逐元素取最大值，整个计算链路从设计初期就没有预留和per_sample_weights做张量交互的分支，不存在“临时加个乘法就能支持”的改造空间。

语义层面的天然不兼容 #

对比sum/average模式的加权逻辑就能明白差异的根源：

• sum/average属于线性归约运算，加权过程是w_i * e_i逐元素相乘后累加，整个运算链路连续可导，就算要从加权求和转成加权平均，只要除以权重总和即可，全程满足线性运算的封闭性，没有逻辑断层。
• max操作本质是逐元素选取输入张量的最大值，属于带断点的非线性运算，完全没法照搬线性加权的逻辑实现有意义的加权效果：

  举个最直白的反例：同个序列下两个嵌入值分别为e1=2、e2=3，对应样本权重w1=0.9、w2=0.1。如果硬套“乘完权重再取max”的逻辑，计算结果是max(0.9*2, 0.1*3) = 1.8，完全偏离了“高权重样本对结果影响更大”的直觉——原值更大的e2因为权重低，加权后的结果反而比e1更低，整个计算结果没有可解释的业务价值，碰到负权重的场景结果更是会完全失控。
  这点也和你提到的补充说明一致：加权求和可以通过除以权重和简单转成加权平均，但max操作根本没有对应的等效线性变换路径，所谓“加权max”从定义上就没有统一的、被行业认可的标准实现。

梯度计算与算子生态的连锁障碍 #

你猜测的和梯度、ONNX算子实现相关的顾虑完全是成立的：

• 梯度链路的问题：max本身的梯度就只会反向传播到被选中的最大值元素上，属于带动态判断的非连续梯度。如果硬要加入per_sample_weights，相当于在max计算前先做一次乘法，会动态改变哪个元素会被选中的结果，不仅嵌入向量的梯度传播路径会随机跳转，权重本身的梯度计算也没有稳定的闭式表达，训练过程很容易出现梯度抖动。
• 跨框架算子兼容的问题：
  • sum模式下的加权逻辑可以拆成Mul(嵌入查找结果, per_sample_weights) -> ReduceSum的标准算子链，就算要实现加权平均，也只需要在后面多接一个ReduceMean除以权重和即可，不管是ONNX符号定义、还是各类推理引擎都是原生支持的，不存在适配障碍。
  • 如果要给max模式加加权，对应算子链是Mul(嵌入查找结果, per_sample_weights) -> ReduceMax，抛开语义问题不说，端侧框架对这类带前置乘法的动态max算子的优化、梯度支持本来就很不完善，根本没有规模化落地的技术基础。你之前观察到的Mul后接ReduceMean顺畅、接ReduceMax有问题的现象，本质上就是运算的线性属性差异带来的，不是技术实现上的临时bug。

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