Part 2: Adapting ResNet for Tabular Data

import torch
import torch.nn as nn
import torch.nn.functional as F

class TabularResidualBlock(nn.Module):
    """
    Residual block for tabular data using fully connected layers.

    Implements the core ResNet skip connection (H(x) = F(x) + x) adapted for tabular data:
    - Uses Linear layers (tabular data has no spatial structure like images)
    - Uses BatchNorm1d for 1D feature vectors
    - Adds dropout for regularization (common in tabular deep learning)

    Architecture: x -> [Linear -> BN1d -> ReLU -> Dropout -> Linear -> BN1d] -> + x -> ReLU -> Dropout
    """
    def __init__(self, d_model, dropout=0.1):
        """
        Args:
            d_model: Feature dimension (must be same for input/output)
            dropout: Dropout probability for regularization
        """
        super().__init__()

        # Two fully connected layers (analogous to two convs in image ResNet)
        self.fc1 = nn.Linear(d_model, d_model)
        self.bn1 = nn.BatchNorm1d(d_model)  # BatchNorm1d for tabular (not 2d for images)
        self.dropout1 = nn.Dropout(dropout)

        self.fc2 = nn.Linear(d_model, d_model)
        self.bn2 = nn.BatchNorm1d(d_model)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x):
        """
        Args:
            x: (batch_size, d_model) tensor of features
        Returns:
            (batch_size, d_model) tensor after residual transformation
        """
        # Save input for skip connection
        residual = x

        # Main path: F(x) = fc2(dropout(relu(bn(fc1(x)))))
        out = self.fc1(x)
        out = self.bn1(out)
        out = F.relu(out)
        out = self.dropout1(out)  # Regularization

        out = self.fc2(out)
        out = self.bn2(out)

        # Skip connection: F(x) + x (the key innovation)
        out = out + residual
        out = F.relu(out)
        out = self.dropout2(out)

        return out

# Test tabular residual block
block = TabularResidualBlock(d_model=128)
x = torch.randn(32, 128)  # Batch of 32 samples, 128 features
output = block(x)
print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")
print("Tabular residual block works! ✓")

class TabularResNet(nn.Module):
    """
    ResNet architecture for tabular data with categorical embeddings.

    Suitable for:
    - High-dimensional tabular data (e.g., OCSF with 300+ fields)
    - Mixed categorical and numerical features
    - Self-supervised learning on unlabelled data
    - Embedding extraction for anomaly detection
    """
    def __init__(
        self,
        num_numerical_features,
        categorical_cardinalities,  # List of unique values per categorical feature
        d_model=256,
        num_blocks=6,
        dropout=0.1,
        num_classes=None,  # None for embedding-only mode
    ):
        """
        Args:
            num_numerical_features: Number of continuous/numerical features
                (e.g., 50 for network_bytes_in, duration, etc.)
                ⚠️ Data-determined, not a hyperparameter to tune

            categorical_cardinalities: List of unique values per categorical feature
                (e.g., [100, 50, 200, 1000] for user_id, status_id, entity_id, etc.)
                Length of list = number of categorical features
                ⚠️ Data-determined, not a hyperparameter to tune

            d_model: Internal hidden dimension (typically 128-512)
                ✓ HYPERPARAMETER - Tune this
                - Larger = more capacity, more parameters
                - Common values: 128, 256, 512
                - This is your final embedding dimension
                - Try: [128, 256, 512] and pick based on validation performance

            num_blocks: Number of stacked residual blocks (typically 4-12)
                ✓ HYPERPARAMETER - Tune this
                - More blocks = deeper network, more capacity
                - Diminishing returns beyond 8-10 blocks for most tabular data
                - Start with 6, increase if underfitting
                - Try: [4, 6, 8, 10] and monitor validation loss

            dropout: Dropout probability for regularization (typically 0.1-0.3)
                ✓ HYPERPARAMETER - Tune this
                - Higher = more regularization, less overfitting
                - 0.1 = light, 0.2 = moderate, 0.3 = heavy
                - Try: [0.1, 0.2, 0.3] based on overfitting behavior

            num_classes: Number of output classes for supervised learning
                ⚠️ Task-determined, not a hyperparameter to tune
                - None = embedding-only mode (for anomaly detection)
                - int = classification mode (e.g., 10 for 10-class problem)
        """
        super().__init__()

        self.num_numerical = num_numerical_features
        self.num_categorical = len(categorical_cardinalities)

        # Categorical feature embeddings
        # Note: d_model // 4 is a fixed ratio (could be a hyperparameter in advanced setups)
        # Lower-dimensional embeddings for categorical features is common practice
        self.embeddings = nn.ModuleList([
            nn.Embedding(cardinality, d_model // 4)
            for cardinality in categorical_cardinalities
        ])

        # Numerical feature projection
        # Note: d_model // 2 is a fixed ratio (could be a hyperparameter in advanced setups)
        self.numerical_projection = nn.Linear(num_numerical_features, d_model // 2)

        # Total input dimension after embeddings
        embedding_dim = (d_model // 4) * len(categorical_cardinalities) + (d_model // 2)

        # Initial projection to d_model
        self.initial_projection = nn.Linear(embedding_dim, d_model)
        self.initial_bn = nn.BatchNorm1d(d_model)

        # Residual blocks
        self.blocks = nn.ModuleList([
            TabularResidualBlock(d_model, dropout)
            for _ in range(num_blocks)
        ])

        # Final layer norm (for embedding extraction)
        self.final_norm = nn.LayerNorm(d_model)

        # Optional classification head
        self.classifier = None
        if num_classes is not None:
            self.classifier = nn.Sequential(
                nn.Linear(d_model, d_model // 2),
                nn.ReLU(),
                nn.Dropout(dropout),
                nn.Linear(d_model // 2, num_classes)
            )

    def forward(self, numerical_features, categorical_features, return_embedding=False):
        """
        Args:
            numerical_features: (batch_size, num_numerical) tensor
            categorical_features: (batch_size, num_categorical) tensor of indices
            return_embedding: If True, return embedding vector instead of classification

        Returns:
            If return_embedding=True: (batch_size, d_model) embedding tensor
            Otherwise: (batch_size, num_classes) logits (if classifier exists)
        """
        batch_size = numerical_features.size(0)

        # Process categorical features through embeddings
        cat_embeddings = []
        for i, embedding_layer in enumerate(self.embeddings):
            cat_embeddings.append(embedding_layer(categorical_features[:, i]))

        if cat_embeddings:
            cat_embed = torch.cat(cat_embeddings, dim=1)
        else:
            cat_embed = torch.empty(batch_size, 0)

        # Process numerical features
        num_embed = self.numerical_projection(numerical_features)

        # Concatenate all features
        x = torch.cat([cat_embed, num_embed], dim=1)

        # Initial projection
        x = self.initial_projection(x)
        x = self.initial_bn(x)
        x = F.relu(x)

        # Residual blocks
        for block in self.blocks:
            x = block(x)

        # Final normalization
        embedding = self.final_norm(x)

        # Return embedding or classification
        if return_embedding or self.classifier is None:
            return embedding
        else:
            return self.classifier(embedding)

# Example usage
# Numerical: network_bytes_in, duration, etc.
num_numerical = 50
# Categorical: user_id, status_id, entity_id, etc.
categorical_cardinalities = [100, 50, 200, 1000]

model = TabularResNet(
    num_numerical_features=num_numerical,
    categorical_cardinalities=categorical_cardinalities,
    d_model=256,
    num_blocks=6,
    num_classes=None  # Embedding mode for anomaly detection
)

# Create dummy data
batch_size = 32
numerical_data = torch.randn(batch_size, num_numerical)
categorical_data = torch.randint(0, 50, (batch_size, len(categorical_cardinalities)))

# Get embeddings
embeddings = model(numerical_data, categorical_data, return_embedding=True)
print(f"\nTabular ResNet Test:")
print(f"Numerical input shape: {numerical_data.shape}")
print(f"Categorical input shape: {categorical_data.shape}")
print(f"Embedding output shape: {embeddings.shape}")
print(f"Total parameters: {sum(p.numel() for p in model.parameters()):,}")