Great news -- transformers do this already. The way transformers generate discrete tokens is by using a linear layer to do classification over the vocabulary given the transformer outputs. The transformer outputs are just dense vectors in the same shape as the input token embeddings (i.e., of shape =
(batch_size, seq_len, embed_dim)).
You could train the model with something like MSE loss. Transformers don't necessarily have to predict discrete tokens (e.g., the MAE paper use a transformer to predict pixel values directly, and trains with MSE). Cosine-distance could also work as a loss function for embedding prediction, as used in this paper which has a transformer model that produces an embedding output.
However, it's important to note that it's common for people to convert the data into some discrete format first. For example, for applications like image or audio generation, you often have a VQ-VAE create a "code-book" of embeddings, and you'd use a transformer to model those codes as a discrete "vocabulary". See this paper for an example. Finally, models like imageGPT instead quantize color values to create a discrete "vocabulary".