Sku Side Feature Problem

Hello,

While training the retrieval-based model, when I add side features, the model stops learning and I get much worse results than just id-based information. I use information such as cr, ctr, product score from sku features.

I forward my features to dense layer --> get embedding
Then i concat my sku embeddings and sku dense embedings --> forward tfrs.task

My features like this

SKU_INPUT_COLS = {}

cols = [
"customerreviewscore"  
]



for col in cols:
    SKU_INPUT_COLS[col] = {
                            "value":[1,2,3,4],
                            "type":"bucket",
                            "dtype":tf.int32
                        }

My dense feature implemantation layer

import numpy as np
import pandas as pd
import tensorflow as tf

# create dense layer
def get_dense_layer(INPUT_COLS, size=8, lr=0.1):    

    # input layer
    inputs = {}
    for col, value  in INPUT_COLS.items():
        inputs[col] = tf.keras.layers.Input(name=col, 
                                            shape=(), 
                                            dtype=value["dtype"])


    # feature columns
    feature_columns = {}
    for col, value in INPUT_COLS.items():

        if value["type"] == "categorical":
            categorical_col = tf.feature_column.categorical_column_with_vocabulary_list(col, vocabulary_list=value["value"])
            input_layer = tf.feature_column.indicator_column(categorical_col)
            
        elif value["type"] == "bucket":
            numeric_col = tf.feature_column.numeric_column(col)
            input_layer = tf.feature_column.bucketized_column(numeric_col, boundaries=value["value"])
            
        else:
            input_layer = tf.feature_column.numeric_column(col)

        feature_columns[col] = input_layer


    # the constructor for DenseFeatures takes a list of numeric columns
    dnn_inputs = tf.keras.layers.DenseFeatures(feature_columns.values())(inputs)
    outputs = tf.keras.layers.Dense(size,
                                    activation=tf.nn.relu,
                                    input_shape=(len(INPUT_COLS),))(dnn_inputs)


    dense_model = tf.keras.Model(inputs=inputs, outputs=outputs)
    dense_model.compile(optimizer=tf.keras.optimizers.Adagrad(lr))
    return dense_model

My retriveval model code

from typing import Dict, Text
import tensorflow as tf
import tensorflow_recommenders as tfrs


class MyRetreivalModel(tfrs.Model):
    # We derive from a custom base class to help reduce boilerplate. Under the hood,
    # these are still plain Keras Models.

    def __init__(
            self,
            user_model: tf.keras.Model, # just embedding layer
            sku_model: tf.keras.Model, # just embedding layer
            sku_dense_model: tf.keras.Model,
            task: tfrs.tasks.Retrieval,
            use_candidate_sampling_probability: bool):
        super().__init__()

        
        # sku_dense_layers.
        self.sku_dense_layers = tf.keras.Sequential()
        self.sku_dense_layers.add(tf.keras.layers.Dense(64, activation="relu"))
        self.sku_dense_layers.add(tf.keras.layers.BatchNormalization())
        self.sku_dense_layers.add(tf.keras.layers.Dense(64))
        

        # Set up a retrieval task.
        self.task = task
        self.use_candidate_sampling_probability =  use_candidate_sampling_probability


    def compute_loss(self, features: Dict[Text, tf.Tensor], training=False) -> tf.Tensor:
        
        # Define how the loss is computed.
        user_embeddings = self.user_model(features["userid"])
        sku_embeddings = self.sku_model(features["sku"])

        
        #### sku
        # concat embeddings
        if self.sku_dense_model is not None:
            sku_dense_embeddings = self.sku_dense_model(features)
            sku_embeddings = tf.concat([
                                            sku_embeddings,
                                            sku_dense_embeddings,
                                        ],  axis=1)
            sku_embeddings = self.sku_dense_layers(sku_embeddings)
        

        
        ## task
        return self.task(user_embeddings,
                                   sku_embeddings,
                                   candidate_sampling_probability=features['candidate_sampling_probabilities'])

How should I proceed here?

Start by reproducing the reported behavior from get_dense_layer and MyRetreivalModel.compute_loss, comparing the id-only model with the version that bucketizes and concatenates SKU features. Inspect the feature inputs, embedding dimensions, and retrieval task call shown in the issue. Done would require identifying the cause of the degraded learning and verifying a documented fix against the baseline.