diff --git a/deeptab/architectures/ft_transformer.py b/deeptab/architectures/ft_transformer.py index 6fa3a0da..49c10975 100644 --- a/deeptab/architectures/ft_transformer.py +++ b/deeptab/architectures/ft_transformer.py @@ -1,4 +1,3 @@ -import numpy as np import torch.nn as nn from deeptab.core import BaseModel @@ -85,7 +84,7 @@ def __init__( ) # pooling - n_inputs = np.sum([len(info) for info in feature_information]) + n_inputs = self.embedding_layer.n_tokens self.initialize_pooling_layers(config=config, n_inputs=n_inputs, hidden_size=self.hparams.d_model) def forward(self, *data): diff --git a/deeptab/architectures/mambattention.py b/deeptab/architectures/mambattention.py index 709b01e6..2790d864 100644 --- a/deeptab/architectures/mambattention.py +++ b/deeptab/architectures/mambattention.py @@ -1,4 +1,3 @@ -import numpy as np import torch from deeptab.core import BaseModel @@ -100,10 +99,10 @@ def __init__( ) if self.shuffle_embeddings: - self.register_buffer("perm", torch.randperm(self.embedding_layer.seq_len)) + self.register_buffer("perm", self.embedding_layer.feature_permutation()) # pooling - n_inputs = np.sum([len(info) for info in feature_information]) + n_inputs = self.embedding_layer.n_tokens self.initialize_pooling_layers(config=config, n_inputs=n_inputs, hidden_size=self.hparams.d_model) def forward(self, *data): diff --git a/deeptab/architectures/mambular.py b/deeptab/architectures/mambular.py index be1bf593..c1cca7af 100644 --- a/deeptab/architectures/mambular.py +++ b/deeptab/architectures/mambular.py @@ -1,4 +1,3 @@ -import numpy as np import torch from deeptab.core import BaseModel @@ -81,10 +80,10 @@ def __init__( ) if self.hparams.shuffle_embeddings: - self.register_buffer("perm", torch.randperm(self.embedding_layer.seq_len)) + self.register_buffer("perm", self.embedding_layer.feature_permutation()) # pooling - n_inputs = np.sum([len(info) for info in feature_information]) + n_inputs = self.embedding_layer.n_tokens self.initialize_pooling_layers(config=config, n_inputs=n_inputs, hidden_size=self.hparams.d_model) def forward(self, *data): diff --git a/deeptab/architectures/tabtransformer.py b/deeptab/architectures/tabtransformer.py index fff4b299..0a7f3258 100644 --- a/deeptab/architectures/tabtransformer.py +++ b/deeptab/architectures/tabtransformer.py @@ -1,4 +1,3 @@ -import numpy as np import torch import torch.nn as nn @@ -113,10 +112,10 @@ def __init__( ) # pooling - # Only cat/embedding features pass through the encoder whose output - # gets pooled (numerical features are concatenated in separately), - # so n_inputs excludes num_feature_info. - n_inputs = int(np.sum([len(info) for info in feature_information[1:]])) + # Only cat/embedding features (plus the CLS token, if enabled) pass + # through the encoder whose output gets pooled (numerical features are + # concatenated in separately), so n_inputs excludes num_feature_info. + n_inputs = self.embedding_layer.n_tokens self.initialize_pooling_layers(config=config, n_inputs=n_inputs, hidden_size=self.hparams.d_model) def forward(self, *data): diff --git a/deeptab/architectures/tabularnn.py b/deeptab/architectures/tabularnn.py index da8dea7e..0304628f 100644 --- a/deeptab/architectures/tabularnn.py +++ b/deeptab/architectures/tabularnn.py @@ -1,6 +1,5 @@ from dataclasses import replace -import numpy as np import torch import torch.nn as nn @@ -84,7 +83,7 @@ def __init__( self.norm_f = get_normalization_layer(temp_config) # pooling - n_inputs = int(np.sum([len(info) for info in feature_information])) + n_inputs = self.embedding_layer.n_tokens self.initialize_pooling_layers(config=config, n_inputs=n_inputs, hidden_size=self.hparams.dim_feedforward) def forward(self, *data): diff --git a/deeptab/configs/models/mambattention_config.py b/deeptab/configs/models/mambattention_config.py index c816b44e..f3524c9c 100644 --- a/deeptab/configs/models/mambattention_config.py +++ b/deeptab/configs/models/mambattention_config.py @@ -57,7 +57,8 @@ class MambAttentionConfig(BaseModelConfig): BC_layer_norm : bool, default=False Whether to apply layer normalization to B-C matrices. shuffle_embeddings : bool, default=False - Whether to shuffle embeddings before passing to Mamba layers. + Whether to shuffle the feature tokens before passing them to the Mamba layers. The CLS token, + if used, keeps its position. head_layer_sizes : list, default=[] Sizes of the fully connected layers in the model's head. head_dropout : float, default=0.5 diff --git a/deeptab/configs/models/mambular_config.py b/deeptab/configs/models/mambular_config.py index 1a978c9f..2ddf9d3e 100644 --- a/deeptab/configs/models/mambular_config.py +++ b/deeptab/configs/models/mambular_config.py @@ -51,7 +51,8 @@ class MambularConfig(BaseModelConfig): BC_layer_norm : bool, default=False Whether to use layer norm on the B and C matrices shuffle_embeddings : bool, default=False - Whether to shuffle embeddings before being passed to Mamba layers. + Whether to shuffle the feature tokens before passing them to the Mamba layers. The CLS token, + if used, keeps its position. head_layer_sizes : list, default=[] Sizes of the layers in the model's head. head_dropout : float, default=0.5 diff --git a/deeptab/nn/blocks/common.py b/deeptab/nn/blocks/common.py index 48fcb058..21595d8a 100644 --- a/deeptab/nn/blocks/common.py +++ b/deeptab/nn/blocks/common.py @@ -639,8 +639,19 @@ def __init__(self, num_feature_info, cat_feature_info, emb_feature_info, config) cat_feature_info : dict Dictionary where keys are categorical feature names and values are the number of categories for each feature. + emb_feature_info : dict + Dictionary where keys are external embedding feature names and values describe their + dimensions. config : Config Configuration object containing all required settings. + + Attributes + ---------- + seq_len : int + Number of feature tokens: one per numerical, categorical and external embedding feature. + n_tokens : int + Length of the token sequence returned by :meth:`forward`: ``seq_len``, plus one when + ``use_cls`` adds a CLS token. """ super().__init__() @@ -659,8 +670,11 @@ def __init__(self, num_feature_info, cat_feature_info, emb_feature_info, config) self.embedding_bias = getattr(config, "embedding_bias", False) self.cat_encoding = getattr(config, "cat_encoding", "int") - # Sequence length - self.seq_len = len(num_feature_info) + len(cat_feature_info) + # Sequence length: one token per numerical, categorical, and external + # embedding feature. n_tokens additionally counts the CLS token, so it + # matches the sequence length returned by forward(). + self.seq_len = len(num_feature_info) + len(cat_feature_info) + len(emb_feature_info) + self.n_tokens = self.seq_len + int(self.use_cls) # Initialize numerical embeddings based on embedding_type if self.embedding_type == "ndt": @@ -850,6 +864,25 @@ def forward(self, num_features, cat_features, emb_features): return x + def feature_permutation(self): + """Draw a random permutation of the feature tokens returned by :meth:`forward`. + + Every feature token (numerical, categorical, and external embedding + features) is shuffled, while the CLS token, when enabled, keeps its + configured position so that ``"cls"`` pooling still finds it. + + Returns + ------- + Tensor + Index tensor of length ``n_tokens`` for the sequence dimension. + """ + perm = torch.randperm(self.seq_len) + if not self.use_cls: + return perm + if self.cls_position == 0: + return torch.cat([torch.zeros(1, dtype=perm.dtype), perm + 1]) + return torch.cat([perm, torch.full((1,), self.seq_len, dtype=perm.dtype)]) + def check_plr_embedding_compatibility(self, feature_info: tuple): # List of incompatible preprocessing terms for PLR embedding incompatible_terms = ["ple", "one-hot", "polynomial", "splines", "sigmoid", "rbf"] diff --git a/docs/core_concepts/custom_models.md b/docs/core_concepts/custom_models.md index de78adcd..c1ac7d05 100644 --- a/docs/core_concepts/custom_models.md +++ b/docs/core_concepts/custom_models.md @@ -200,14 +200,14 @@ class MyEmbeddedModel(BaseModel): self.save_hyperparameters(ignore=["feature_information"]) self.embedding_layer = EmbeddingLayer(*feature_information, config=config) - n_features = sum(len(info) for info in feature_information) - input_dim = n_features * self.hparams.d_model + # n_tokens counts every token the layer returns, including a CLS token if enabled. + input_dim = self.embedding_layer.n_tokens * self.hparams.d_model self.head = nn.Linear(input_dim, num_classes) def forward(self, *data): - x = self.embedding_layer(*data) # (batch, n_features, d_model) - x = x.reshape(x.shape[0], -1) # flatten to (batch, n_features * d_model) + x = self.embedding_layer(*data) # (batch, n_tokens, d_model) + x = x.reshape(x.shape[0], -1) # flatten to (batch, n_tokens * d_model) return self.head(x) ``` diff --git a/docs/model_zoo/stable/mambattention.md b/docs/model_zoo/stable/mambattention.md index a5611e3f..c2986ded 100644 --- a/docs/model_zoo/stable/mambattention.md +++ b/docs/model_zoo/stable/mambattention.md @@ -40,6 +40,8 @@ The default config uses `d_model=64`, `n_layers=4`, `n_heads=8`, `n_attention_la With `shuffle_embeddings=True`, the token permutation is drawn once at construction and registered as the `perm` buffer. It is saved with the model and restored on load, so a reloaded model reads features in the same order it was trained with. Earlier versions did not save the permutation, so their shuffled models cannot recover the original order and fail strict loading; retrain them. +The permutation covers every feature token, including external-embedding tokens, and keeps the CLS token (`use_cls=True`) at its configured position so `"cls"` pooling still reads it. Versions up to 2.1.0 drew the permutation over the numerical and categorical tokens only, so the remaining positions were dropped: every external-embedding token and, with `use_cls=True`, either the last feature token (`cls_position=0`, which also moved the CLS token away from position 0) or the CLS token itself (`cls_position=1`). Shuffled models saved with `use_cls=True` or external embeddings by those versions fail to load because their `perm` buffer has a different length; retrain them. + > **Configuration note:** Pooling operates on width-`d_model` tokens, not TabulaRNN's `dim_feedforward` width. The local Mamba residual blocks also use the shared fixes for dilated convolution padding, missing-`mambapy` sequential fallback, and `AD_weight_decay=False` parameter markers. Final normalization after pooling is selected by `norm` and is separate from normalization inside the residual blocks. ## Practical Config diff --git a/docs/model_zoo/stable/mambular.md b/docs/model_zoo/stable/mambular.md index fdb002f9..024ebe56 100644 --- a/docs/model_zoo/stable/mambular.md +++ b/docs/model_zoo/stable/mambular.md @@ -39,6 +39,8 @@ The default config uses `d_model=64`, `n_layers=4`, `d_state=128`, `d_conv=4`, ` With `shuffle_embeddings=True`, the token permutation is drawn once at construction and registered as the `perm` buffer. It is saved with the model and restored on load, so a reloaded model reads features in the same order it was trained with. Earlier versions did not save the permutation, so their shuffled models cannot recover the original order and fail strict loading; retrain them. +The permutation covers every feature token, including external-embedding tokens, and keeps the CLS token (`use_cls=True`) at its configured position so `"cls"` pooling still reads it. Versions up to 2.1.0 drew the permutation over the numerical and categorical tokens only, so the remaining positions were dropped: every external-embedding token and, with `use_cls=True`, either the last feature token (`cls_position=0`, which also moved the CLS token away from position 0) or the CLS token itself (`cls_position=1`). Shuffled models saved with `use_cls=True` or external embeddings by those versions fail to load because their `perm` buffer has a different length; retrain them. + > **Configuration note:** In the local Mamba block, `dilation=3` with `d_conv=3` gives padding 6 in each directional convolution. Padding and the existing crop preserve sequence length; previously, the forward convolution ignored dilation. If `use_pscan=True` but `mambapy` cannot be imported, the block switches to its sequential scan rather than calling a missing function. `AD_weight_decay=False` marks `A_log` and `D` for zero weight decay; module-based optimizer construction respects those markers independently of the bias/norm exemption flag. Learnable pooling uses the sequence width `d_model`, not a `dim_feedforward` field. For `d_model=16`, attention pooling learns a width-16 weight vector. This keeps the shared pooling helper independent of TabulaRNN's recurrent width. diff --git a/tests/architectures/test_architecture_contracts.py b/tests/architectures/test_architecture_contracts.py index ec9090b5..4020cce6 100644 --- a/tests/architectures/test_architecture_contracts.py +++ b/tests/architectures/test_architecture_contracts.py @@ -199,6 +199,71 @@ def test_shuffled_embeddings_permutation_is_checkpointed(architecture_name): torch.testing.assert_close(restored(features, [], []), model(features, [], []), rtol=0, atol=0) +def _small_sequence_config(architecture_name, **options): + model_class = next(model_class for model_class in model_classes if model_class.__name__ == architecture_name) + config = get_model_config(model_class) + for name, value in {"d_model": 8, "n_layers": 1, "d_state": 4, "n_heads": 2, "dim_feedforward": 8}.items(): + if hasattr(config, name): + setattr(config, name, value) + for name, value in options.items(): + setattr(config, name, value) + return model_class, config + + +_NUM_AND_EMB_INFO = ( + {f"value_{index}": {"dimension": 1} for index in range(3)}, + {}, + {"text": {"dimension": 5}}, +) + + +@pytest.mark.parametrize("architecture_name", ["Mambular", "MambAttention"]) +@pytest.mark.parametrize( + "options, feature_info", + [ + pytest.param({"use_cls": True, "cls_position": 0}, (*_NUM_AND_EMB_INFO[:2], {}), id="cls-first"), + pytest.param({"use_cls": True, "cls_position": 1}, (*_NUM_AND_EMB_INFO[:2], {}), id="cls-last"), + pytest.param({}, _NUM_AND_EMB_INFO, id="external-embeddings"), + ], +) +def test_shuffled_embeddings_keep_every_token(architecture_name, options, feature_info): + model_class, config = _small_sequence_config(architecture_name, shuffle_embeddings=True, **options) + model = model_class(feature_info, config=config).eval() + embeddings = [torch.randn(4, 5)] if feature_info[2] else [] + tokens = model.embedding_layer([torch.randn(4, 1) for _ in range(3)], [], embeddings) + + assert model.embedding_layer.n_tokens == tokens.shape[1] + assert sorted(model.perm.tolist()) == list(range(tokens.shape[1])) + if options.get("use_cls"): + cls_index = 0 if options["cls_position"] == 0 else tokens.shape[1] - 1 + assert model.perm[cls_index].item() == cls_index + + +@pytest.mark.parametrize("architecture_name", ["Mambular", "MambAttention"]) +def test_shuffled_embeddings_use_external_embeddings(architecture_name): + model_class, config = _small_sequence_config(architecture_name, shuffle_embeddings=True) + model = model_class(_NUM_AND_EMB_INFO, config=config).eval() + features = [torch.randn(4, 1) for _ in range(3)] + with torch.no_grad(): + first = model(features, [], [torch.randn(4, 5)]) + second = model(features, [], [torch.randn(4, 5)]) + assert not torch.allclose(first, second) + + +@pytest.mark.parametrize( + "architecture_name", ["FTTransformer", "Mambular", "MambAttention", "TabTransformer", "TabulaRNN"] +) +@pytest.mark.parametrize("cls_position", [0, 1]) +def test_learned_flatten_pooling_counts_cls_token(architecture_name, cls_position): + model_class, config = _small_sequence_config( + architecture_name, use_cls=True, cls_position=cls_position, pooling_method="learned_flatten" + ) + feature_info = ({"value": {"dimension": 1}}, {"color": {"dimension": 1, "categories": 3}}, {}) + model = model_class(feature_info, num_classes=2, config=config) + output = model([torch.randn(4, 1)], [torch.randint(0, 3, (4, 1))], []) + assert output.shape == (4, 2) + + @pytest.mark.parametrize("model_type", ["mLSTM", "sLSTM"]) def test_tabularnn_recurrent_options_preserve_rows(model_type): from deeptab.architectures.tabularnn import TabulaRNN