[BUG] Fix quantile support in Samformer

[Faakhir30]

Reference Issues/PRs

Refs ##2340 (comment)

What does this implement/fix? Explain your changes.

Samformer:

# blocks of code from samformer
        self.linear_forecaster = nn.Linear(
            self.max_encoder_length, self.max_prediction_length
        )  # noqa: E501
...
       out = self.linear_forecaster(out)

        out = out.transpose(1, 2)

        target_predictions = out[:, :, -1]  # (batch_size, max_prediction_length)
...
        if self.n_quantiles > 1:
            target_predictions = target_predictions.unsqueeze(-1).expand(
                -1, -1, self.n_quantiles
            )
        elif self.n_quantiles == 1:
            target_predictions = target_predictions.unsqueeze(-1)
        return {"prediction": target_predictions}

This doesnt seem right to me, I think its just generating max_prediction_length single quantile and .expand is just replicating that single quantile to n_quantile times

each quantile should have its own parameters and receive its own gradient, allowing the model to learn different quantiles. i.e.,

self.linear_forecaster = nn.Linear(
            self.max_encoder_length, self.max_prediction_length * n_quantiles
        ) 

PR checklist

• The PR title starts with either [ENH], [MNT], [DOC], or [BUG]. [BUG] - bugfix, [MNT] - CI, test framework, [ENH] - adding or improving code, [DOC] - writing or improving documentation or docstrings.
• Added/modified tests
• Used pre-commit hooks when committing to ensure that code is compliant with hooks. Install hooks with pre-commit install.
  To run hooks independent of commit, execute pre-commit run --all-files

[fkiraly]

I am confused - I think your initial claim (the models do not support quantiles) is incorrect.

If a quantile loss is used at the last/output layer, then the predictions indeed become quantile predictions.

This is because the argmin of a quantile loss expectation (true distribution vs variable in the armin) is the (generative) quantile. Similarly, if you take multiple quantiles in the quantile loss, you get multiple quantiles in the forecast.

[fkiraly]

These model arcitectures are not "quantile-aware".

All that needs to be done for that is to use a quantile loss (or a distributional loss such as log-loss) in the output layer.

Which, as far as I can see in the current implementation, is exactly what happens?

[fkiraly]

Please correct me explicitly though if you think I am wrong.

[phoeenniixx]

I think this is my mistake @fkiraly. See this comment: #2340 (comment)

I had a misconception that for quantile support, just using loss at the output layer is not enough, it should be used somehow in the deeper layers as well. My bad

[phoeenniixx]

I thought in the model, the gradients will be diluted by the time they reach the inner layers and it is not necessary that it is completely aware of the quantiles. But now when I think in more detail, it is an issue with extremely deep networks (where we might need to use auxiliary heads to compute the loss in the middle of the layers as well). I didnt think of it by the chain rule and realised these are not that deep models. And by thinking that some models just support point pred losses, It made me feel that my thinking is right. But now I understand, it is just the case of the output layer not supporting the loss

[phoeenniixx]

• Increasing dimention by output_dim = self.prediction_length * self.n_quantiles of linear head doesnt make model quantile-aware

This was my thinking :)

[Faakhir30]

Increasing dimention by output_dim = self.prediction_length * self.n_quantiles of linear head doesnt make model quantile-aware

Alright, now that we have agreed that above statement is false, I think timexer and Dlinear were already having correct support, reverted those.
I just have one final comment regarding Samformer:

# blocks of code from samformer
        self.linear_forecaster = nn.Linear(
            self.max_encoder_length, self.max_prediction_length
        )  # noqa: E501
...
       out = self.linear_forecaster(out)

        out = out.transpose(1, 2)

        target_predictions = out[:, :, -1]  # (batch_size, max_prediction_length)
...
        if self.n_quantiles > 1:
            target_predictions = target_predictions.unsqueeze(-1).expand(
                -1, -1, self.n_quantiles
            )
        elif self.n_quantiles == 1:
            target_predictions = target_predictions.unsqueeze(-1)
        return {"prediction": target_predictions}

This still doesnt seem right to me, I think its just generating max_prediction_length single quantile and .expand is just replicating that single quantile to n_quantile times

each quantile should have its own parameters and receive its own gradient, allowing the model to learn different quantiles. i.e.,

self.linear_forecaster = nn.Linear(
            self.max_encoder_length, self.max_prediction_length * n_quantiles
        ) 

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (main@e7c7965). Learn more about missing BASE report.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2357   +/-   ##
=======================================
  Coverage        ?   87.34%           
=======================================
  Files           ?      171           
  Lines           ?    10075           
  Branches        ?        0           
=======================================
  Hits            ?     8800           
  Misses          ?     1275           
  Partials        ?        0           

Flag	Coverage Δ
cpu	87.34% <100.00%> (?)
pytest	87.34% <100.00%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[fkiraly]

I thought in the model, the gradients will be diluted by the time they reach the inner layers and it is not necessary that it is completely aware of the quantiles. But now when I think in more detail, it is an issue with extremely deep networks (where we might need to use auxiliary heads to compute the loss in the middle of the layers as well). I didnt think of it by the chain rule and realised these are not that deep models. And by thinking that some models just support point pred losses, It made me feel that my thinking is right. But now I understand, it is just the case of the output layer not supporting the loss

I see - it may make sense to use different networks or at least heads for different quantiles so the model is better. But it also does not feel strictly necessary, more of a "one can try".

However, I think all that matters from the API perspective is whether the last layer reasonably claims to produce quantiles (not whether the model is actually good). For this, having a multi-quantile-loss in the last layer is already sufficient - even if it may or may not produces the best quantile predictions that could be obtained.

In this, there is also an idea for an interesting wrapper or operation: split the NN at some point or add heads onto it for different quantiles. Though maybe a lower prio tangent compared to ptf2 release.