Update README.md

README.md

@@ -32,9 +32,7 @@ widget:
 ---
 
 # Graphcore/gptj-mnli
-
 This model is the fine-tuned version of [EleutherAI/gpt-j-6B](https://huggingface.co/EleutherAI/gpt-j-6B) on the [GLUE MNLI dataset](https://huggingface.co/datasets/glue#mnli).
-
 MNLI dataset consists of pairs of sentences, a *premise* and a *hypothesis*.
 The task is to predict the relation between the premise and the hypothesis, which can be:
 - `entailment`: hypothesis follows from the premise,
@@ -51,6 +49,34 @@ For example:
 mnli hypothesis: Your contributions were of no help with our students' education. premise: Your contribution helped make it possible for us to provide our students with a quality education. target: contradiction <|endoftext|>
 ```
 
+## Model description
+
+GPT-J 6B is a transformer model trained using Ben Wang's [Mesh Transformer JAX](https://github.com/kingoflolz/mesh-transformer-jax/). "GPT-J" refers to the class of model, while "6B" represents the number of trainable parameters.
+
+<figure>
+
+| Hyperparameter       | Value      |
+|----------------------|------------|
+| \\(n_{parameters}\\) | 6053381344 |
+| \\(n_{layers}\\)     | 28&ast;    |
+| \\(d_{model}\\)      | 4096       |
+| \\(d_{ff}\\)         | 16384      |
+| \\(n_{heads}\\)      | 16         |
+| \\(d_{head}\\)       | 256        |
+| \\(n_{ctx}\\)        | 2048       |
+| \\(n_{vocab}\\)      | 50257/50400&dagger; (same tokenizer as GPT-2/3)  |
+| Positional Encoding  | [Rotary Position Embedding (RoPE)](https://arxiv.org/abs/2104.09864) |
+| RoPE Dimensions      | [64](https://github.com/kingoflolz/mesh-transformer-jax/blob/f2aa66e0925de6593dcbb70e72399b97b4130482/mesh_transformer/layers.py#L223) |
+<figcaption><p><strong>&ast;</strong> Each layer consists of one feedforward block and one self attention block.</p>
+<p><strong>&dagger;</strong> Although the embedding matrix has a size of 50400, only 50257 entries are used by the GPT-2 tokenizer.</p></figcaption></figure>
+
+The model consists of 28 layers with a model dimension of 4096, and a feedforward dimension of 16384. The model
+dimension is split into 16 heads, each with a dimension of 256. Rotary Position Embedding (RoPE) is applied to 64
+dimensions of each head. The model is trained with a tokenization vocabulary of 50257, using the same set of BPEs as
+GPT-2/GPT-3.
+
+[EleutherAI/gpt-j-6B](https://huggingface.co/EleutherAI/gpt-j-6B), our starting point for finetuning, is trained on [the Pile](https://pile.eleuther.ai), a large-scale curated dataset created by [EleutherAI](https://www.eleuther.ai).
+
 ## Fine-tuning and validation data
 Fine tuning is done using the `train` split of the GLUE MNLI dataset and the performance is measured using the [validation_mismatched](https://huggingface.co/datasets/glue#mnli_mismatched) split.