Question
Can token embeddings be arranged as small image patches and used as inputs to a joint-embedding predictive model?
The proposed representation maps each token embedding to a 16 × 16 × 3 tile. A passage becomes a grid of tiles. The model sees part of that grid and predicts representations for the masked positions.
Approach
A context encoder processes the visible tiles. A predictor estimates the representations of masked tiles. A target encoder receives the full input and supplies the comparison targets; its weights follow an exponential moving average of the context encoder.
The objective would compare predicted and target representations, rather than score the prediction over the token vocabulary. That removes a vocabulary-sized output head from this objective. It does not remove the input embedding table or establish an advantage over token-based models.
Representation
The mapping from embedding dimensions to tile positions determines which values are spatial neighbours. A direct reshape is a useful starting point. Other mappings would need to be compared against it under the same training conditions.
A change in layout is not automatically a meaningful change in the model. An input projection may absorb some rearrangements. Any comparison needs to distinguish the effect of the mapping from the flexibility of the input stem.
Evaluation
A useful comparison would keep the corpus, parameter budget, training steps, and evaluation tasks fixed. It should include an untrained control and a token-based baseline.
The latent representation also needs to be checked for collapse. A small prediction error alone would not show that the model learned useful distinctions.
Open questions
- Does the spatial arrangement contribute anything beyond an input projection?
- How much information is retained by the embedding-to-tile mapping?
- Which masking scheme makes sense for a grid whose contents originated as a sequence?
- Do learned representations transfer to a downstream task at a comparable compute budget?