I have created some document embeddings which were then used further in text classification tasks. After revisiting my code I was unsure about the workflow I used to train the document embeddings.

At the moment I am creating the document embeddings based on the complete corpus available at the time of training. After the training is done, I evaluate the model by looking whether it creates useful similarities between the document embeddings. Those embeddings are used then in machine learning models and that's where the embeddings will be split into train, test and validation sets.

Now my question is: Where is the right time to split the data? Should I do it before creating the document embeddings to prevent data leakage? I have used the mentioned approach because I viewed the creation of the document embeddings as a preprocessing step, so the computer can work with textual data. However, after I have put some thought into it, I think it's the wrong approach. I wanted to hear from more experienced NLP practitioners how they approach this task. Sorry for this very basic question.



1 Answer 1



If you are training the document-embedding model, then split the data before you convert the text into embeddings.

If you are using a pre-trained document-embedding model, then it won't matter and it is pre-processing step that it doesn't matter when you execute it.

Pipeline when training your own document-embedding model

  1. Split your text data into train/validate/test sets.
  2. Use your train set to train the document-embedding model.
  3. Use your trained document-embedding model to convert train and validation sets to train your other model (e.g. classification model).
  4. Test your final model by using your trained document-embedding model to convert the test set and test the trained final (classification) model.
  • $\begingroup$ Thank you for the answer. Does the same pipeline for training my own document-embedding model apply if I am using transfer learning? There I would use the train set to continue training a pre-trained document-embedding model and use step 3 and 4 afterwards like you have described. $\endgroup$
    – Sento
    Dec 12, 2019 at 11:10
  • 1
    $\begingroup$ Yes, even if you do transfer learning you should split it beforehand $\endgroup$ Dec 12, 2019 at 11:31

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.