Double log_softmax on logit in rnnt_loss_tf function

[thanatl]

In rnnt_loss_tf function, you take the tf.nn.log_softmax(logits) on line 238, and you take the tf.nn.log_softmax(logits) again in the compute_rnnt_loss_and_grad_helper on line 159

is this intentional or an error?

reference file:
https://github.com/TensorSpeech/TensorFlowASR/blob/main/tensorflow_asr/losses/rnnt_losses.py#L150

[nglehuy]

@thanatl Oh, I didn't notice, let me dive deeper to that code.

[mjurkus]

Sorry, for the ignorance, but what is the effect of double log_softmax?

[thanatl]

@mjurkus well, I just recheck, compute log_softmax multiple time output the same result as compute log_softmax once.

[nglehuy]

@thanatl Cool, then we should only use 1 time 😄

[nglehuy]

@thanatl @mjurkus I just found out that if we use 2 log_softmax as before, we can use larger batch size (I tested on TPUs, using 2 log_softmax allows me to use batch size 2 where single log_softmax only 1).
@monatis Can you tell us why you use 2 log_softmax at the first place 😄 ?

[monatis]

@usimarit Well, this line originally comes from the borrowed code, so I don't know the exact intention of the original author, but it may be related to an issue of numerical stability imo, as they don't give the same result always --see this quick test.

[nglehuy]

@monatis I just checked out the author's repo, I only see 1 log_softmax.

[monatis]

Oh then it must be my mistake when copying some lines around to fix the broken code, but I don't remember exactly how it happened. Sorry for the inconvenience.