add AutoInt model class

AnFreTh · AnFreTh · commit 2bba25912fd7 · 2025-03-09T20:51:44.000+01:00
diff --git a/README.md b/README.md
@@ -76,7 +76,9 @@ Mambular is a Python package that brings the power of advanced deep learning arc
 | `TabulaRNN`      | A Recurrent Neural Network for Tabular data, introduced [here](https://arxiv.org/pdf/2411.17207).                                                   |
 | `MambAttention`  | A combination between Mamba and Transformers, also introduced [here](https://arxiv.org/pdf/2411.17207).                                             |
 | `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).                                              |
+| `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).                      |
+
 
 
 
diff --git a/mambular/base_models/__init__.py b/mambular/base_models/__init__.py
@@ -10,8 +10,10 @@
 from .tabm import TabM
 from .tabtransformer import TabTransformer
 from .tabularnn import TabulaRNN
+from .autoint import AutoInt
 
 __all__ = [
+    "AutoInt",
     "MLP",
     "NDTF",
     "NODE",
diff --git a/mambular/base_models/autoint.py b/mambular/base_models/autoint.py
@@ -0,0 +1,187 @@
+import torch.nn as nn
+from ..arch_utils.layer_utils.embedding_layer import EmbeddingLayer
+from .utils.basemodel import BaseModel
+import torch.nn.init as nn_init
+import numpy as np
+from ..configs.autoint_config import DefaultAutoIntConfig
+
+
+class AutoInt(BaseModel):
+    """
+    AutoInt: Automatic Feature Interaction Learning via Self-Attentive Neural Networks.
+
+    This model uses multi-head self-attention layers to learn feature interactions for tabular data.
+    It supports key-value compression for memory efficiency and is compatible with embedding-based
+    feature encodings.
+
+    Parameters
+    ----------
+    feature_information : tuple
+        A tuple containing information about numerical features, categorical features,
+        and any additional embeddings. Expected format: `(num_feature_info, cat_feature_info, embedding_feature_info)`.
+    num_classes : int, default=1
+        Number of output classes. For regression, this should be set to `1`.
+    config : DefaultAutoIntConfig, optional
+        Configuration object containing hyperparameters such as `d_model`, `n_heads`, `n_layers`,
+        dropout rates, and compression settings.
+    **kwargs : dict
+        Additional arguments passed to the `BaseModel`.
+
+    Attributes
+    ----------
+    embedding_layer : EmbeddingLayer
+        Module that processes numerical and categorical features into embeddings.
+    kv_compression : float or None
+        The proportion of key-value compression. If `None`, no compression is applied.
+    kv_compression_sharing : str or None
+        Defines how key-value compression is shared across layers. Options:
+        - `"layerwise"`: One shared compression layer for all layers.
+        - `"headwise"`: Separate key compression per head.
+        - `"key-value"`: Separate compression layers for `k` and `v`.
+    shared_kv_compression : nn.Linear or None
+        Shared key-value compression layer, used when `kv_compression_sharing="layerwise"`.
+    layers : nn.ModuleList
+        A list of transformer-based attention layers, each consisting of:
+        - `attention`: Multi-head self-attention module.
+        - `linear`: Fully connected layer for projection.
+        - `norm0`: Layer normalization.
+    last_norm : nn.LayerNorm or None
+        Final normalization layer applied before output if `prenormalization` is enabled.
+    head : nn.Linear
+        Output layer mapping from the processed feature representation to the final predictions.
+    """
+
+    def __init__(
+        self,
+        feature_information: tuple,  # (num_feature_info, cat_feature_info, embedding_feature_info)
+        num_classes=1,
+        config: DefaultAutoIntConfig = DefaultAutoIntConfig(),  # noqa: B008
+        **kwargs,
+    ):
+        super().__init__(config=config, **kwargs)
+        self.save_hyperparameters(ignore=["feature_information"])
+        self.returns_ensemble = False
+
+        # Embedding layer
+        self.embedding_layer = EmbeddingLayer(*feature_information, config=config)
+        n_inputs = np.sum([len(info) for info in feature_information])
+
+        # Key-Value Compression
+        self.kv_compression = config.kv_compression
+        self.kv_compression_sharing = config.kv_compression_sharing
+
+        def make_kv_compression():
+            compression = nn.Linear(
+                n_inputs,
+                int(n_inputs * config.kv_compression),
+                bias=False,
+            )
+            nn_init.xavier_uniform_(compression.weight)
+            return compression
+
+        self.shared_kv_compression = (
+            make_kv_compression()
+            if self.kv_compression and self.kv_compression_sharing == "layerwise"
+            else None
+        )
+
+        # Transformer-based Interaction Layers
+        self.layers = nn.ModuleList()
+        for layer_idx in range(config.n_layers):
+            layer = nn.ModuleDict(
+                {
+                    "attention": nn.MultiheadAttention(
+                        embed_dim=config.d_model,
+                        num_heads=config.n_heads,
+                        dropout=config.attn_dropout,
+                        batch_first=True,
+                    ),
+                    "linear": nn.Linear(config.d_model, config.d_model, bias=False),
+                    "norm0": nn.LayerNorm(config.d_model),
+                }
+            )
+
+            if self.kv_compression and self.shared_kv_compression is None:
+                layer["key_compression"] = make_kv_compression()
+                if self.kv_compression_sharing == "headwise":
+                    layer["value_compression"] = make_kv_compression()
+                else:
+                    assert self.kv_compression_sharing == "key-value"
+
+            self.layers.append(layer)
+
+        # Final Normalization & Output Head
+        self.last_norm = (
+            nn.LayerNorm(config.d_model) if getattr(config, "prenorm", False) else None
+        )
+
+        self.head = nn.Linear(config.d_model * n_inputs, num_classes)
+
+    def _get_kv_compressions(self, layer):
+        """
+        Returns the correct key-value compression layers based on the sharing strategy.
+
+        Parameters
+        ----------
+        layer : nn.ModuleDict
+            The transformer layer containing possible key-value compression modules.
+
+        Returns
+        -------
+        tuple of (nn.Linear or None, nn.Linear or None)
+            The key compression and value compression layers, or `(None, None)` if no compression is applied.
+        """
+        return (
+            (self.shared_kv_compression, self.shared_kv_compression)
+            if self.shared_kv_compression is not None
+            else (
+                (layer["key_compression"], layer["value_compression"])
+                if "key_compression" in layer and "value_compression" in layer
+                else (
+                    (layer["key_compression"], layer["key_compression"])
+                    if "key_compression" in layer
+                    else (None, None)
+                )
+            )
+        )
+
+    def forward(self, *data):
+        """
+        Forward pass of the AutoInt model.
+
+        Parameters
+        ----------
+        *data : tuple
+            Input tuple of tensors containing numerical features, categorical features, and embeddings.
+
+        Returns
+        -------
+        Tensor
+            The output predictions of the model.
+        """
+        x = self.embedding_layer(*data)  # Shape: (N, J, d_model)
+
+        for layer in self.layers:
+            x_residual = x  # Store original input for residual connection
+
+            # Apply normalization before attention if prenormalization is enabled
+            x_residual = layer["norm0"](x_residual)
+
+            # Retrieve key-value compression layers
+            key_compression, value_compression = self._get_kv_compressions(layer)
+
+            # Multihead Attention
+            x_residual, _ = layer["attention"](x_residual, x_residual, x_residual)
+
+            # Apply residual connection
+            x = x + x_residual
+
+            # Apply the linear transformation
+            x_residual = layer["linear"](x)
+            x = x + x_residual  # Second residual connection
+
+        if self.last_norm:
+            x = self.last_norm(x)  # Final normalization if prenormalization is used
+
+        x = x.flatten(1)  # Flatten from (N, J, d_model) to (N, J * d_model)
+        return self.head(x)  # Final prediction
diff --git a/mambular/configs/__init__.py b/mambular/configs/__init__.py
@@ -10,9 +10,11 @@
 from .tabm_config import DefaultTabMConfig
 from .tabtransformer_config import DefaultTabTransformerConfig
 from .tabularnn_config import DefaultTabulaRNNConfig
+from .autoint_config import DefaultAutoIntConfig
 from .base_config import BaseConfig
 
 __all__ = [
+    "DefaultAutoIntConfig",
     "DefaultFTTransformerConfig",
     "DefaultMLPConfig",
     "DefaultMambAttentionConfig",
@@ -25,5 +27,5 @@
     "DefaultTabMConfig",
     "DefaultTabTransformerConfig",
     "DefaultTabulaRNNConfig",
-    "BaseConfig"
+    "BaseConfig",
 ]
diff --git a/mambular/configs/autoint_config.py b/mambular/configs/autoint_config.py
@@ -0,0 +1,49 @@
+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 DefaultAutoIntConfig(BaseConfig):
+    """Configuration class for the AutoInt model with predefined hyperparameters.
+
+    Parameters
+    ----------
+    d_model : int, default=128
+        Dimensionality of the transformer model.
+    n_layers : int, default=4
+        Number of transformer layers.
+    n_heads : int, default=8
+        Number of attention heads in the transformer.
+    attn_dropout : float, default=0.2
+        Dropout rate for the attention mechanism.
+    transformer_dim_feedforward : int, default=256
+        Dimensionality of the feed-forward layers in the transformer.
+    prenorm : bool, default=False
+        Whether to apply normalization before last layer.
+    bias : bool, default=True
+        Whether to use bias in linear layers.
+    cat_encoding : str, default="int"
+        Method for encoding categorical features ('int', 'one-hot', or 'linear').
+    kv_compression : float, default=0.5
+        Compression ratio for key-value pairs.
+    kv_compression_sharing : str, default='key-value'
+        Sharing strategy for key-value compression ('headwise', or 'key-value').
+    """
+
+    # Architecture Parameters
+    d_model: int = 128
+    n_layers: int = 4
+    n_heads: int = 8
+    attn_dropout: float = 0.2
+    fprenorm: bool = False
+    transformer_dim_feedforward: int = 256
+    bias: bool = True
+
+    use_cls: bool = False
+    cat_encoding: str = "int"
+
+    kv_compression: float = 0.5
+    kv_compression_sharing: str = "key-value"
diff --git a/mambular/models/__init__.py b/mambular/models/__init__.py
@@ -25,8 +25,12 @@
     TabTransformerRegressor,
 )
 from .tabularnn import TabulaRNNClassifier, TabulaRNNLSS, TabulaRNNRegressor
+from .autoint import AutoIntClassifier, AutoIntLSS, AutoIntRegressor
 
 __all__ = [
+    "AutoIntClassifier",
+    "AutoIntLSS",
+    "AutoIntRegressor",
     "MLPLSS",
     "NDTFLSS",
     "NODELSS",
diff --git a/mambular/models/autoint.py b/mambular/models/autoint.py
@@ -0,0 +1,64 @@
+from ..base_models.autoint import AutoInt
+from ..configs.autoint_config import DefaultAutoIntConfig
+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 AutoIntRegressor(SklearnBaseRegressor):
+    __doc__ = generate_docstring(
+        DefaultAutoIntConfig,
+        model_description="""
+        AutoInt regressor. This class extends the SklearnBaseRegressor
+        class and uses the AutoInt model with the default AutoInt
+        configuration.
+        """,
+        examples="""
+        >>> from mambular.models import AutoIntRegressor
+        >>> model = AutoIntRegressor(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=AutoInt, config=DefaultAutoIntConfig, **kwargs)
+
+
+class AutoIntClassifier(SklearnBaseClassifier):
+    __doc__ = generate_docstring(
+        DefaultAutoIntConfig,
+        """AutoInt Classifier. This class extends the SklearnBaseClassifier class
+        and uses the AutoInt model with the default AutoInt configuration.""",
+        examples="""
+        >>> from mambular.models import AutoIntClassifier
+        >>> model = AutoIntClassifier(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=AutoInt, config=DefaultAutoIntConfig, **kwargs)
+
+
+class AutoIntLSS(SklearnBaseLSS):
+    __doc__ = generate_docstring(
+        DefaultAutoIntConfig,
+        """AutoInt for distributional regression.
+        This class extends the SklearnBaseLSS class and uses the
+        AutoInt model with the default AutoInt configuration.""",
+        examples="""
+        >>> from mambular.models import AutoIntLSS
+        >>> model = AutoIntLSS(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=AutoInt, config=DefaultAutoIntConfig, **kwargs)
