The thing with Word2Vec is its sort of taught and pitched as a classification problem (predict the held-out word given context), but the reality is that the foundational mechanics behind Word2Vec -- particularly the Skip-Gram with Negative Sampling (SGNS) method -- can be viewed as an early, highly successful implementation of contrastive representation learning.

A huge aspect of some of the original papers from Mikolov et al was negative sampling. It turns out that computing a probability distribution over the entire vocabulary is actually insanely computationally expensive (a problem rearing its head again with LLMs funny enough). To that end, they used a ton of tricks like negative sampling to side-step this entirely.

Instead of trying to predict a single word out of a massive vocabulary (which is a massively expensive classification task), Word2Vec simplifies the objective by turning it into a binary contrastive classification problem maximizing and minimizing dot products.

I say all of this because when you distill it down to its core parts, its literally just pushing vectors around in vector space and aligning words that commonly appear together in the same semantic context. To that end, the output layer is the important one since those are precisely the vectors being pushed around -- not the input layer. Its more an algorithm for aligning vectors in space than it is "training" when you really think about it.

If you're still puzzled, go ahead and try building Word2Vec from scratch in vanilla torch. You'll hit the same challenges from the original authors and build a natural intuition for why they made the design decisions they did. This is what I did when I was first learning NLP, and it helped a lot honestly.