I converted the way a tinyllama model works with the following script:
https://huggingface.co/Josephgflowers/Liquid-Metal-Tinyllama-Test-1/blob/main/LM.py
Continued training on a Cinder dataset for 1 epoch.
This model could use more training to fully utilize the new features.
- Overall Flow in the Model
Each of these modules is integrated into the model’s modified decoder layer (ModifiedLlamaDecoderLayer). Here’s a high-level outline of the sequence in which they operate within the decoder:
Step 1: Adaptive RMSNorm normalizes the input while applying an adaptive scaling based on the global context of each input batch.
Step 2: Token Mixing performs a local convolution across tokens in the sequence, helping to capture intra-sequence dependencies.
Step 3: Post-Attention Adaptive RMSNorm applies adaptive normalization after attention processing.
Step 4: The output is passed through the model’s MLP (multilayer perceptron) layer for further feature transformation.
Step 5: SEBlock performs global channel-wise recalibration to enhance or suppress certain channels based on the context.
Let’s break down how these components contribute to the model’s overall performance:
- Component-Level Contributions
Adaptive RMSNorm
Purpose: Provides context-sensitive normalization, allowing the model to scale features dynamically based on the input’s global context.
Effect on Model: Makes the normalization process adaptable rather than static, which can improve the model’s ability to generalize across diverse inputs. This is especially useful in language models where different prompts may require different emphasis on specific features.
Performance Impact: Adaptive scaling helps maintain stability in training, as it smooths out variations while retaining sensitivity to input-specific details. This can lead to improved convergence and robustness, especially in complex tasks.
Token Mixing
Purpose: Blends information across tokens within each feature channel through depthwise convolution across the sequence dimension.
Effect on Model: By capturing local dependencies within the sequence, Token Mixing complements self-attention’s global scope, giving the model a better understanding of local patterns and relationships.
Performance Impact: Improves the model’s intra-sequence awareness, which can be particularly beneficial in processing structured or position-sensitive data. This layer’s lightweight nature makes it a low-cost way to add a degree of locality that can enhance overall performance.
SEBlock (Squeeze-and-Excitation Block)
Purpose: Performs adaptive channel-wise recalibration by scaling each feature channel based on global context.
Effect on Model: SEBlock helps the model emphasize or suppress specific features across all tokens, adapting the channel importance to match the input context.
Performance Impact: Boosts the model’s expressiveness by allowing it to dynamically adjust which features are most relevant for each input. This helps improve generalization, especially when handling varied inputs with different feature relevances, such as conversations with shifting topics.