Add missing BGE-M3 sparse and ColBERT projection weights
Hello,
The current mlx-community/bge-m3-mlx-8bit repository contains the quantized XLM-RoBERTa backbone, but it does not include the BGE-M3-specific sparse_linear and colbert_linear projection heads. As a result, the published MLX model can generate dense embeddings, but the sparse and multi-vector functionality described for BGE-M3 cannot be reproduced directly from the files currently included in the repository.
I verified that the official sparse_linear.pt from BAAI/bge-m3 is compatible with the MLX backbone after conversion to safetensors. Dense and sparse representations can then be produced from a single backbone forward pass, with negligible additional cost for the sparse projection.
Could you please add the following files to the MLX model repository:
- sparse_linear.safetensors
- ideally, colbert_linear.safetensors
The weights should be converted directly from the corresponding official files in BAAI/bge-m3, without additional training or modification. A simple and convenient key layout would be:
weight
bias
For the sparse projection, the output is computed from last_hidden_state as follows:
sparse_weights = mx.maximum( hidden_states @ weight.T + bias, 0)
Same as - sparse_weights = nn.relu( hidden_states @ weight.T + bias)
It would also be useful to add a short model-card example showing how to obtain dense and sparse outputs from the same forward pass.
I have implemented and documented the MLX-side integration, including a reusable sparse projection helper and an optimized fused-attention path for XLM-RoBERTa, here:
MLX XLM-RoBERTa fused SDPA and BGE-M3 sparse projection documentation
The corresponding implementation is available in the same branch:
https://github.com/anfedoro/mlx-embeddings/tree/optimize/xlm-roberta-fused-sdpa
On an Apple M4 Max, the sparse projection itself adds less than one millisecond per batch in the tested configuration; nearly all runtime remains in the encoder forward. Including the missing projection weights would therefore make the published BGE-M3 MLX model functionally complete without introducing meaningful runtime overhead.
Thank you.