How to replace argmax with softmax so that I dont get zero gradient?

[velkyvont]

I'm trying to modify MyFunc (see below) so that I dont get zero gradient at all input values. The FAQ states:

The jax.nn submodule also has smooth versions of other common rank-based functions, for example jax.nn.softmax() can replace uses of jax.numpy.argmax(), jax.nn.soft_sign() can replace uses of jax.numpy.sign(), jax.nn.softplus() can replace uses of jax.nn.relu(), etc.

So I thought if I modify - my original code:

if result == 1.0:
    if fPoints >= dPoints:
        balance += fOdds*1000

with:

fav = np.array([fPoints.primal, dPoints.primal])	
favC = nn.softmax(np.array(fav))
if result == 1.0 and favC[0] >= 0.5:
    balance += fOdds*1000

then I would finally not get zero gradients but this doesn't work either. Could someone help please?

from jax import grad
import jax.nn as nn
import jax.numpy as jnp
import numpy as np
import json

# Example data - in real case I have 150,000+ rows
data = jnp.array([[ 1.  ,  1.06,  9.77,  5.  ,  3.  ,  2.  ,  6.  , 12.  ,  4.  ,
       10.  ,  1.  ,  7.  ,  1.  , 12.  , 10.  , 12.  ,  4.  , 10.  ,
        8.  ,  7.  , 11.  ,  5.  ,  9.  ,  3.  ,  6.  , 12.  ,  6.  ,
        5.  ,  3.  ,  5.  ,  9.  ,  8.  ,  9.  , 10.  , 11.  , 12.  ,
        1.  ,  2.  ,  3.  ,  4.  ,  5.  ,  6.  ,  7.  ,  0.  ,  0.  ,
        0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,
        0.  ,  0.  ,  0.  ], 
        [ 1.  ,  1.33,  3.33,  5.  ,  3.  ,  2.  ,  6.  , 12.  ,  4.  ,
       10.  ,  1.  ,  7.  ,  1.  , 12.  , 10.  , 12.  ,  4.  , 10.  ,
        8.  ,  7.  , 11.  ,  5.  ,  9.  ,  3.  ,  6.  , 12.  ,  6.  ,
        5.  ,  3.  ,  5.  ,  9.  ,  8.  ,  9.  , 10.  , 11.  , 12.  ,
        1.  ,  2.  ,  3.  ,  4.  ,  5.  ,  6.  ,  7.  ,  0.  ,  0.  ,
        0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,
        0.  ,  0.  ,  0.  ], 
        [ 2.  ,  1.65,  2.07,  5.  ,  3.  ,  2.  ,  6.  , 12.  ,  4.  ,
       10.  ,  1.  ,  7.  ,  1.  , 12.  , 10.  , 12.  ,  4.  ,  8.  ,
        6.  ,  5.  ,  9.  ,  3.  ,  7.  ,  1.  ,  4.  , 10.  ,  4.  ,
        3.  ,  1.  ,  3.  ,  7.  , 10.  , 11.  , 12.  ,  1.  ,  2.  ,
        3.  ,  4.  ,  5.  ,  6.  ,  7.  ,  8.  ,  9.  ,  0.  ,  0.  ,
        0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,
        0.  ,  0.  ,  0.  ]])

        
def MyFunc(coefs, data):
	balance = float(len(data)*-1000)
	
	for row in data:
		result = row[0]
		fOdds = row[1]
		dOdds = row[2]

		for p in range (0, 14):
			s = (row[3+p]-1).astype(int)
			h = (row[17+p]).astype(int)
			r = (row[43+p]).astype(int)
				
			bCoef = coefs[p]
			sCoef = coefs[14 + (p * 12) + s]
			hCoef = coefs[182 + (p * 12) + h]
			
			pStrength = bCoef * sCoef * hCoef
			
			houseKoefs = jnp.zeros(12)
			for h in range (0, 12):
				houseKoefs = houseKoefs.at[h].set(houseKoefs[h]+pStrength)
		
		for h in range (0, 12):
			hSign = (row[31+h]-1).astype(int)
			houseKoefs = houseKoefs.at[h].set(houseKoefs[h]*coefs [364 + (h*12) + hSign])
							
		fPoints = 0.0
		dPoints = 0.0
			
		for h in range (0, 12):
			fPoints += houseKoefs[h] * coefs[508+h]
			dPoints += houseKoefs[h] * coefs[520+h]
	
		fav = np.array([fPoints.primal, dPoints.primal])	
		favC = nn.softmax(np.array(fav))
		dog = np.array([dPoints.primal, fPoints.primal])	
		dogC = nn.softmax(np.array(dog))
	
		if result == 1.0 and favC[0] >= 0.5:
			balance += fOdds*1000
		elif result == 2.0 and dogC[0] > 0.5:
			balance += dOdds*1000

	return balance


derivFunction = grad (MyFunc)
coefs = np.random.sample(532)
# here I just get a list of 532 zeros instead of the derivatives...
print (derivFunction(coefs, data))
coefs = np.random.sample(532)
print (derivFunction(coefs, data))

[jakevdp]

I don't quite follow your full code example, but I think the way to do a "soft" version of your snippet:

if fPoints >= dPoints:
  balance += fOdds*1000

would be something like this:

import jax.nn
balance += fOdds * 1000 * jax.nn.sigmoid(fPoints - dPoints)

This uses the Sigmoid function in place of your (implied) step function, so that if fPoints is much smaller than dPoints, the resulting operation is balance += 0, and if fPoints is much larger than dPoints, the resulting operation is balance += fOdds * 1000, and the output varies smoothly between these two extremes.

The recommendation for softmax in the docs you linked to is as a replacement for argmax, which your code doesn't use, so softmax is not particularly relevant here.

Hope that helps!

[velkyvont]

@jakevdp Thanks for the answer. My original code had comparison if fPoints >= dPoints which can be effectively written as if argmax ([fPoints, dPoints]) == 0 - this is why I was thinking about softmax.

[jakevdp]

Ah, I see what you're saying. In that sense, your original logic expressed in terms of argmax is something like this:

weights = jnp.array([fOdds * 1000, 0.0])
i = jnp.argmax(jnp.array([fPoints, dPoints]))
balance += weights[i]

And the softmax alternative would be

i = nn.softmax(jnp.array([fPoints, dPoints]))
balance += jnp.dot(weights, i)