关于gumbel softmax的使用

[Seeker98]

作者你好，我注意到计算cls policy的gumbel softmax并不是一个确定性算子，而且代码中没有出现训练和推理的不同设定。请问在推理中仍然使用gumbel softmax而不是某个对应的特定算子（argmax, onehot）的考量是什么？

[csguoh]

Yes, during inference we keep using the gumble softmax since we find the randomness of this operator doesn't fluctuate performance too much. So for simplicity, we keep this during inference. If you want to totally eliminate the randomness, it is suggested to use deterministic sampling during training instead of directly switching to argmax at inference.

[Seeker98]

Thanks for the instant reply.

As I noticed that the module route that produces pred_route acts as some kind of classifier, and before the logits (pred_route) are fed into gumbel_softmax, a log_softmax is applied.
Numerically, this means the differences between the elements of such a logits vector are further compressed, which can also be observed in the forward hook.

From my opinion, when sampling with gumbel_softmax from such 'compressed difference' logits makes the model unfavorable (at least in terms of performance) because it struggles to determine which group a pixel should belong to.
Since Mamba itself is unidirectional, if some pixels meaningful to the current pixel are placed later in the sequence due to such ambiguous grouping, this clearly affects the model.

Based on the above findings, I wonder if there are any exchangeable insights here, as you are indeed an expert in this field.

[csguoh]

Hi,
In the appendix in the paper, we had give some visualization of how pixels belonging to the same group is re-ordered. The results show the classifier indeed has learnt the semantics. For the unidirectional nature of Mamba, we use the semantic prompts as a complement for the problem you mentioned. Hope my reply helps.

[Seeker98]

Wow, your answer's perspective is really insightful, while in the above I've been stuck in the training dynamic matters :(