include trompt

Author: AnFreTh

README.md

@@ -78,6 +78,7 @@ Mambular is a Python package that brings the power of advanced deep learning arc
 | `NDTF`           | A neural decision forest using soft decision trees. See [Kontschieder et al.](https://openaccess.thecvf.com/content_iccv_2015/html/Kontschieder_Deep_Neural_Decision_ICCV_2015_paper.html) for inspiration. |
 | `SAINT`          | Improve neural networs via Row Attention and Contrastive Pre-Training, introduced [here](https://arxiv.org/pdf/2106.01342).                         |
 | `AutoInt`        | Automatic Feature Interaction Learning via Self-Attentive Neural Networks introduced [here](https://arxiv.org/abs/1810.11921).                      |
+| `Trompt `        | Trompt: Towards a Better Deep Neural Network for Tabular Data introduced [here](https://arxiv.org/abs/2305.18446).                                  |
 
 
 

mambular/arch_utils/layer_utils/importance.py

@@ -0,0 +1,28 @@
+import torch.nn as nn
+import torch
+
+
+class ImportanceGetter(nn.Module):  # Figure 3 part 1
+    def __init__(self, P, C, d):
+        super().__init__()
+        self.colemb = nn.Parameter(torch.empty(C, d))
+        self.pemb = nn.Parameter(torch.empty(P, d))
+        torch.nn.init.normal_(self.colemb, std=0.01)
+        torch.nn.init.normal_(self.pemb, std=0.01)
+        self.C = C
+        self.P = P
+        self.d = d
+        self.dense = nn.Linear(2 * self.d, self.d)
+        self.laynorm1 = nn.LayerNorm(self.d)
+        self.laynorm2 = nn.LayerNorm(self.d)
+
+    def forward(self, O):
+        eprompt = self.pemb.unsqueeze(0).repeat(O.shape[0], 1, 1)
+
+        dense_out = self.dense(torch.cat((self.laynorm1(eprompt), O), dim=-1))
+
+        dense_out = dense_out + eprompt + O
+
+        ecolumn = self.laynorm2(self.colemb.unsqueeze(0).repeat(O.shape[0], 1, 1))
+
+        return torch.softmax(dense_out @ ecolumn.transpose(1, 2), dim=-1)

mambular/arch_utils/trompt_utils.py

@@ -0,0 +1,54 @@
+import torch.nn as nn
+import torch
+from .layer_utils.embedding_layer import EmbeddingLayer
+from .layer_utils.importance import ImportanceGetter
+import numpy as np
+
+
+class Expander(nn.Module):  # Figure 3 part 3
+    def __init__(self, P):
+        super().__init__()
+        self.lin = nn.Linear(1, P)
+        self.relu = nn.ReLU()
+        self.gn = nn.GroupNorm(2, P)
+
+    def forward(self, x):
+        res = self.relu(self.lin(x.unsqueeze(-1)))
+
+        return x.unsqueeze(1) + self.gn(torch.permute(res, (0, 3, 1, 2)))
+
+
+class TromptCell(nn.Module):
+    def __init__(self, feature_information, config):
+        super().__init__()
+        C = np.sum([len(info) for info in feature_information])
+        self.enc = EmbeddingLayer(
+            *feature_information,
+            config=config,
+        )
+        self.fe = ImportanceGetter(config.P, C, config.d_model)
+        self.ex = Expander(config.P)
+
+    def forward(self, *data, O=None):
+        x_res = self.ex(self.enc(*data))
+
+        M = self.fe(O)
+
+        return (M.unsqueeze(-1) * x_res).sum(dim=2)
+
+
+class TromptDecoder(nn.Module):
+    def __init__(self, d, d_out):
+        super().__init__()
+        self.l1 = nn.Linear(d, 1)
+        self.l2 = nn.Linear(d, d)
+        self.relu = nn.ReLU()
+        self.laynorm1 = nn.LayerNorm(d)
+        self.lf = nn.Linear(d, d_out)
+
+    def forward(self, x):
+        pw = torch.softmax(self.l1(x).squeeze(-1), dim=-1)
+
+        xnew = (pw.unsqueeze(-1) * x).sum(dim=-2)
+
+        return self.lf(self.laynorm1(self.relu(self.l2(xnew))))

mambular/base_models/__init__.py

@@ -11,8 +11,10 @@
 from .tabtransformer import TabTransformer
 from .tabularnn import TabulaRNN
 from .autoint import AutoInt
+from .trompt import Trompt
 
 __all__ = [
+    "Trompt",
     "AutoInt",
     "MLP",
     "NDTF",

mambular/base_models/trompt.py

@@ -0,0 +1,56 @@
+import torch.nn as nn
+import torch
+from ..arch_utils.get_norm_fn import get_normalization_layer
+from ..arch_utils.layer_utils.embedding_layer import EmbeddingLayer
+from ..configs.trompt_config import DefaultTromptConfig
+from .utils.basemodel import BaseModel
+from ..arch_utils.trompt_utils import TromptCell, TromptDecoder
+import numpy as np
+
+
+class Trompt(BaseModel):
+
+    def __init__(
+        self,
+        feature_information: tuple,  # Expecting (num_feature_info, cat_feature_info, embedding_feature_info)
+        num_classes=1,
+        config: DefaultTromptConfig = DefaultTromptConfig(),  # noqa: B008
+        **kwargs,
+    ):
+        super().__init__(config=config, **kwargs)
+        self.save_hyperparameters(ignore=["feature_information"])
+        self.returns_ensemble = True
+
+        # embedding layer
+        self.cells = nn.ModuleList(
+            TromptCell(feature_information, config) for _ in range(config.n_cycles)
+        )
+        self.decoder = TromptDecoder(config.d_model, num_classes)
+        self.init_rec = nn.Parameter(torch.empty(config.P, config.d_model))
+        self.n_cycles = config.n_cycles
+
+    def forward(self, *data):
+        """Defines the forward pass of the model.
+
+        Parameters
+        ----------
+        data : tuple
+            Input tuple of tensors of num_features, cat_features, embeddings.
+
+        Returns
+        -------
+        Tensor
+            The output predictions of the model.
+        """
+        O = self.init_rec.unsqueeze(0).repeat(data[0][0].shape[0], 1, 1)
+        outputs = []
+
+        for i in range(self.n_cycles):
+            O = self.cells[i](*data, O=O)
+            # print(O.shape)
+            # print(self.tdown(O).shape)
+            outputs.append(self.decoder(O))
+
+        out = torch.stack(outputs, dim=1).squeeze(-1)
+        # preds = out.mean(dim=1)
+        return out

mambular/configs/__init__.py

@@ -11,9 +11,11 @@
 from .tabtransformer_config import DefaultTabTransformerConfig
 from .tabularnn_config import DefaultTabulaRNNConfig
 from .autoint_config import DefaultAutoIntConfig
+from .trompt_config import DefaultTromptConfig
 from .base_config import BaseConfig
 
 __all__ = [
+    "DefaultTromptConfig",
     "DefaultAutoIntConfig",
     "DefaultFTTransformerConfig",
     "DefaultMLPConfig",

mambular/configs/trompt_config.py

@@ -0,0 +1,27 @@
+from collections.abc import Callable
+from dataclasses import dataclass, field
+import torch.nn as nn
+from ..arch_utils.transformer_utils import ReGLU
+from .base_config import BaseConfig
+
+
+@dataclass
+class DefaultTromptConfig(BaseConfig):
+    """Configuration class for the Trompt model with predefined hyperparameters.
+
+    Parameters
+    ----------
+    d_model : int, default=128
+        Dimensionality of the transformer model.
+    n_cycles : int, default=6
+        Number of cycles in the Trompt model.
+    n_cells : int, default=4
+        Number of cells in each cycle.
+    P : int, default=128
+        Number of steps in the Trompt model.
+    """
+
+    d_model: int = 128
+    n_cycles: int = 6
+    n_cells: int = 4
+    P: int = 128

mambular/models/__init__.py

@@ -26,8 +26,12 @@
 )
 from .tabularnn import TabulaRNNClassifier, TabulaRNNLSS, TabulaRNNRegressor
 from .autoint import AutoIntClassifier, AutoIntLSS, AutoIntRegressor
+from .trompt import TromptClassifier, TromptLSS, TromptRegressor
 
 __all__ = [
+    "TromptClassifier",
+    "TromptLSS",
+    "TromptRegressor",
     "AutoIntClassifier",
     "AutoIntLSS",
     "AutoIntRegressor",

mambular/models/trompt.py

@@ -0,0 +1,64 @@
+from ..base_models.trompt import Trompt
+from ..configs.trompt_config import DefaultTromptConfig
+from ..utils.docstring_generator import generate_docstring
+from .utils.sklearn_base_classifier import SklearnBaseClassifier
+from .utils.sklearn_base_lss import SklearnBaseLSS
+from .utils.sklearn_base_regressor import SklearnBaseRegressor
+
+
+class TromptRegressor(SklearnBaseRegressor):
+    __doc__ = generate_docstring(
+        DefaultTromptConfig,
+        model_description="""
+        Trompt regressor. This class extends the SklearnBaseRegressor
+        class and uses the Trompt model with the default Trompt
+        configuration.
+        """,
+        examples="""
+        >>> from mambular.models import TromptRegressor
+        >>> model = TromptRegressor(d_model=64, n_layers=8)
+        >>> model.fit(X_train, y_train)
+        >>> preds = model.predict(X_test)
+        >>> model.evaluate(X_test, y_test)
+        """,
+    )
+
+    def __init__(self, **kwargs):
+        super().__init__(model=Trompt, config=DefaultTromptConfig, **kwargs)
+
+
+class TromptClassifier(SklearnBaseClassifier):
+    __doc__ = generate_docstring(
+        DefaultTromptConfig,
+        """Trompt Classifier. This class extends the SklearnBaseClassifier class
+        and uses the Trompt model with the default Trompt configuration.""",
+        examples="""
+        >>> from mambular.models import TromptClassifier
+        >>> model = TromptClassifier(d_model=64, n_layers=8)
+        >>> model.fit(X_train, y_train)
+        >>> preds = model.predict(X_test)
+        >>> model.evaluate(X_test, y_test)
+        """,
+    )
+
+    def __init__(self, **kwargs):
+        super().__init__(model=Trompt, config=DefaultTromptConfig, **kwargs)
+
+
+class TromptLSS(SklearnBaseLSS):
+    __doc__ = generate_docstring(
+        DefaultTromptConfig,
+        """Trompt for distributional regression.
+        This class extends the SklearnBaseLSS class and uses the
+        Trompt model with the default Trompt configuration.""",
+        examples="""
+        >>> from mambular.models import TromptLSS
+        >>> model = TromptLSS(d_model=64, n_layers=8)
+        >>> model.fit(X_train, y_train, family="normal")
+        >>> preds = model.predict(X_test)
+        >>> model.evaluate(X_test, y_test)
+        """,
+    )
+
+    def __init__(self, **kwargs):
+        super().__init__(model=Trompt, config=DefaultTromptConfig, **kwargs)