MistralLite: A Fine-Tuned 7B Model for Enhanced Long Context Handling

MistralLite is a fine-tuned version of the Mistral-7B-v0.1 language model, designed to excel in processing long contexts (up to 32K tokens). This model leverages adapted Rotary Embeddings and a sliding window mechanism during fine-tuning, which significantly improves its performance on tasks that require handling extensive text inputs. MistralLite is particularly useful for applications like long context line retrieval, topic summarization, question-answering, and more.

Key Technical Enhancements

• Adapted Rotary Embedding: The model uses a modified rotary embedding with rope_theta = 1000000, which allows it to better handle longer sequences.
• Sliding Window Mechanism: During fine-tuning, MistralLite employs a sliding window of size 16384 tokens. This technique helps the model maintain context coherence over long inputs.

Performance and Deployment

MistralLite can be deployed on a single AWS g5.2x instance using the Hugging Face Text Generation Inference (TGI) endpoint, making it ideal for resource-constrained environments that require high performance. You can also serve MistralLite directly using TGI Docker containers or other serving methods like vLLM. For Python users, the model is compatible with the Hugging Face transformers library and supports FlashAttention-2 for optimized inference.

Comparison with Mistral-7B-Instruct-v0.1

| Model | Fine-tuned on long contexts | Max context length | Rotary Embedding adaptation | Sliding Window Size | | --- | --: | --: | --: | --: | | Mistral-7B-Instruct-v0.1 | up to 8K tokens | 32K | rope_theta = 10000 | 4096 | | MistralLite | up to 16K tokens | 32K | rope_theta = 1000000 | 16384 |

Motivation and Development

The development of MistralLite was driven by the need to improve long context handling, a critical capability for many real-world applications. While Mistral-7B-Instruct-v0.1 excelled on short context benchmarks, its performance on longer contexts (beyond 4096 tokens) was less competitive. By fine-tuning the original model with a focus on long context tasks, MistralLite significantly boosts performance in this area.

Evaluation Results

MistralLite was evaluated against several benchmarks designed to assess long context handling capabilities. The results show:

1. Topic Retrieval: Improved accuracy in identifying and retrieving relevant topics from extensive text inputs.
2. Summarization: Enhanced ability to generate concise summaries of long documents.
3. Question-Answering: Better performance in answering questions that require understanding large amounts of context.

Example Usage

To use MistralLite, you can employ the following prompt template:

<|prompter|>What are the main challenges to support a long context for LLM?</s><|assistant|>

This format ensures the model processes the input correctly and provides coherent responses.

Conclusion

MistralLite represents a significant step forward in handling long contexts, making it a valuable tool for applications that require processing extensive text inputs. Its efficient deployment options and compatibility with popular libraries make it accessible to a wide range of users.