# Remove subwords from BERT output

I'm trying to build a multilingual WSD system with BERT on top as the embedding layer. In order to have better performances, after BERT finishes its job (and performs Transfer Learning), I need to remove the subwords from its output. Is there a way to do so?
I've tried to detach the model from the network's architecture, doing something like this... but I need to do this as a custom layer and I'm not 100% sure that this is even right

class Bert:
def __init__(self):
input_word_ids = tf.keras.layers.Input(shape=(None,), dtype=tf.int32, name="input_word_ids")
segment_ids = tf.keras.layers.Input(shape=(None,), dtype=tf.int32, name="segment_ids")
bert = hub.KerasLayer("https://tfhub.dev/tensorflow/bert_multi_cased_L-12_H-768_A-12/1", trainable=False, name="BERT")
pooled_output, sequence_output = bert([input_word_ids, input_mask, segment_ids])
self.model = tf.keras.models.Model(inputs=[input_word_ids, input_mask, segment_ids], outputs=[pooled_output, sequence_output])
self.model.summary()

print("Starting BERT prediction...")
pool_embs, all_embs = self.model.predict(
verbose=1,
batch_size=64
)
del pool_embs
to_return = []
print("Conversion\nSoftware version 2.0...")
for i in tqdm(range(len(positional_ids))):
indexes_to_extrapolate = indexes_to_extrapolate[:63] if len(indexes_to_extrapolate) > 64 else indexes_to_extrapolate
new_version = tf.gather(all_embs[i], tf.constant(indexes_to_extrapolate))
if train_mode and new_version.shape[0] < 64:
#Means that, originally, there has to be a padding!
#And, if there is, it can surely be found in the first position of the needed_padding!
how_much_iteration = 64 - new_version.shape[0]
if how_much_iteration > 0:
for iteratore in range(how_much_iteration):
new_version = tf.concat([new_version, tf.constant(all_embs[i][tmp_padding_for_iteration], shape=(1,768))], 0)
with open("registro_shape.txt","a") as registro:
registro.write("Shape --> " +str(new_version.shape)+"\n")
if new_version.shape[0] > 64:
print("wth")
to_return.append(new_version)
return tf.stack(to_return)


EDIT: I'll try to contextualize the case with more information regarding the architecture of the network. In particular, this is the architecture of the network that I'm trying to build for the WSD task. Note that the network should perform a multitask learning task:

1. Bert
2. BiLSTM
3. Attention Layer
4. 3 outputs layer

self.tokenizatore = FullTokenizer(bert_path,do_lower_case=False)

input_word_ids = tf.keras.layers.Input(shape=(None,), dtype=tf.int32,name="input_word_ids")

segment_ids = tf.keras.layers.Input(shape=(None,), dtype=tf.int32,name="segment_ids")

print("dopwnloading BERT...")
bert = hub.KerasLayer("https://tfhub.dev/tensorflow/bert_multi_cased_L-12_H-768_A-12/1", trainable=False)
pooled_output, sequence_output = bert([input_word_ids, input_mask, segment_ids])
LSTM = tf.keras.layers.Bidirectional(
tf.keras.layers.LSTM(
units=hidden_size,
dropout=dropout,
recurrent_dropout=recurrent_dropout,
return_sequences=True,
return_state=True
)
)(sequence_output)
LSTM = self.produce_attention_layer(LSTM)
LSTM = tf.keras.layers.Dropout(0.5)(LSTM)

babelnet_output = tf.keras.layers.Dense(outputs_size[0], activation="softmax", name="babelnet")(LSTM)
domain_output = tf.keras.layers.Dense(outputs_size[1], activation="softmax", name="domain")(LSTM)
lexicon_output = tf.keras.layers.Dense(outputs_size[2], activation="softmax", name="lexicon")(LSTM)

def produce_attention_layer(self, LSTM):
"""
Produces an Attention Layer like the one mentioned in the Raganato et al. Neural Sequence Learning Models for Word Sense Disambiguation,
chapter 3.2
:param lstm: The LSTM that will be used in the task
:return: The LSTM that was previously given in input with the enhancement of the Attention Layer
"""
hidden_states = tf.keras.layers.Concatenate()([LSTM[1],LSTM[3]])
ripetitore = tf.keras.layers.RepeatVector(tf.keras.backend.shape(LSTM[0])[1])(hidden_states)
u = tf.keras.layers.Dense(1, activation="tanh")(ripetitore)
attivazione = tf.keras.layers.Activation('softmax')(u)  # We are using a custom softmax(axis = 1) loaded in this notebook
dotor = tf.keras.layers.Multiply()([LSTM[0],attivazione])

return dotor

• What do you mean with "remove" subwords from the output? BERT's tokenization is subword-level, which means that the outputs are subword represetations. Do you mean you want to combine subword representations into word representations?
– noe
Mar 16, 2020 at 8:07
• @ncasas yes, exactly! Because, as far as I've understood, when BERT's output is transferred to a BiLSTM Neural Network, the subwords shouldn't be there Mar 16, 2020 at 8:09
• With a normal BiLSTM, you obtain as many outputs as inputs. One option would be to only take into account the output at the last subword of each word in the loss function of your WSD task (and also at inference).
– noe
Mar 16, 2020 at 9:15
• @ncasas I've added a bit of code of my WSD system's architecture, hoping that it might be helpful. From the BiLSTM, I gain an output of 5 elements, that are then passed to the attention layer, then to the dropout and then processed with the softmax. Also, If it can help: when I go in prediction, I have an output of shape (number of sentences, max seq len, number of labels per task) Mar 16, 2020 at 9:25
• With all this setup, you end up with the same number of outputs as input. Subwords won't disappear magically unless you do something specific about that. I keep my suggestion: keep only the positions of the last subword in each word; given that you have a BiLSTM, you can keep the concatenation of the first and last subwords in each word.
– noe
Mar 16, 2020 at 12:06