How to fix embedding of zero index to zero vector in tensorflow?

user_2179098 0
Viewed 311

I want to fix zero index embedding vector to zero vector.

For example,

embedding layer input: [0, 1, 2, 0, 3]
embedding layer output:
[
  [0, 0, 0, 0, 0],
  [ (some vector) ],
  [ (some vector) ],
  [0, 0, 0, 0, 0],
  [ (some vector) ]
]

I found mask_zero option in Keras Embedding layer, but it's not one I want...

FYI, my purpose to do this is:

  • Input: Tensor → shape (batch_size, input_size)
  • Look-up embedding using the input tensor → shape (batch_size, input_size, embedding_size)
  • Do sum pooling for axis 1 → shape (batch_size, 1, embedding_size)

How can I implement this? I tried to mask the input tensor using tf.ragged.boolean_mask, but it is slower and I think it is overkill.

keras tensorflow deep-learning machine-learning python
1 Answers

You can do this by embedding a simple function into a Lambda Layer.

def switch_layer(inputs):
    
    inp, emb = inputs
    zeros = tf.zeros_like(inp)
    ones = tf.ones_like(inp)
    
    inp = tf.keras.backend.switch(inp > 0, ones, zeros)
    inp = tf.expand_dims(inp, -1)
    
    return inp * emb

Here the full model:

embedding_size = 4
X = np.asarray([[0, 1, 2, 0, 3]])

inp = Input((None,))
x = Embedding(4, embedding_size)(inp)
x = Lambda(switch_layer)([inp, x])

model = Model(inp, x)

pred = model.predict(X)

The 'switch function' receives as input the original integer sequences and the embedding mapping. It switches the embedding vectors into zeros where the original input integers are equal to 0 otherwise it does nothing.

As you can see, the final output (pred) is a tensor of shape (batch_size, input_size, embedding_size) with vectors equal to zeros in the positions where the original input are 0.

user_103750490
Related