`embed_data` appears to use wrong modality for shifted values, destroying continuous information at C→D transitions

[keesjandevries]

Summary

In embed_data (model.py:217-223), the autoregressive value shift and the modality shift are inconsistent. Values are shifted right by 1, but the modality sequence uses the current position's modality instead of the previous position's modality. This causes continuous values to be embedded via disc_embedding (with clamping) whenever a continuous dimension is followed by a categorical dimension, permanently destroying the continuous information.

Details

x_shifted = torch.cat([
    torch.zeros_like(x_in[:, :1]), x_in[:, :-1]   # values shifted right by 1
], dim=-1)
modality_shifted = torch.cat([
    torch.full_like(x_modality[:, :1], MODALITY_CATEGORICAL, dtype=torch.long),
    x_modality[:, 1:]                               # NOT shifted — uses current position's modality
], dim=-1)

At position j, x_shifted[j] holds the value from position j-1, but modality_shifted[j] holds the modality of position j (not j-1). The embedding function is then chosen based on this mismatched modality:

modality_mask = (modality_shifted == modality) & torch.isfinite(x_shifted)
x[modality_mask] = embedding_fn(x_shifted[modality_mask])

When position j-1 is continuous and position j is categorical, disc_embedding is applied to a continuous value. Because disc_embedding clamps to [0, max_num_classes - 1] and casts to long, all continuous values outside that small integer range map to the same embedding, irreversibly losing information.

Reproduction

Given a flat sequence with alternating modalities (e.g., [continuous, categorical, continuous, categorical]), change a continuous value at a position followed by a categorical position. The transformer hidden states will be identical before and after the change, because both values clamp to the same disc_embedding index.

# Example: amount=99.95 and amount=999999.0 both produce disc_embedding(3)
# when max_num_classes=4, because int(99.95)=99 and int(999999)=999999
# both clamp to max_num_classes - 1 = 3.

Expected behavior

The modality used for embedding should match the value's modality (the previous position), not the current position. This would be:

modality_shifted = torch.cat([
    torch.full_like(x_modality[:, :1], MODALITY_CATEGORICAL, dtype=torch.long),
    x_modality[:, :-1]   # shifted right by 1, matching x_shifted
], dim=-1)

This ensures cont_embedding is applied to continuous values and disc_embedding to categorical values, preserving information regardless of modality transitions.

[fdraxler]

I checked the code and can confirm that this is not as intended. As far as I can see, the predictions and likelihood computations were not affected; I will double check what version of the code we ran for the experiments in the paper. We might actually have underperformed due to the broken embedding.