Visualizing 🪆 Models

Architecture

The suggested Matryoshka Architecture Suite, outlined in the paper, essentially builds nested sub-models, such that each one is composed of its predecessors:

The Junction Mechanism

To deal with the size mismatch between the different submodels, they suggest the following “Junction” mechanism:

We do 2 things here:

1. match the norm of the output, to the norm of the next sub-models’ embedding of the input.
2. concat the normalized output from the previous model, to the input embedding, and feed both to the next submodel.

The Key Idea

The nested structure that we visualized here, is what allows the Matryushka models to be modular - we essentially trained one model (the largest one), and recieved the smaller models “for free” (hence the 36% reduction in training compute). The benefits of the structure are also related to sharing the KV-cache,and allowing for easier inter-model tasks, namely Speculative Decoding and Model Distillation.