AFM-4.5B-Base-KDA-NoPE

AFM-4.5B-Base-KDA-NoPE

A hybrid attention variant of AFM-4.5B-Base combining Kimi Delta Attention (KDA) with NoPE (No Positional Encoding) full-attention layers in a 3:1 ratio. This architecture balances efficiency with performance through knowledge distillation.

⚠️ Research Model: This is an experimental model released for research purposes. For production use, see AFM-4.5B.

More details available in our blog post here: https://www.arcee.ai/blog/distilling-kimi-delta-attention-into-afm-4-5b-and-the-tool-we-used-to-do-it

Overview

Following the Kimi Linear architecture pattern, this model interleaves KDA layers with periodic full-attention layers (using NoPE) in a 3:1 ratio. This hybrid structure reduces memory and KV-cache usage while preserving global information flow via the full attention layers.

Key characteristics:

• 3:1 KDA to full-attention ratio
• Full attention layers use NoPE (No Positional Encoding)
• Trained up to 32k sequence length
• Better short-context performance than pure KDA
• Reduced memory footprint compared to full attention

Architecture

Benchmark Results

Performance compared to the teacher model and other configurations:

Key Findings

• Math advantage: The hybrid recovers significantly more math performance (36.5%) than pure KDA (26.8%)
• Knowledge benchmarks: Performs comparably to KDA-Only on MMLU, ARC, and HellaSwag
• Efficiency: Maintains efficiency gains from KDA while preserving global reasoning via NoPE layers

Long-Context Performance (NIAH)

The hybrid model shows distinct long-context behavior:

• 100% single-needle retrieval up to 32k
• Sharp performance cliff past 32k training length
• Near-zero performance beyond training context (vs. smooth degradation for KDA-Only)

The NoPE full-attention layers appear responsible for the hard cutoff—they haven't seen positions beyond 32k during training. KDA layers generalize more naturally to longer sequences.

Usage

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_id = "arcee-ai/AFM-4.5B-Base-KDA-NoPE"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

prompt = "The theory of relativity states that"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)

outputs = model.generate(
    input_ids,
    max_new_tokens=100,
    do_sample=True,
    temperature=0.7,
    top_p=0.95
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Training Details

• Method: Knowledge distillation from AFM-4.5B-Base using DistillKit
• Teacher: AFM-4.5B-Base (full attention)
• Student Architecture: Hybrid 3:1 KDA:NoPE
• Training Length: 32k sequence length

Comparison: Hybrid vs KDA-Only

Choose Hybrid for better short-context reasoning, especially math. Choose KDA-Only for more predictable long-context degradation.

Intended Use

This model is intended for:

• Research into hybrid attention architectures
• Studying linear/full attention tradeoffs
• Exploring NoPE attention in hybrid configurations
• Benchmarking efficiency vs. capability tradeoffs

License

AFM-4.5B is released under the Apache-2.0 license.