Gemma 3n is a new member of the Gemma family with free weights that was released during Google I/O. It's dedicated to on-device (edge) inference and supports image and text input, with audio input. Google has released an app that can be used for inference on the phone.
What is clear from the documentation, is that this model is stuffed to the brim with architectural innovations: Per-Layer Embedding (PLE), MatFormer Architecture, Conditional Parameter Loading.
Unfortunately, there is no paper out for the model yet. I assume that this will follow at some point, but so far I had some success poking around in the model file. I thought I'd share my findings so far, maybe someone else has more insights?
The provided .task file is actually a ZIP container of tflite models. It can be unpacked with ZIP.
| TF_LITE_PREFILL_DECODE | 2.55 GB | Main language model component for text generation |
| TF_LITE_PER_LAYER_EMBEDDER | 1.23 GB | Per-layer embeddings from the transformer |
| TF_LITE_EMBEDDER | 259 MB | Input embeddings |
| TF_LITE_VISION_ENCODER | 146 MB | Vision Encoding |
| TF_LITE_VISION_ADAPTER | 17 MB | Adapts vision embeddings for the language model? |
| TOKENIZER_MODEL | 4.5 MB | Tokenizer |
| METADATA | 56 bytes | general metadata |
The TFlite models can be opened in a network visualizer like netron.app to display the content.
The model uses an inner dimension of 2048 and has 35 transformer blocks. Tokenizer size is 262144.
First, one interesting find it that is uses learned residual connections. This paper seems to be related to this: https://arxiv.org/abs/2411.07501v3 (LAuReL: Learned Augmented Residual Layer)
The FFN is projecting from 2048 to 16384 with a GeGLU activation. This is an unusually wide ratio. I assume that some part of these parameters can be selectively turned on and off to implement the Matformer architecture. It is not clear how this is implemented in the compute graph though.
A very interesting part is the per-layer embedding. The file TF_LITE_PER_LAYER_EMBEDDER contains very large lookup tables (262144x256x35) that will output a 256 embedding for every layer depending on the input token. Since this is essentially a lookup table, it can be efficiently processed even on the CPU. This is an extremely interesting approach to adding more capacity to the model without increasing FLOPS.
The embeddings are applied in an operation that follows the FFN and are used as a gate to a low rank projection. The residual stream is downprojected to 256, multiplied with the embedding and then projected up to 2048 again. It's a bit like a token-selective LoRA. In addition there is a gating operation that controls the overall weighting of this stream.
I am very curious for further information. I was not able to find any paper on this aspect of the model. Hopefully, google will share more information.
English (US) · 
