Skip to content

Fix per-feature OOV index shift in FeatureEmbeddingLayer #613

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
mkolodner-sc wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Fix per-feature OOV index shift in FeatureEmbeddingLayer #613

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 8 additions & 8 deletions gigl/src/common/models/pyg/nn/models/feature_embedding.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,13 +68,12 @@ def __init__(
self.__padding_idx = padding_idx
self.__feature_padding_value_map = feature_padding_value_map

# whether to add 1 to the whole tensor, so all elements in tensor is >= 0 for nn.Embedding input
# Since tft.compute_and_apply_vocabulary will be 0 based and use -1 as OOV padding
self.__plus_one = False
self.__oov_idx: Optional[int] = None
assert oov_idx is None or oov_idx >= -1, "oov_idx has to be >= -1"
if oov_idx and oov_idx == -1:
self.__plus_one = True
# Per-feature flag: shift indices by +1 when int_domain.min == -1, which is the
# single-OOV-bucket case where tft.compute_and_apply_vocabulary uses -1 for OOV.
# When num_oov_buckets > 1, TFT assigns OOV indices starting at vocab_size (min == 0),
# so no shift is needed — applying +1 would push the max OOV index out of bounds.
self.__feature_plus_one: dict[str, bool] = {}

for feature_name, emb_dim in features_to_embed.items():
feat_dim = get_feature_len_from_fixed_len_feature(
Expand All @@ -99,6 +98,7 @@ def __init__(
"If int_domain.min_value is -1, oov_idx must also be -1"
)
vocab_size = feat_schema.int_domain.max - feat_schema.int_domain.min + 1
self.__feature_plus_one[feature_name] = feat_schema.int_domain.min == -1

feature_padding_idx: Optional[int]
if (
Expand All @@ -113,7 +113,7 @@ def __init__(
].index(feature_padding_value)
else:
feature_padding_idx = self.__padding_idx
if self.__plus_one and feature_padding_idx is not None:
if self.__feature_plus_one[feature_name] and feature_padding_idx is not None:
feature_padding_idx = feature_padding_idx + 1
self.__feature_embedding_layers[feature_name] = nn.Embedding(
num_embeddings=vocab_size,
Expand All @@ -133,7 +133,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
x_to_emb = filter_features(
feature_schema=self.__feature_schema, feature_names=[feature], x=x
).long() # embedding layer takes LongTensor
if self.__plus_one:
if self.__feature_plus_one[feature]:
x_to_emb = x_to_emb + 1
emb_layer = self.__feature_embedding_layers[feature]
emb = emb_layer(x_to_emb)
Expand Down