Update README.md

README.md

@@ -24,6 +24,67 @@ This model is suitable for tasks that require complex decision-making and optimi
 
 The model is constructed with several primary components:
 1. **Transformer**: The transformer has encoder and decoder layers with rotary positional encoding and Mixture of Experts (MoE) to improve generalization and reduce computational cost by routing only parts of the data to certain experts. GELU and SwiGLU activation functions are alternated between the experts.
+
+### Multi-Token Prediction with Beam Search
+
+Multi-token prediction in a language model involves generating multiple tokens in sequence, rather than one token at a time. This can improve the fluency and coherence of generated text by allowing the model to "look ahead" and consider multiple possible continuations at each step.
+
+**Beam Search** is a popular decoding algorithm used for multi-token prediction that allows the model to explore multiple potential sequences and choose the most likely one based on the overall probability. Here's how it works:
+
+1. **Initialization**:
+   - Start with a single "beam" (sequence) that contains the initial token, typically the beginning-of-sequence (`<sos>`) token.
+
+2. **Expansion**:
+   - At each time step, the model generates a probability distribution over the vocabulary for each sequence in the beam.
+   - For each sequence, it expands by predicting the next possible tokens, creating new sequences for each possible token.
+
+3. **Scoring**:
+   - Calculate the score for each expanded sequence by taking the sum (or average) of log probabilities for all tokens in the sequence. Log probabilities are used to avoid underflow and ensure stable computation.
+
+4. **Selection**:
+   - Keep only the top-k sequences with the highest scores (known as the "beam width" or "beam size") and discard the rest. This limits the number of sequences kept at each step, focusing only on the most promising ones.
+
+5. **Repeat**:
+   - Continue expanding and scoring until reaching the desired sequence length or the end-of-sequence (`<eos>`) token.
+
+6. **Final Output**:
+   - After a set number of steps, or if all sequences end with `<eos>`, select the sequence with the highest score as the final output.
+
+This process allows the model to generate more fluent and accurate sequences by considering multiple potential continuations at each step and selecting the best overall sequence.
+
+---
+
+### Brief Overview of the Transformer Architecture
+
+The Transformer architecture, introduced in the paper "Attention is All You Need," is a powerful neural network design for handling sequential data, especially in natural language processing tasks. Transformers are known for their parallelism and ability to capture long-range dependencies in data.
+
+#### Key Components of the Transformer
+
+1. **Embeddings and Positional Encoding**:
+   - The input tokens are embedded into dense vectors. Since Transformers do not inherently encode the sequence order (as opposed to RNNs), they require **positional encodings**. These encodings are added to the embeddings to provide information about the token positions in the sequence.
+   
+2. **Multi-Head Self-Attention**:
+   - Each token in a sequence attends to every other token, capturing dependencies regardless of distance. Multiple attention heads allow the model to focus on different parts of the sequence, extracting varied features.
+   - In self-attention, the model computes **query**, **key**, and **value** vectors for each token. The output is a weighted sum of values, where the weights are determined by the similarity between the query and key vectors.
+
+3. **Feedforward Neural Networks**:
+   - After self-attention, a position-wise feedforward neural network is applied to each token independently. This network consists of two linear layers with a ReLU or GELU activation function in between.
+   
+4. **Layer Normalization and Residual Connections**:
+   - To improve learning stability, **layer normalization** is applied. Residual connections help the model to learn effectively by adding the input of a layer to its output, allowing gradients to flow more easily during backpropagation.
+   
+5. **Stacking of Layers**:
+   - The Transformer consists of **multiple encoder and decoder layers**. Each encoder layer is identical and consists of self-attention and feedforward layers. The decoder layers include an additional cross-attention mechanism to attend to the encoder's output.
+
+6. **Final Linear and Softmax Layer**:
+   - The final output of the decoder layer is passed through a linear layer, projecting it onto the vocabulary size. A **softmax** function then converts the output into a probability distribution over the vocabulary, from which the next token is selected or sampled.
+
+#### Encoder-Decoder Structure
+
+- **Encoder**: The encoder processes the input sequence into a contextualized representation that captures relationships between tokens. It consists of multiple layers of self-attention and feedforward networks.
+- **Decoder**: The decoder generates the output sequence by attending to both the encoded input representation (using cross-attention) and previously generated tokens (using self-attention). The decoder's output is used to predict the next token in the sequence.
+
+
 2. **Representation Network**: This module encodes the Transformer output to generate a state representation, reducing dimensionality and making it suitable for further processing.
 3. **Dynamics Network**: This module predicts the next state given a current state and an action. It uses layer normalization and a GELU activation function.
 4. **Prediction Network**: Predicts both the policy logits and value estimates for a given state. It outputs the probabilities of different actions as well as a single scalar value.
@@ -46,6 +107,9 @@ thought_1 = {P1, ... , PN}
 
 The model explores and exploits thoughts, policies, actions, and tokens, and learning happens at each step of granularity. 
 
+
+
+
 ## Training Details
 
 The model is trained with the following components and techniques: