Hi Zhihan!



 Zhihan_Yang:

Therefore, it’s unclear to me why adding log to it makes it unstable.


First try some numerical experiments.
Use the double-precision version of the naive expression as your
assumed “correct” result.  (The double-precision calculation will also
have the numerical overflow issue, but it won’t set in as soon.)
true_result = torch.log (torch.softmax (alpha * torch.tensor ([-1.0, 0.0, 1.0]).double(), dim = 0))

Then compare the single-precision results from
torch.log (torch.softmax (alpha * torch.tensor ([-1.0, 0.0, 1.0]), dim = 0))
# and
torch.log_softmax (alpha * torch.tensor ([-1.0, 0.0, 1.0]), dim = 0)

with one another and with your “true” result for increasing values of alpha,
say, alpha = 2, 5, 10, 20, 50, 100, ..., and see how the results behave.
Second, write two of your own versions of log_softmax(), one where you
just use the naive log (softmax()) approach, and a second where you
apply the “log-sum-exp trick” to the sum of exponentials in the denominator
of the formula for softmax().
Does the log-sum-exp trick significantly reduce the overflow issue?  This
is pretty much how log_softmax() is implemented in pytorch.
Best.
K. Frank