Remove special case for momentum, I think theano will factor the zero multiplication out of the graph(?)