EvanZhuang / MetaTree

Official implementation of MetaTree: Learning a Decision Tree Algorithm with Transformers
https://arxiv.org/abs/2402.03774
MIT License
97 stars 5 forks source link

To regression problem? #3

Open intMinsu opened 2 weeks ago

intMinsu commented 2 weeks ago

Hi, thanks for opening the code. Can I use this model to solve regression problem?

EvanZhuang commented 2 weeks ago

Thanks for your interest!

The native code supports classification only. I think with some modifications it should be able to handle regression, basically changing the BCE classification loss into a regression loss, and substituting label embedding from one-hot classes into something continuous.