# For prediction problems, why cant we simply use softmax as activation for hidden layers and no activation function for output layer

I am a beginner who is learning how to use neural networks for prediction problems, and almost all tutorials I have found use relu for hidden layer and sigmoid function for output later. My question is, if we simply use softmax for hidden layer and no activation function for output layer, can we get the same result? My inspiration for this is a neural network tutorial by statquest that can be found here https://www.youtube.com/watch?v=CqOfi41LfDw&t=1s. Please keep it simple. I am neither a statistician nor a data scientist

The short answer is "no;" we'll go through why. I'm assuming that you're using linear/affine layers; i.e. of the form $$X\theta + \beta$$ (matrix operations) with weights and bias $$\theta, \beta$$ respectively. Furthermore, I'll assume we're in the classification regime, given that you're talking about networks with softmax activation, which are generally used for classification (as opposed to regression).

I will try to keep this as high-level and intuitive as possible with links to more rigorous resources.

Tl;dr: Would X work? Maybe. Would it work "well? Probably not.

1) No activation function for output layer.

In regression, you probably wouldn't need/want an activation function, whereas for classification, people generally use activation functions to normalize the outputs to behave like a categorical distribution, so that cross-entropy loss can be used against one-hot labels. This also lends itself to an easy way to output model predictions: simply take the argmax of the outputs. More on the last-layer activation function here and here.

Concretely, in the binary classification case, the sigmoid activation given by $$\sigma(x) = 1/(1 + exp(-x))$$ is nice because it "squashes" all real numbers to the range $$(0, 1)$$. This means we get a "probability-like" score (this is not exactly accurate, but is an okay conceptual model for now) that a particular example belongs to a particular class, for which cross-entropy loss is the standard option. The softmax function, which can be thought of as a multi-dimensional generalization of the sigmoid function, does something similar in higher dimensions (some discussion of the differences and impact on applications here).

Why cross-entropy loss? A derivation demonstrating the correspondence between maximum likelihood estimation (MLE) and cross-entropy for a Bernoulli response variable can be found in Shalev-Shwartz & Ben-David (2014), Sec. 24.1, pg. 343. In summary, a single neural network neuron with a softmax activation can actually be thought of as performing multinomial logistic regression; the cross-entropy loss corresponds to the negative log-likelihood of the data.

It is theoretically possible/legal to have no activation function, and use a loss like MSE to perform classification. I don't know how this would shake out even in toy examples. In any case, this would be highly non-standard and I'm not sure what rule would be best to use to transform the outputs into a discrete prediction. I would not recommend this unless you have a very good mathematical reason to do this.

2) Softmax for hidden layers.

Using an activation function like softmax technically gives the neural network the requisite non-linearity required to approximate any* function of interest (the caveats are outside the scope of this answer, but some variants of the theorem are stated here with relevant citations). However, the key problem is that this can impact the optimization process. This is known as the vanishing gradient problem.

Most neural networks are optimized using iterative first-order methods. "Iterative" means that we start with some estimate of our weights and update them many times. "First-order" means that we're updating the weights using derivatives (or gradients in higher dimensions); i.e. with gradient descent, for which I find a hill-descending analogy useful.

Basically, the derivative of the softmax/sigmoid functions at places "far" from zero is close to zero. This is called "saturation." This means when we take the derivative of the whole network, the derivative is small, so our model doesn't update very much (and therefore doesn't "learn" much). This follows diretly from the multiplicative nature of the chain rule. As we add more layers with small derivatives -- this problem gets worse and worse.

To see this, note that $$\sigma'(x) = \sigma(x) (1 - \sigma(x))$$ (a good exercise). The softmax has a Jacobian matrix instead; calling the Jacobian $$\mathbf{J}$$, and calling the softmax output $$S$$, the entries are given by

$$\mathbf{J}_{ij} = S_i (\delta_{ij} - S_j)$$

where $$\delta_{ij}$$ is the indicator $$\mathbf{1}[i = j]$$, i.e. the Kronecker delta. The derivation is a great exercise, and was a common machine learning exam question for me. (Solution)

I'll leave the rest to you: to reach the same conclusion (small derivative magnitude) as we did for the sigmoid function, given that $$S_i \in [0, 1]$$ for all $$i$$, and that $$\sum_i S_i = 1$$, plug in a few values of $$S_i$$ and see how $$S_j$$ are constrained. Then examine the resultant values of $$\mathbf{J}_{ij}$$. Can you determine when $$\mathbf{J}_{ij}$$ is maximized? What about the conditions under which $$\mathbf{J}_{ij}$$ is "small?"

Hence, using an activation like ReLU ($$\text{ReLU}(x) = \max\{0, x\}$$), which trivially has derivative 1 in the non-negative domain, sidesteps this problem. However, using the same derivative-based analysis, there's another issue with the ReLU function that comes up -- can you see what it is? (Answer)

Small tips about asking ML questions in the future. For future reference, it's helpful if you specify the type of task you're trying to solve specifically (i.e. are you predicting house prices? Are you predicting whether a patient has X conditions), as well as more details about the model (i.e. what's the structure of the neural network, as exactly as possible? what are the inputs and outputs?)

You can, and there is a sense in which doing so would be analogous to a linear probability model after layers of feature extraction, much as using a sigmoid activation function on the output neuron is analogies to logistic regression after layers of feature extraction. In fact, you don’t even have to limit yourself to sigmoid activation in the hidden layers to be analogous to a linear probability model.

A common criticism of a linear probability model is that the output is not guaranteed to be restricted to $$[0,1]$$, and your neural network that has all sigmoid activation functions in the hidden layers and then no activation function in the final neuron can suffer from this same issue. If this seems like a dealbreaker, then you might not like a model like you’ve proposed that forces the output to be restricted to $$[0,1]$$, as the sigmoid activation functions in the hidden layers need not force the final output to be restricted to $$[0,1]$$, even though the values in the hidden neurons will have that restriction.

Cross-entropy (XEnt) loss is common for these classification neural networks. To be analogous to a linear probability model, square loss would be used. Further, XEnt presents problems with predictions outside of the interval, $$[0,1]$$, due to the $$\log(\hat y)$$ and $$\log(1-\hat y)$$ terms (where $$\hat y$$ is the prediction).