Update README.md

README.md

@@ -9,21 +9,23 @@ Zamba-2-2.7B is a hybrid model between state-space models and transformers. It b
 2.) Instead of a single shared attention block, we utilize two shared attention blocks which are interleaved in an ABAB pattern through the network.
 3.) We apply a LoRA projector to each shared MLP block allowing the network to specialize the MLPs at each shared layer with a minimal increase in total parameter count.
 
-Zamba was trained using next-token prediction. It uses the Mistral v0.1 tokenizer. Zamba2-2.7B was pre-trained on 3T tokens of text and code data sourced from open web-datasets. Subsequently in a second phase, Zamba was annealed on a mixture of 100B high-quality tokens.
+Zamba was trained using next-token prediction. It uses the Mistral v0.1 tokenizer. Zamba2-2.7B was pre-trained on 3T tokens of text and code data sourced from open web-datasets. Subsequently, in a second phase, Zamba2-2.7B was annealed on a mixture of 100B high-quality tokens.
 
-Note: this is a temporary HuggingFace implementation of Zamba 3B and is designed for specific use cases. It may not be fully compatible with all frameworks and tools intended to interface with HuggingFace models.
+Note: this is a temporary HuggingFace implementation of Zamba2-2.7B. It may not yet be fully compatible with all frameworks and tools intended to interface with HuggingFace models.
+
+A pure Pytorch implementation of Zamba2-2.7B may be found [here](https://github.com/Zyphra/Zamba2).
 
 ## Quick start
 
 ### Presequities
 
-To download Zamba 3B, clone Zyphra's fork of transformers:
+To download Zamba2-2.7B, clone Zyphra's fork of transformers:
 1. `git clone https://github.com/Zyphra/transformers_zamba2.git`
 2. `cd transformers_zamba2`
 3. Install the repository: `pip install -e .`
 
 
-You can run the model without using the optimized Mamba kernels, but it is **not** recommended as it will result in significantly higher latency. 
+You can run the model without using the optimized Mamba kernels, but it is **not** recommended as it will result in significantly higher latency and memory usage. 
 
 To run on CPU, please specify `use_mamba_kernels=False` when loading the model using ``AutoModelForCausalLM.from_pretrained``.
 
@@ -44,13 +46,12 @@ outputs = model.generate(**input_ids, max_new_tokens=100)
 print(tokenizer.decode(outputs[0]))
 ```
 
-## Model Details [to update!]
-
-Zamba utilizes a unique hybrid SSM architecture. This architecture consists of a backbone of Mamba layers interspersed with a shared attention layer. This attention has shared weights to minimize the parameter cost of the model. We find that concatenating the original model embeddings to the input to this attention block improves performance, likely due to better maintenance of information across depth. 
+## Model Details
 
+Zamba2-2.7B utilizes and extends our original Zamba hybrid SSM-attention architecture. The core Zamba architecture consists of a backbone of Mamba layers interleaved with a shared attention layer. This attention has shared weights to minimize the parameter cost of the model. We find that concatenating the original model embeddings to the input to this attention block improves performance, likely due to better maintenance of information across depth. The Zamba2 architecture also applies LoRA projection matrices to the shared MLP to gain some additional expressivity in each block and allow each shared block to specialize slightly to its own unique position while keeping the additional parameter overhead small. 
 
 <center>
-<img src="https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/IGK562oVTFSOQbpLavu7E.png" width="300" alt="Zamba architecture">
+<img src="https://cdn-uploads.huggingface.co/production/uploads/65c05e75c084467acab2f84a/XrEIEBxd0fqIgh3LyArAV.png" width="300" alt="Zamba architecture">
 </center>