Update README.md

README.md

@@ -77,6 +77,7 @@ widget:
 <h1 style="font-size: 42px">GPT-JT<h1/>
 
 # Model Summary
+
 We present GPT-JT, a fork of GPT-6B, trained on 3.53 billion tokens, that outperforms most 100B+ parameter models at classification. 
 GPT-JT was trained with a new decentralized algorithm on computers networked with 1Gbps interconnect, in contrast with typical 100Gbps-1.6Tbps data center networks.
 GPT-JT is a bidirectional dense model, which processes the prompt with bidirectional attention to fully leverage the context information, and uses causal attention only for token generation.
@@ -84,6 +85,7 @@ GPT-JT is a bidirectional dense model, which processes the prompt with bidirecti
 ***Please try out our [Online Demo](https://huggingface.co/spaces/togethercomputer/GPT-JT)!***
 
 # Quick Start
+
 ```python
 from transformers import pipeline
 pipe = pipeline(model='togethercomputer/GPT-JT-6B-v1')
@@ -98,20 +100,55 @@ tokenizer = AutoTokenizer.from_pretrained("togethercomputer/GPT-JT-6B-v1")
 model = AutoModelForCausalLM.from_pretrained("togethercomputer/GPT-JT-6B-v1")
 ```
 
-# Training Data
+# Training Details
+
+## UL2 Training Objective
+
+We train GPT-J using UL2 training objective [1][2].
+The usual GPT model, including GPT-J, uses the lower left causal mask to do autoregressive generation, so for each token, it can only see the context information before itself.
+In order to fully leverage the context information, we continue training with UL2 training objectives, and uses the lower right causal mask with prefix -- using bidirectional attention for the prompt and causal attention for token generation.
+
+$$ 
+\begin{bmatrix}
+1 & 0 & 0 & 0 & 0 \\
+1 & 1 & 0 & 0 & 0 \\
+1 & 1 & 1 & 0 & 0 \\
+1 & 1 & 1 & 1 & 0 \\
+1 & 1 & 1 & 1 & 1 
+\end{bmatrix}
+
+\begin{bmatrix}
+1 & 1 & 1 & 0 & 0 \\
+1 & 1 & 1 & 0 & 0 \\
+1 & 1 & 1 & 0 & 0 \\
+1 & 1 & 1 & 1 & 0 \\
+1 & 1 & 1 & 1 & 1 
+\end{bmatrix}  
+$$
+
+## Data
+
 We fine-tune [GPT-J-6B](https://huggingface.co/EleutherAI/gpt-j-6B) on NI, P3, COT, the pile data.
 - [Natural-Instructions](https://github.com/allenai/natural-instructions)
 - [P3](https://huggingface.co/datasets/Muennighoff/P3)
 - [MMLU-COT](https://github.com/jasonwei20/flan-2/blob/main/mmlu-cot.json)
 - [the pile](https://huggingface.co/datasets/the_pile)
 
-We first conduct training for 2.62 billion tokens using the UL2 loss, followed by 0.92 billion tokens with a mixture of the above datasets: 5% of COT, 20% of P3, 20% of NI, and 55% of the Pile.
+We first conduct training for 2.62 billion tokens using the UL2 loss on the Pile, followed by 0.92 billion tokens with a mixture of the above datasets: 5% of COT, 20% of P3, 20% of NI, and 55% of the Pile.
+
+## Hyperparameters
 
-# Hyperparameters
-We used AdamW with a learning rate of 1e-5 and global batch size of 64.
+We used AdamW with a learning rate of 1e-5 and global batch size of 64 (16 for each data parallel worker).
 We used mix-precision training where the activation is in FP16 while the optimizer states are kept in FP32.
 We use both data parallelism and pipeline parallelism to conduct training.
 During training, we truncate the input sequence to 2048 tokens, and for input sequence that contains less than 2048 tokens, we concatenate multiple sequences into one long sequence to improve the data efficiency.
 
-# Infrastructure
-We used [the Together Research Computer](https://together.xyz/) to conduct training. 
+## Infrastructure
+
+We used [the Together Research Computer](https://together.xyz/) to conduct training. 
+
+# References
+
+[1]: Tay, Yi, Mostafa Dehghani, Vinh Q. Tran, Xavier Garcia, Dara Bahri, Tal Schuster, Huaixiu Steven Zheng, Neil Houlsby, and Donald Metzler. "Unifying Language Learning Paradigms." arXiv preprint arXiv:2205.05131 (2022).
+
+[2]: Tay, Yi, Jason Wei, Hyung Won Chung, Vinh Q. Tran, David R. So, Siamak Shakeri, Xavier Garcia et al. "Transcending scaling laws with 0.1% extra compute." arXiv preprint arXiv:2210.11399 (2022).