"IF statement" in a Neural Network In classification problems, it is frequent to have classes with different properties. For example, I came across a problem where I needed to classify the following images in a single network:

*

*White mugs

*Black mugs

*Blue mugs

*Glasses

As my approach was getting more complex, the accuracy between glass vs mugs (due to the lights) was increasing, but after looking at the confusion matrix, black vs blue vs white seemed to overfit. Likewise as my model was getting simpler, glass was confused for white and blue for black and black for white.
How can I write a NN such as:

If it is a mug, then classify mugs only, else it is a glass.

If this concept makes any sense (let aside the mug oversampling issue), I would appreciate any reference to papers or how can this be implemented in theory.
 A: You can make your network aware that high probabilities of having a glass should make it detrimental to output probabilities of mug colors, a soft IF statement implemented with a gating mechanism.
At the end of your network $\mathcal N$ you could add two branches: $\mathcal G$ (for glass) & $\mathcal M$ (for mug).

*

*$\mathcal G$ outputs a probability that decides if the image contains or not glasses.

*$\mathcal M$ outputs three probabilities, which pertain to the colors of the mugs.

*$\mathcal M$ probabilities are multiplied by $(1 -\mathcal G)$ probabilities, making it aware if said object is, in fact, not a mug.

You end up having to optimize for two losses: the binary cross-entropy of glass-vs-mug and the categorical cross-entropy for white-vs-black-vs-blue mugs. You can, of course, fuse them into a single categorical cross-entropy with the probabilities: $\mathcal M_{\text White}, \mathcal M_{\text Black}, \mathcal M_{\text Blue}, \mathcal G$, keeping the rest of the branches intact.
