Lite Oute 2 Mamba2Attn 250M

Introduction

We're excited to introduce our latest model, the Lite Oute 2 Mamba2Attn 250M. This is our third generation model featuring the new Mamba2 architecture with attention layers.

Model Variants

The Lite Oute 2 Mamba2Attn 250M model comes in two variants:

• Lite-Oute-2-Mamba2Attn-250M-Instruct
• Lite-Oute-2-Mamba2Attn-250M-Base

Architecture Overview

Lite Oute 2 features 250 million parameters and the following architectural specifications:

• d_model: 1024
• n_layer: 32
• vocab_size: 32768
• num_heads: 16
• head_dim: 128
• attn_layer_idx: [6, 12, 18, 24]

Attention Layer Integration

Training Process

The model was pre-trained on 30 billion tokens using a balanced mixture of datasets:

• 50% dclm-baseline-1.0
• 50% fineweb-edu

The training was conducted on single NVIDIA 4090 and NVIDIA H100 GPUs, with the following key parameters:

• Max learning rate: 4e-4
• Min learning rate: 1e-4
• Block size: 4096
• Token batches: ~100k tokens

To enhance training efficiency and performance, we employed various optimization techniques, including:

• Paged AdamW 32-bit optimizer
• Mixed precision training with bfloat16
• Gradient accumulation
• Flash Attention 2

Instruction Training

For instruction training, we first trained the model with Supervised Fine-tuning (SFT):

• Max learning rate: 2e-5
• Scheduler: Cosine
• Max block size: 4096
• Token batches: ~100k tokens

After the SFT training, we further refined the model using DPO training:

• Max learning rate: 9e-7
• Scheduler: Cosine
• Max block size: 4096
• Token batches: ~100k tokens

Benchmark Results

Instruct model:

Base model:

Interfacing with the Instruct Model

To ensure optimal performance, please use the following template when interacting with the model:

<s>system\n
{System message}</s>
<s>user\n
{Your question or message}</s>
<s>assistant\n
{The model's response}</s>

When prompting the model for a response, your input should be structured as follows:

<s>system\n
{System message}</s>
<s>user\n
{Your question or message}</s>
<s>assistant\n

This format ensures the best performance, as the model was fine-tuned using this specific template.

Recommended Settings

• Temperature: 0.1 - 0.4
• Repetition Penalty: 1.10 - 1.12

If you're using a Hugging Face tokenizer, make sure to use tokenizer.apply_chat_template or adjust your interface to match the template above.

Let's look at some real world examples:

Instruct Model

Conclusion

The Lite-Oute-2-Mamba2Attn-250M-Instruct model demonstrates notable improvements over its predecessors, particularly in zero-shot and few-shot learning tasks. Despite its smaller size compared to Lite-Oute-1-300M-Instruct, it consistently outperforms both older models across most benchmarks. The model shows enhanced capabilities in tasks requiring common sense reasoning, language understanding, and general knowledge. Real-world examples indicate that the model can provide coherent and relevant responses to various queries, though there's room for improvement in accuracy and depth of knowledge.

Strengths:

• Adaptability: The model demonstrates the ability to handle various types of queries, from simple factual questions to more complex, multi-faceted inquiries.
• Structured Responses: For more detailed questions, the model often provides organized answers with numbered points, enhancing readability and comprehension.
• Context Utilization: When given additional context or information, the model shows capability in extracting and using relevant details in its responses.

Weaknesses:

• Inconsistent Accuracy: The model sometimes provides incomplete or partially incorrect information, indicating potential gaps in its knowledge base or reasoning abilities.
• Lack of Depth: In some instances, especially for simpler queries, the model's responses lack depth or additional context that could be beneficial to the user.
• Verbosity Control: The model occasionally struggles with balancing conciseness and detail, either providing overly brief answers or unnecessarily repeating given information.

Fine-Tuning for Specific Tasks

For optimal task-specific performance, it is recommended to fine-tune the model on datasets relevant to the intended use case. Fine-tuning can significantly improve the model's accuracy, relevance, and depth of knowledge in specialized domains or for particular types of tasks. This process allows the model to adapt its pre-trained knowledge to the nuances and requirements of specific applications, potentially mitigating some of the weaknesses observed in general use.

Limitations and Considerations

Users should be aware that Lite Oute 2 is a relatively small language model with only 250 million parameters. As such, it has limitations in its knowledge, reasoning capabilities, and overall performance. This model can and will produce nonsensical or incorrect outputs at times. It is not a replacement for larger, more sophisticated models or human expertise. Users should always critically evaluate the model's outputs and not rely on them for critical decision-making without proper verification. This model is best suited for lightweight applications or as a starting point for further fine-tuning and development.