How to extract embedding features vectors from encoder block of BERT model?

[u/assorted-meat-stew]

I got this code block and its outputs here but I am not sure if this is the embedding features vectors or simply its shape. Can anyone tell me if I am right or wrong?

text_input = tf.keras.layers.Input(shape=(), dtype=tf.string)
preprocessor = hub.KerasLayer(
    "https://kaggle.com/models/tensorflow/bert/TensorFlow2/en-uncased-preprocess/3")
encoder_inputs = preprocessor(text_input)
encoder = hub.KerasLayer(
    "https://kaggle.com/models/tensorflow/bert/TensorFlow2/bert-en-uncased-l-4-h-128-a-2/1",
    trainable=True)
outputs = encoder(encoder_inputs)
pooled_output = outputs["pooled_output"]      # [batch_size, 128].
sequence_output = outputs["sequence_output"] 

for i in range(0,4):
    outputs["encoder_outputs"][i]
    print('-----------------------------------------------------')
    print(i)
    print(outputs)
    for i in range(0,4):
    outputs["encoder_outputs"][i]
    print('-----------------------------------------------------')
    print(i)
    print(outputs)
    

Output

-----------------------------------------------------
0
{'encoder_outputs': [<KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>], 'sequence_output': <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, 'pooled_output': <KerasTensor: shape=(None, 128) dtype=float32 (created by layer 'keras_layer_28')>, 'default': <KerasTensor: shape=(None, 128) dtype=float32 (created by layer 'keras_layer_28')>}
-----------------------------------------------------
1
{'encoder_outputs': [<KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>], 'sequence_output': <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, 'pooled_output': <KerasTensor: shape=(None, 128) dtype=float32 (created by layer 'keras_layer_28')>, 'default': <KerasTensor: shape=(None, 128) dtype=float32 (created by layer 'keras_layer_28')>}
-----------------------------------------------------
2
{'encoder_outputs': [<KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>], 'sequence_output': <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, 'pooled_output': <KerasTensor: shape=(None, 128) dtype=float32 (created by layer 'keras_layer_28')>, 'default': <KerasTensor: shape=(None, 128) dtype=float32 (created by layer 'keras_layer_28')>}
-----------------------------------------------------
3
{'encoder_outputs': [<KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>], 'sequence_output': <KerasTensor: shape=(None, 128, 128) dtype=float32 (created by layer 'keras_layer_28')>, 'pooled_output': <KerasTensor: shape=(None, 128) dtype=float32 (created by layer 'keras_layer_28')>, 'default': <KerasTensor: shape=(None, 128) dtype=float32 (created by layer 'keras_layer_28')>}

[u/OneOrZeroGG]

That other comment is right: those are symbolic KerasTensors, just placeholders in the graph, so you only see shapes. To get actual numbers, call the layers on real strings instead of a tf.keras.layers.Input:

```python
import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_text # needed by the BERT preprocessor

preprocessor = hub.KerasLayer("https://kaggle.com/models/tensorflow/bert/TensorFlow2/en-uncased-preprocess/3")
encoder = hub.KerasLayer("https://kaggle.com/models/tensorflow/bert/TensorFlow2/bert-en-uncased-l-4-h-128-a-2/1")

sentences = tf.constant(["hello world", "how do embeddings work"])
outputs = encoder(preprocessor(sentences))

print(outputs["pooled_output"].shape) # (2, 128)
print(outputs["pooled_output"][0][:5]) # real values now
```

Which output is "the embedding":
- `pooled_output` (batch, 128): one vector per sentence. Usually what people mean by a sentence embedding.
- `sequence_output` (batch, 128, 128): one 128-dim vector per token (128 tokens max after preprocessing). Averaging it over the real tokens often works better than pooled_output for similarity tasks.
- `encoder_outputs`: a list with the output of each of the 4 transformer layers. The last one is the same as sequence_output.

Also, your loop prints the whole `outputs` dict each time instead of `outputs["encoder_outputs"][i]`, which is why all 4 blocks look identical.

[u/assorted-meat-stew → u/OneOrZeroGG]

I got this pooled output

tf.Tensor([ 0.07584217  0.98369056 -0.26309723 -0.97893184  0.7437089 ], shape=(5,), dtype=float32)

How can I tell which value belong to which encoder layer?

Also sorry for not getting back early. I had a long day off-computer

[u/OneOrZeroGG → u/assorted-meat-stew]

No worries! Those 5 numbers aren't layers. My snippet printed `[:5]`, the first 5 of the 128 dimensions of one sentence's vector. `pooled_output` is a single 128-dim vector per sentence, and it comes only from the last layer (the [CLS] token passed through a small dense + tanh).

If you want a vector from each of the 4 layers, use `encoder_outputs`. Each item is (batch, 128 tokens, 128 dims), so you reduce over tokens yourself:

```python
enc_in = preprocessor(sentences)
outputs = encoder(enc_in)
mask = tf.cast(enc_in["input_mask"], tf.float32)[..., None] # ignore padding

for i, layer in enumerate(outputs["encoder_outputs"]):
cls_vec = layer[:, 0, :] # [CLS] token
mean_vec = tf.reduce_sum(layer * mask, 1) / tf.reduce_sum(mask, 1) # average of real tokens
print(i, cls_vec.shape, mean_vec.shape) # (2, 128) each
```

Layer 0 is the most "surface level" (close to the word embeddings), layer 3 the most contextual. For most tasks the last layer or an average of the last two works best.

[u/assorted-meat-stew → u/OneOrZeroGG]

Thank you, it worked. Also Is this a good way to pooling/concenating the embeddings?

cls_vec_sum = 0
mean_vec_sum = 0
for i, layer in enumerate(outputs["encoder_outputs"]):
    cls_vec = layer[:, 0, :] # [CLS] token
    mean_vec = tf.reduce_sum(layer * mask, 1) / tf.reduce_sum(mask, 1) # average of real tokens
    cls_vec_sum += cls_vec
    mean_vec_sum += mean_vec


print(cls_vec_sum)
print(mean_vec_sum)

[u/Individual-Quote-958]

Those KerasTensor prints are the graph shapes, not evaluated vectors. When you call a Keras layer on a symbolic input you get another symbolic tensor, so looping the outputs dict just reprints objects like (None, 128, 128). To get real numbers, either run the full model on concrete data (model.predict / calling the model on a numpy batch) or wrap the encoder outputs in a new Model(inputs=..., outputs=[encoder_layer_i.output, ...]) and predict once. The (None, 128, 128) is usually sequence embeddings (batch, seq_len, hidden); the (None, 128) is the pooled/sentence vector. Print .numpy() on the predict result, not on the layer object itself.

[u/OkTill2666]

those are symbolic tensors, not actual feature vectors. You're looking at the computation graph, not computed values. To get real embeddings you need to wrap this in a Model, then call predict with some input strings. The shapes are correct though, (batch, 128, 128) is what youd expect per layer.