[ENH] `TransformedTargetRegressor` should mesh with linear or smooth transformations to produce exact `pdf` and `log_pdf`

[joshdunnlime]

Is your feature request related to a problem? Please describe.

LogLoss is a very commonly used proper scoring rule (e.g. logistic regression and NNs) and has efficient implementations for most distributions. The TransformTargetRegressor (TTR) as implementation does not include this. The TTR is lacking a log_pdf method, likely due to the numerical complications of mapping the log_pdf across domains. E.g. the transform function must be "a one-to-one, monotonic function with a differentiable inverse".

In reality, many of the functions used to transform the target have trivial differentiable inverses and therefore, trivial solutions to mapping the log_pdf. The MinMaxScaler internals for the proposed in SupportScaler is one such - it is a linear transformation so the derivative of the inverse is a constant.

Here is an walkthrough:

import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler

from skpro.metrics import LogLoss, CRPS
from skpro.regression.xgboostlss import XGBoostLSS
from skpro.regression.compose import TransformedTargetRegressor

size = 1000

X = np.random.normal(1, 1, size)
y = 2 * X + np.random.normal(0, 0.5, size)

Xy = pd.DataFrame(
    {"target": y, "feature": X},
    index=range(size)
)

# use lower n_trails for speed-up
params = {"n_trials": 2, "dist": "Normal"}

# simple example with no scaling - logloss works
xgb = XGBoostLSS(**params)
xgb.fit(X=Xy[["feature"]], y=Xy["target"])
p = xgb.predict(Xy[["feature"]])
pp = xgb.predict_proba(Xy[["feature"]])

print(CRPS()(y_true=Xy["target"], y_pred=pp))
print(LogLoss()(y_true=Xy["target"], y_pred=pp))

0.25535
0.56026

Now use the TTR:

xgb = XGBoostLSS(**params)
mms = MinMaxScaler()
pipe = TransformedTargetRegressor(regressor=xgb, transformer=mms)

# pipeline with TTR logloss fails
pipe.fit(X=Xy[["feature"]], y=Xy["target"])
pp = pipe.predict_proba(Xy[["feature"]])
p = pipe.predict(Xy[["feature"]])

print(CRPS()(y_true=Xy["target"], y_pred=pp))
print(LogLoss()(y_true=Xy["target"], y_pred=pp))

0.24274
inf

CRPS returns results. When trying to perform the LogLoss on a TTR we get inf. We can do a manual transformation and get the LogLoss on the new domain:

xgb = XGBoostLSS(**params)
mms = MinMaxScaler()
y_ = mms.fit_transform(Xy[["target"]])

xgb.fit(X=Xy[["feature"]], y=y_)
pp = xgb.predict_proba(Xy[["feature"]])
p = xgb.predict(Xy[["feature"]])
p_ = mms.inverse_transform(p.reshape(-1, 1))

crps = CRPS()(y_true=y_, y_pred=pp)
print(crps)
ll = LogLoss()(y_true=y_, y_pred=pp)
print(ll)

0.01859
-2.0715

Both the CRPS and LogLoss can be adjusted back to their respective domains by using the Change of Variable (applying the Jacobian of the transform appropriately).

jac = np.abs(1 / mms.scale_)
print((CRPS()(y_true=y_, y_pred=pp) * jac)[0])
print((ll + np.log(jac))[0])

0.25338
0.54385

I believe the mathematics behind this are correct. We multiply by the Jacobian of the inverse transform for CRPS. For LogLoss we add the log-Jacobian to log_pdf.

Describe the solution you'd like

1. Add the whole sk-style transformer to the TransformedDistribution class. This will allow more properties/method to be accessed - currently, it only has the inverse_transform as a callable. Adding the whole instance would allow access to the parameters. In the case of MinMaxScaler above, this allows us to access the .scale_ property and directly compute the deriative if the inverse.
2. Add the log_pdf method to the TranformedDistribution, with changes-of-variables Jacobain.

Bonus ideas:

A) Consider adding automatic differentiation of the inverse for more complex transformations.
B) Consider adding a _dinv_dx method (I'm sure there is a better name) to the transformer so that a user can directly formulate a derivative and skip auto-diff.

I see no reason to add the changes-of-variables version of the CRPS - it is purely here to help illustrate that the changes-of-variables score are somewhat consistent.

Describe alternatives you've considered

The "manual" option is the only other idea. This is prohibitive as it prevents users from getting meaningful LogLoss scores for comparison across different pipelines.

I believe the average user would be unaware of the "changes-of-variables" so would miss out on potentially using either TTR or LogLoss. Furthermore, it complements the other proposal of the SupportScaler. Again, this is a simple transformation so adding LogLoss for this is trivial. This allows LogLoss as a metric to GridSearch across multiple distributions with differing supports.

Additional context

[joshdunnlime]

Some questions/points for discussion:

1. Is this a useful enough feature?
2. Does the theory/mathematics seem sound?
3. Do we wish to support only linear transforms or to use one of the ideas from the bonus section to implement non-linear transformations?

[fkiraly]

Is this a useful enough feature?

Probably, if you are talking about the option to have smoothed distributions which allow pdf. Currently, the TTR uses the TransformedDistribution which is using quantile points.

How about the following approach: we add a quantile parametrized distribution that allows many quantiles and not just 3 like the current implementation? A smooth such QPD would allow log-loss, and any losses that require continuous distirbutions.

Does the theory/mathematics seem sound?

Yes - from a software perspective though it works only if we have access to the Jacobian or autodiff functionality of the transformation.

That is not true for generic transformations implemented in the scikit-learn interface, but it is typically true for transformations implemented in torch.

What one would need is some kind of merge of the two types of interfaces? E.g., scikit-learn compatible transformations that also expose gradients? Feels like a bigger project, unless you already know of a place where we can get the most common scikit-learn transformations with autodiff feature.

Do we wish to support only linear transforms or to use one of the ideas from the bonus section to implement non-linear transformations?

To me, it feels that software-wise, the difficulty will be equal. Since the feature or API endpoint we need to add is that a transformation is aware of its own Jacobian, and that the Jacobian is inspectable through a consistent API endpoint.

[fkiraly]

I think here is a good starter question: what is currently the best way to get a collection of transformations (sklearn-like or similar) together with autodiff interface? Or at least access to the Jacobian? And/or a tag that tells us whether it is monotonic.

Is the best way to reimplement everything? Or is there another option?

[joshdunnlime]

I'll slightly edit the questions:

What's the easiest way to get a set of sklearn-like transformations that return the Jacobian?

Check for hasattr(transformer, 'scale_') OR explicitly check for the following four classes:
StandardScaler, MinMaxScaler, MaxAbsScaler, RobustScale.

As these are linear transforms the Jacobian of the inverse is simply 1 / transformer.scale_. The hasattr-scale has the benefit of allowing a user to implement their own linear transformer (again SupportScaler is a great example).

For linear transformations, anything other than using the scale_ (or some equivalent) is overkill.

The only real change to facilitate this is giving the TransformedDistribution class full access to the transformer (instead of currently just the inverse_transform).

What's the "best way" (for non-linear)?

Automatic differentiation would probably be the nicest solution but this means rewriting a load of transformers. It would likely mean adding heavy new dependancies. Users who want to implement their own custom transformers would have to use torch, JAX or similar.

Another option: Numerical differentiation to the rescue.

If we expose the full transformer on the TransformedDistribution, this gives us access to both y (original) and y_t (transformed). Therefore, you can use np.gradient to compute the point-wise values of d(f-1(y))/dy pretty quickly and easily. I have try this out on a couple well defined functions and it seems to work well.

[joshdunnlime]

I have added a PR. It is fairly rough but I think this gives a good indication of how we can implement this and not need to make sweeping changed to existing transformers.

I have added a notebook to the PR. This is intended to show/check how adding the jacobian to the LogLoss (and also CRPS) gives the scores on the transformed domain. This can be deleted in future commits.

[fkiraly]

Check for hasattr(transformer, 'scale_') OR explicitly check for the following four classes:
StandardScaler, MinMaxScaler, MaxAbsScaler, RobustScale.

That would work, but it is not "API level" - when designing unified APIs, checking for specific classes should be avoided. Though it would be only for interaction with a single class, still it feels like there would be too much coupling. The best composition design is obtained when only the "generic API" of a certain type is assumed.

What's the "best way" (for non-linear)?

Automatic differentiation would probably be the nicest solution but this means rewriting a load of transformers. It would likely mean adding heavy new dependancies. Users who want to implement their own custom transformers would have to use torch, JAX or similar.

Is this really the best way? I can think of extending sklearn transformations with a method transform_diff or similar.

Another option: Numerical differentiation to the rescue.

Hm, that could work! Although we would have to live with numerical inaccuracies then.

[fkiraly]

PS: log-loss is a useful metric, but it can be very misleading in case the underlying data distribution has mass. It can end up with very high variance and unbounded values. CRPS imo is a better general choice.

But of course this feature request is still a very good idea imo.

[joshdunnlime]

That would work, but it is not "API level" - when designing unified APIs, checking for specific classes should be avoided. Though it would be only for interaction with a single class, still it feels like there would be too much coupling. The best composition design is obtained when only the "generic API" of a certain type is assumed.

That was why I opted for the scale_ attr check for linear transformers. It respects the sklearn API and allows users to extend it.

Is this really the best way? I can think of extending sklearn transformations with a method transform_diff or similar.

I agree, this is a nice solution. I like the option of giving the user ability to manually define the derivative. When I said "best" I meant from a user point of view (someone could easy mess up the derivative). The trade-off is the added dependancy.

How would you suggest adding transform_diff to an skpro transformer? Subclass, composition or maybe a Mixin?

Hm, that could work! Although we would have to live with numerical inaccuracies then.

Agreed. I think in most practical use case there will be enough instances i of y and y_t to get reasonable gradients for each i. These transformations should also be "well-behaved" (one-to-one, monotonic function with a differentiable inverse). We can add a warning if the number instances are below a certain threshold.

[fkiraly]

That was why I opted for the scale_ attr check for linear transformers. It respects the sklearn API and allows users to extend it.

I think that is a good idea, though I do not see this as a "general API" that is enforced by scikit-learn - e.g., there is no check in check_estimator of scikit-learn, it is only a very loosely encouraged convention.

I like the option of giving the user ability to manually define the derivative. When I said "best" I meant from a user point of view (someone could easy mess up the derivative). The trade-off is the added dependancy.

The user would just import the class and use it, a developer would need to implement the derivative. scikit-learn is unlikely to accept this kind of extension to their API, so one could use the visitor pattern? https://refactoring.guru/design-patterns/visitor

How would you suggest adding transform_diff to an skpro transformer? Subclass, composition or maybe a Mixin?

There currently are no skpro transformers - but we could actually write a base class, the topic has come up multiple times already (e.g., transforming distributions or parameter estimation in distributions). That might be a medium sized project though - so maybe not the easiest of all the clean ways to achieve what we want.

I think in most practical use case there will be enough instances i of y and y_t to get reasonable gradients for each i.

Wait, can you explain to me why we would need that?

In my mental model, the number of test instances is not relevant for obtaining of approximative derivatives, as long as we have unrestricted and fast access to the function we want to differentiate.

I was thinking we just naively substitute y[i] and y[i] + small_delta and y[i] - small_delta and get a numerical from a symmetric difference quotient, or a central finite difference of higher order https://en.wikipedia.org/wiki/Finite_difference_coefficient

[fkiraly]

How would you suggest adding transform_diff to an skpro transformer? Subclass, composition or maybe a Mixin?

Playing around with this idea, I think an easy way to make this happen would be something similar to a mixin, using the visitor pattern (see above) but adding the visitor method in a subclass that lives in skpro: we could inherit from sklearn transformers, and add the transform_diff method in skpro.

Whenever an sklearn transformer is passed to TransformedTargetRegressor, it is coerced to the corresponding skpro variant, and internally the transform_diff method is used, if available - otherwise a numerical differentiation is used.

I am not entirely comfortable with this option yet, but I do think it is a good sign that the combination of our ideas seems to produce something that is not too hard to create and not API breaking.

Thoughts appreciated!