Model no longer learning - any ideas?

Sorry, this is a drive-by comment. I’m not able review the particulars.

You say “The data also seems fine.” I my experience, that’s usually the problem even though I didn’t spot it at first. I suggest two things to sanity check that it’s not the data:

1. Try to train another, simpler model on the same data. Can a different model find a signal in that data?
2. Try to train your model on a contrived example. Can that model find a signal when it’s definitely present?

If you find that the answers to both those questions are “yes”, then truly something odd is happening. But in my own work, I find that at least one of those answers is usually “no”.

I replicated the issue. What seems weird (to me, not really an expert), is that I observe too many NaN values in the tensors in the output of the training step.

Is there any chance that the training data/labels are in a wrong format or that a kind of missing normalisation step is required?

Sean Moriarity wrote me that I should try axon 0.4.
With axon 0.4.1 the model is learning, the loss is going down during training.

Here is the diff:

<   {:axon, "~> 0.6"},
<   {:nx, "~> 0.7"},
<   {:exla, "~> 0.7"},
---
>   {:axon, "~> 0.4.1"},
>   {:nx, "~> 0.4.2"},
>   {:exla, "~> 0.4.2"},

< optimizer = Polaris.Optimizers.adamw(learning_rate: 1.0e-3)
---
> optimizer = Axon.Optimizers.adamw(1.0e-3)

https://github.com/preciz/not_learning/blob/master/model_learning.livemd

But now I would be happiest if newest Axon version would also learn. Any ideas on that?

There seems to be something wrong with binary_cross_entropy. Try categorical_cross_entropy for your loss function, it should be equivalent for two outputs:

loss =
  &Axon.Losses.categorical_cross_entropy(
    &1,
    &2,
    reduction: :mean
  )

Screenshot 2024-04-15 at 18.47.421200×1596 76.5 KB

If the change works on your end, could you report the issue to the Axon Github?

Awesome, with categorical_cross_entropy and softmax output activation the loss is going down and accuracy is >0.99.
Thank you, I will open an issue on Axon repo then.