Consider tile/token embeddings instead of one-hot

[schrum2]

This was used in the single-example diffusion example:
https://ojs.aaai.org/index.php/AAAI/article/view/28865

[schrum2]

MLM looks at the whole sentence. Can use something similar to train a tile embedding model, but should only look at a localized window around the block instead of the whole scene. So, given a 3 x 3 patch of blocks, masking the center block, try to predict what the center block is.

[schrum2]

This issue focuses specifically on learning a latent representation for the tokens based on their surrounding tokens/tiles.

First, a dataset would need to be made that maps 3x3 patches to their center token. Then, once that is trained, it could be incorporated into the diffusion training.

[schrum2]

The block2vec method they use in this paper was actually first explained in this paper:
https://arxiv.org/pdf/2106.10155
So, I'm looking there for details

[salasemilio1]

Created a new file named create_tile_level_json_data.py which functions virtually the same as create_level_json_data.py, thought it creates the data in x by x tiles rather than overlapping snapshots of height h. Additionally, I generalized the ascii_data_browser.py file to work with any size level/frame, not just 16x16.

[schrum2]

As discussed, we need to remove the out-of-bounds tokens from the dataset. Once you've done that, follow these guidelines:

https://chatgpt.com/share/682b78c6-546c-800e-9c43-bd1e86439431

You need to make a separate training script, call it "train_block2vec.py", that takes in the dataset you've created, and then saves a resulting model (the .pt file). Once it works, add some nice usability features to the training script, like the ability to specify the input file and output location (make it a directory) via command line parameters using argparse.

[salasemilio1]

(SURF) C:\Users\salas2\Documents\GitHub\MarioDiffusion>python inference_block2vec.py --embedding_path embeddings\block2vec_embeddings.pt --tile_id 1 --top_k 5
Loaded embeddings with shape: torch.Size([13, 64])

Most similar tiles to tile 1:
Tile 10: Cosine similarity = 0.2116
Tile 0: Cosine similarity = 0.2113
Tile 11: Cosine similarity = 0.1778
Tile 4: Cosine similarity = 0.0644
Tile 12: Cosine similarity = 0.0329

(SURF) C:\Users\salas2\Documents\GitHub\MarioDiffusion>python inference_block2vec.py --embedding_path embeddings\block2vec_embeddings.pt --tile_id 2 --top_k 5
Loaded embeddings with shape: torch.Size([13, 64])

Most similar tiles to tile 2:
Tile 10: Cosine similarity = 0.3248
Tile 6: Cosine similarity = 0.2519
Tile 5: Cosine similarity = 0.1231
Tile 0: Cosine similarity = 0.0262
Tile 12: Cosine similarity = 0.0205

(SURF) C:\Users\salas2\Documents\GitHub\MarioDiffusion>python inference_block2vec.py --embedding_path embeddings\block2vec_embeddings.pt --tile_id 9 --top_k 5
Loaded embeddings with shape: torch.Size([13, 64])

Most similar tiles to tile 9:
Tile 8: Cosine similarity = 0.2529
Tile 7: Cosine similarity = 0.0118
Tile 12: Cosine similarity = -0.0239
Tile 3: Cosine similarity = -0.0382
Tile 1: Cosine similarity = -0.0726

[salasemilio1]

We have a model ready for token embeddings, but are waiting until train diffusion has been changed to train the model

[schrum2]

I'm merging various updates on train_diffusion.py to the main branch. Tomorrow, you should be able to update from main and then continue with this issue. First, some notes:

1. I think you should train the model so that the embedding vector size is just 32.
2. In train_diffusion, add a parameter names --block_embedding_model_path whose default value is None. If None, the script behaves the way it currently does. However, if a path to your trained model is provided, then you are essentially replacing the one-hot encoded vectors for each tile with its corresponding block embedding vectors. There are various places in the code where num_tiles or num_channels is mentioned (for Mario, this is 15), but that number will instead by the embedding vector length.
3. There will also be a requirement to eventually change the text_diffusion_pipeline, but put that off for as long as possible, as Oliva is still making changes there. In fact, don't even work on text-conditional model training. Instead, focus on training an unconditional model that uses block embeddings. This will require changing latent_diffusion_pipeline.py

[schrum2]

Save a class in the models directory that represents a Block2Vec model (block2vec.py). The train_block2vec.py should save a model that can be easily loaded. Hugging Face has a save/load convention with save_pretrained and load_from_pretrained methods. Make sure your model code has these methods. If a block2vec model is provided to the LevelDataset, then it will load that model and use it to provide tile vectors instead of one-hot encoding them.