Effect of input normalisation on the softmax function and its gradients.

plot_softmax.py

import autograd.numpy as np
from autograd import grad,elementwise_grad

def softmax(z):
  return (np.exp((z))) / np.sum(np.exp((z)))

nb_of_zs = 200
zs = np.linspace(-10, 10, num=nb_of_zs) # input 
zs_1, zs_2 = np.meshgrid(zs, zs) # generate grid
y = np.zeros((nb_of_zs, nb_of_zs, 2)) # initialize output
# variance norm
sample_var = np.sqrt(( np.var(zs_1)+np.var(zs_2) ) )
print sample_var
# Fill the output matrix for each combination of input z's
for i in range(nb_of_zs):
    for j in range(nb_of_zs):
        w=np.asarray([zs_1[i,j], zs_2[i,j]])
#         y[i,j,:] = elementwise_grad(softmax)(w/sample_var )
        y[i,j,:] = softmax( w )
# Plot the cost function surfaces for both classes
fig = plt.figure()
# Plot the cost function surface for t=1
ax = fig.gca(projection='3d')
surf = ax.plot_surface(zs_1, zs_2, y[:,:,0], linewidth=0, cmap=cm.coolwarm)
ax.view_init(elev=30, azim=70)
cbar = fig.colorbar(surf)
ax.set_xlabel('$z_1$', fontsize=15)
ax.set_ylabel('$z_2$', fontsize=15)
ax.set_zlabel('$y_1$', fontsize=15)
ax.set_title ('$P(t=1|\mathbf{z})$')
cbar.ax.set_ylabel('$P(t=1|\mathbf{z})$', fontsize=15)
plt.grid()
plt.show()