Input dimension mismatch when using categorical variables with >2 classes

[mahmoudibrahim98]

Issue

When using categorical variables with more than 2 classes, there's a dimension mismatch between the expected input size and actual input size after one-hot encoding. This occurs because the current channel assumes all categorical variables are binary (2 classes) by simply doubling the number of categorical columns. The current calculation doesn't account for the expanded dimensions created by one-hot encoding categorical variables.

Current Behavior

The model initialization uses:
input_channels = dataset.channels + len(dataset.categorical_cols)

This leads to errors when categorical variables have more than 2 classes. For example, with:

• 4 numerical features
• 6 binary categorical features (2 classes each)
• 1 categorical feature with 4 classes

The code calculates:

• input_channels = 11 + 7 = 18

But after one-hot encoding, the actual input dimension becomes:

• 4 (numerical) + (6 × 2) + (1 × 4) = 20 dimensions

This causes the error:
RuntimeError: input.size(-1) must be equal to input_size. Expected 18, got 20

Solution

Calculate the total channels by summing the numerical features and the total number of classes across all categorical features:

total_numerical = len(COLUMNS_DICT[data_name]["numerical"])
total_categorical = sum(COLUMNS_DICT[data_name]["categorical_num_classes"])
total_channels = total_numerical + total_categorical
model = RNN(
input_channels = total_channels, # Will account for all classes after one-hot encoding
hidden_channels = total_channels 4 if data_name in ["stock", "energy", "mimiciv", "mimiciii", "hirid"] else 256,
output_channels = total_channels,
...
)
diffusion = MixedDiffusion(
model = model,
channels = total_channels,
...
)

[MuhangTian]

Thank you for your interest in our work! Yes, we should include more flexibility in our code to allow for this, thank you for point this out. I will add it to our repository!

[mahmoudibrahim98]

I noticed as well that the current implementation produces inconsistent output dimensions when sampling data with multi-class (>2) categorical variables. While the input data has the correct number of channels, the sampled output has additional channels due to how categorical variables are handled.

COLUMNS_DICT = { "data_name": { "numerical": [0, 2, 4, 6], # 4 numerical features "categorical": [1, 3, 5, 7, 8, 9, 10, 11], # 8 categorical features "categorical_num_classes": [2, 2, 2, 2, 2, 2, 4, 4] # 6 binary + 2 four-class features } }

• Input data shape: (batch_size, 12, seq_length) # 4 numerical + 8 categorical
• Sampled output shape: (batch_size, 16, seq_length) # 4 extra channels

The issue stems from how categorical_num_classes_expanded is created and used in sampling:

categorical_num_classes_expanded = np.concatenate([ self.categorical_num_classes[i].repeat(self.categorical_num_classes[i]) for i in range(len(self.categorical_num_classes)) ])

This creates an expanded representation:

• Each binary (2-class) variable remains as 1 channel
• Each 4-class variable expands to 4 channels
• Result: 4 (numerical) + 6 (binary) + 2×4 (four-class) = 16 channels

The model internally handles this correctly during training, but the sampling process doesn't properly collapse the multi-class variables back to single channels.

Impact

• Training works correctly
• Sampling produces more channels than expected
• The extra channels contain redundant information from multi-class variables

Current Workaround

The model still functions, but users need to be aware that sampled data will have more channels than input data when using categorical variables with more than 2 classes.

Would you like me to proceed with the fix PR?

[MuhangTian]

Thank you for raising the issue. The expanded representation is to be expected. I checked the current code and believe it should collapse the multi-class variables properly. In particular, categorical_num_classes_expanded = np.concatenate([ self.categorical_num_classes[i].repeat(self.categorical_num_classes[i]) for i in range(len(self.categorical_num_classes)) ]) should produce 2 channels for binary (2-class) variables and 4 channels for 4-class variables, and the function ohe_to_categories should collapse the channels properly by taking argmax on the output, indexed by self.categorical_num_classes. May you provide more information on how this issue is caused?

[mahmoudibrahim98]

The problem is in the function create_samples that returns sampled tensor with extra empty channels.

diffusion.sample() should collapse each categorical feature back to one channel, yielding (B, num_numerical + num_categorical, seq_length) but instead, create_samples does self.channels − len(categorical_features_indices)

The fix is to replace this line

samples = torch.zeros((B, self.channels - len(self.categorical_features_indices), L))

with

samples = torch.zeros((numerical_samples.shape[0],
                       self.num_numerical_features
                       + len(self.categorical_features_indices),
                       self.seq_length))

[MuhangTian]

Thank you for pointing this out. I see your point, but I believe the current suggested edit alone does not solve this problem. In particular, self.channels in create_samples function is equal to dataset.channels + len(dataset.categorical_cols) (see etdiff_train.py). This means self.channels - len(self.categorical_features_indices) equals to dataset.channels, which is the same as self.num_numerical_features + len(self.categorical_features_indices) (this also equals to dataset.channels).

Therefore, unless there is something I am misunderstanding here, the two implementations are the same. Still, thank you for pointing out this issue.

While making the code work for >2 classes is straightforward, after some consideration, I believe it is better to leave the code as it is, since the code is mainly responsible for reproducing the results used in our paper, and it is best to keep the code as it was when the paper was accepted for reproducibility purposes. I can leave the issue open so that others can see about this problem when they want to apply TimeDiff to >2 classes. There are definitely multiple ways to adapt the current code to >2 classes (most likely involving counting the number of channels expanded correctly), and we can leave this task to future users.