The before/after cosine values make the learning step tangible. I would also compare the predicted next-token distributions for cat and dog across a few random seeds, because the objective constrains those predictions more directly than it constrains a unique embedding geometry.

With a trainable output matrix, two runs can solve the tiny task while arranging their internal coordinates differently. Showing low prediction disagreement alongside the cosine measurements would make the result more robust pedagogically: the shared predictive role is learned, while the exact vector arrangement remains a property of the fitted model.