nusdbsystem / ARM-Net

A ready-to-use framework of the state-of-the-art models for structured (tabular) data learning with PyTorch. Applications include recommendation, CRT prediction, healthcare analytics, anomaly detection, and etc.
Apache License 2.0
73 stars 13 forks source link

Multiple output neurons (multiclass classification) support #1

Closed GillianGrayson closed 2 years ago

GillianGrayson commented 2 years ago

Hi!

How to correctly support multiple output neurons in your models (for example for multiclass classification task)?

1) armnet and armnet_1h support multiple output neurons via noutput parameter - there is no problem.

2) For some models, I added support for multiple output neurons by explicitly specifying the noutput parameter in the last MLP-layer, or by adding support for multiple output neurons for the Linear as follows:


class Linear(nn.Module):

    def __init__(self, nfeat, noutput=1):
        super().__init__()
        self.noutput = noutput
        self.weight = nn.Embedding(nfeat, noutput)
        self.bias = nn.Parameter(torch.zeros((1,)))

    def forward(self, x):
        """
        :param x:   {'id': LongTensor B*F, 'value': FloatTensor B*F}
        :return:    linear transform of x
        """
        wights = self.weight(x['id'])
        linear = []
        for i in range(self.noutput):
            a_i = wights[:, :, i]
            a_i_mul_b_i = torch.mul(a_i, x['value'])
            linear.append(a_i_mul_b_i)
        linear = torch.stack(linear, dim=2)

        val = torch.sum(linear, dim=1) + self.bias

        return val

The following models can be adjusted in this way: lr, dnn, afn, gc_arm, dcn, cin, nfm, xdfm, ipnn, kpnn, wd, gat, gcn, dcn+, sa_glu.

3) It is not quite clear how to add support for multiple output neurons for the following models: dfm, fm, hofm, afm because of it is unclear how to modify FactorizationMachine for multiple output.

Could you please comment on the correctness of changing the Linear layer from (2), and how to add support for multiple output neurons for models from (3)?

solopku commented 2 years ago

Hi GillianGrayson,

In our implementation, supporting multiple outputs for multiclass classification tasks is rather easy.

For models using an MLP layer to produce the final prediction, you can achieve so by specifying the number of outputs via the noutput parameter in the MLP layer. This applies to dnn, afn, gc_arm, dcn, cin, nfm, xdfm, ipnn, kpnn, wd, gat, gcn, dcn+, sa_glu. E.g., for gcn, you should update line 57 with noutput specified for MLP, and then directly return y in line 75 without calling the squeeze function.

Comments on (2): It is recommended to use the provided MLP layer, which should be much more efficient than using the for loop.

Comments on (3): Actually, models such as dfm, fm, hofm, afm are originally proposed for binary classification tasks (pls refer to respective papers). Nonetheless, we can also extend these models to support multiclass classification, by simply using our implementation of the FactorizationMachine layer, initializing it with reduce_dim=False, and then feeding the output of this layer to another MLP layer with ninput=nemb, and noutput=nclass to produce the final predictions.

Hope the above can address your concerns :)

GillianGrayson commented 2 years ago

Thank you a lot for a very informative and quick response!