In this project, we implemented a small-scale version of the BBPE (Byte-level BPE) algorithm using a Russian jokes dataset to build a word table for sequence serialization of Russian text. We developed a Decoder Only model inspired by the Llama architecture, which supports attention masks and incorporates cutting-edge transformer-based techniques such as RoPE (Rotary Positional Encoding) or ALiBi (Attention with Linear Biases for Bigger Contexts).
2. Model Architecture
Decoder Only Structure: The model follows a similar structure to the Llama family of models, focusing solely on sequence-to-sequence tasks without including a full generator component.
Support for Attention Masks: The implementation includes attention mask functionality, enabling flexible context windows for processing input sequences.
Positional Encoding Choices:
RoPE (Rotary Positional Encoding): Implemented as an alternative to traditional sinusoidal positional encodings, offering distinct performance characteristics.
ALiBi (Attention with Linear Biases for Bigger Contexts): Integrated for compatibility with state-of-the-art transformer architectures.
Multi-head Attention Mechanism: Utilizes multi-head attention layers, allowing the model to capture complex relationships within input sequences.
Efficient Caching Strategy (GQA optimization): Implemented cache size reduction techniques while maintaining compatibility with KV-cache structures.
3. Training and Evaluation
Training Data: The Russian jokes dataset was used for model training and loss function analysis.
Performance Comparison: The implementation compared RoPE vs. ALiBi positional encodings, observing that RoPE-based models demonstrated superior convergence properties in the loss function.
Model Sizes: Trained three variant models (0.5M, 7M, and 80M parameters) to establish a foundation for future LLM applications based on these architectures.
4. Conclusion
This project successfully implemented a small-scale BPE-based model capable of handling Russian text sequences with attention mechanisms and modern positional encoding techniques. The results provide valuable insights into the development of efficient LLMs tailored for specific language tasks, particularly in the context of humor or conversational applications. Future work could extend this approach to larger datasets and more complex tasks.
Loss curve
The loss curve for nano model with ALiBi position embedding are as follows:
