我正在尝试构建一个基于标量输入的数学操作选择神经网络模型。操作的选择基于神经网络产生的softmax结果。然后，该操作需要应用于标量输入以生成最终输出。到目前为止，我已经设想通过对softmax输出应用argmax和onehot来生成一个掩码，然后将该掩码应用于所有可能操作的连接值矩阵上（如下伪代码所示）。问题是，argmax和onehot似乎都不可微分。我是新手，所以任何建议都将非常感激。提前谢谢。

    #执行softmax        logits  = tf.matmul(current_input, W) + b    softmax = tf.nn.softmax(logits)    #对输入执行所有可能的操作    op_1_val = tf_op_1(current_input)    op_2_val = tf_op_2(current_input)    op_3_val = tf_op_2(current_input)    values = tf.concat([op_1_val, op_2_val, op_3_val], 1)    #创建掩码    argmax  = tf.argmax(softmax, 1)    mask  = tf.one_hot(argmax, num_of_operations)    #通过掩码去除未被选择的操作结果以生成输入    output = values * mask

回答：

我认为这是不可能的。这类似于论文中描述的硬注意力机制。硬注意力机制在图像标注中用于让模型在每个步骤中只关注图像的特定部分。硬注意力不可微分，但有两种方法可以绕过这个问题：

1- 使用强化学习（RL）：RL用于训练做出决策的模型。尽管损失函数不会将任何梯度反向传播到用于决策的softmax，但可以使用RL技术来优化决策。举个简化的例子，可以将损失视为惩罚，并向softmax层中值最大的节点发送一个与惩罚成比例的策略梯度，以便在决策不佳时（导致高损失）降低决策的得分。

2- 使用类似软注意力的方法：不是只选择一个操作，而是根据softmax的权重混合它们。因此，不是使用：

output = values * mask

而是使用：

output = values * softmax

现在，操作将根据softmax选择的程度逐渐收敛到零。与RL相比，这种方法更容易训练，但如果您必须从最终结果中完全去除未选择的操作（将它们完全设为零），它将不起作用。

这是另一个讨论硬注意力和软注意力的答案，您可能会发现有帮助：https://stackoverflow.com/a/35852153/6938290