SAS Deep Learning - "Floating-point exception" Error with MultiHeadAttention Layer

[Gaozhe]

I'm encountering an error when using SAS's deep learning capabilities. I've created a sequential model with a MultiHeadAttention layer for a regression task with a dataset of 1024 rows and 5 columns (4 features and 1 target column).

My model architecture is:

model = Sequential(conn, model_table='MHA_Regression')

# Input layer
model.add(InputLayer(n_channels=4, name='input_layer'))

# Dense layer
model.add(Dense(n=32, act='relu', name='dense2'))

# Multi-head attention layer
model.add(MultiHeadAttention(n=32, n_attn_heads=2, name='mha_layer'))

# Output layers
model.add(Dense(n=16, act='relu', name='pre_output'))
model.add(OutputLayer(n=1, act='IDENTITY', name='output_layer'))

When I try to fit the model, I receive this error:
ERROR: A floating-point exception or access violation has occurred, halting the analysis. This condition is usually caused by improperly scaled inputs, a large learning rate, or exploding gradients. ERROR: The action stopped due to errors.

I have tried to normalize my data and add Reshape layer, but it doesn't work. How to solve the problem?