Retrieval models are often trained on cloud GPUs even when the production target is a laptop, desktop, or edge node. That split creates friction: data must leave the developer machine, experiment loops depend on GPU queue availability, and the final model may need another conversion step before it can run locally. Apple MLX changes that trade-off for teams working on Apple Silicon by making tensor operations, automatic differentiation, and model adaptation run directly on the unified CPU/GPU memory architecture.
The jina-ai/mlx-retrieval project is a compact example of this pattern. It fine-tunes embedding and reranker-style retrieval models on Apple Silicon using MLX, LoRA adapters, contrastive losses such as InfoNCE and NT-Xent, gradient accumulation for large effective batches, streamed JSONL or Elasticsearch training data, and MTEB evaluation during training. The reference README reports roughly 4000-5000 training tokens per second on an M3 Ultra with an effective batch size of 256 and 16 gradient accumulation steps.
The transferable lesson is not the exact repository command. It is the shape of a local retrieval fine-tuning loop: convert or load a compact base model, inject a retrieval adapter, train on query-document pairs with hard negatives, evaluate on retrieval tasks frequently, and export an adapter that can be used by the same local runtime. That loop is useful when data privacy, cost, latency, or rapid iteration makes cloud-first training unattractive.
This page focuses on retrieval-specific fine-tuning rather than generic instruction tuning. Retrieval models need geometry: related queries and documents should become neighbors, hard negatives should move away, prompt markers should preserve asymmetric query/document roles, and evaluation should measure ranking quality rather than language-generation fluency.
MLX is Apple’s array and neural-network framework for Apple Silicon. Its main operational advantage is that arrays can live in unified memory shared by CPU and GPU, reducing the explicit host-device copy management common in discrete GPU workflows. For retrieval fine-tuning, this makes small-to-medium experiments practical on developer hardware, especially when the model is compact and the training loop uses adapters instead of updating every parameter.
The local runtime changes engineering constraints. Instead of maximizing cluster utilization, the training loop must keep the Mac’s unified memory below pressure limits, avoid slow remote data reads, and choose model sizes that leave room for activations, optimizer state, batches, and evaluation. The win is a low-latency experiment loop where data preparation, training, evaluation, and inference can happen in one environment.
The source project uses LoRA-style adaptation for retrieval tasks. LoRA inserts trainable low-rank matrices into selected model layers while freezing most base weights. This reduces optimizer memory and makes each experiment cheaper to store. For retrieval, a single adapter can specialize the embedding space for one domain, corpus, or query style without maintaining a full model copy.
This differs from task-specific multi-adapter systems. The README notes that this implementation trains a single retrieval configuration rather than the multi-LoRA setup used by larger Jina embedding model families. That simplification is appropriate for a focused local experiment: one corpus, one query distribution, one retrieval objective, one adapter artifact.
Retrieval training usually presents the model with a query, a positive document, and one or more negatives. InfoNCE-style losses reward high similarity for the positive pair and low similarity for negatives. Hard negatives are especially valuable because they are plausible but wrong: a document that shares keywords, topic, or structure with the positive but does not answer the query forces the embedding model to learn finer distinctions.
Batch composition matters because in-batch negatives become part of the objective. A larger effective batch exposes each query to more competing documents, but it also increases memory requirements. Gradient accumulation gives the model the statistical effect of a larger batch while processing smaller microbatches that fit in local memory.
A generative model can sound better without retrieving better documents. A retrieval model must improve the neighborhood structure of its embedding space. The training loop should therefore focus on pair construction, negative mining, pooling, prompt markers, and ranking metrics rather than only loss curves or generated text quality.
A practical local setup starts with a model small enough to train comfortably on the target Mac. The source repository uses a Gemma 3 270M model converted to MLX format as the example, then adds LoRA during training. The base model supplies language understanding; the adapter shifts the representation space toward the retrieval task.
# Conceptual command shape based on the source README.
python train.py \
--model gemma-3-270m-mlx \
--batch-size 256 \
--gradient-accumulation-steps 16 \
--steps 2000 \
--eval-steps 100 \
--save-steps 500 \
--eval-tasks NanoMSMARCORetrieval \
--skip-eval-init \
--wandbThe important parameters are not specific to one repository. batch-size and gradient-accumulation-steps define the effective contrastive batch. eval-steps decides how quickly regressions surface. save-steps controls how much experiment history remains available for comparison or rollback.
Asymmetric retrieval benefits from role markers. The README describes query and document prompt tokens, with final embeddings generated by mean pooling while masking ordinary boundary and padding tokens. Preserving role markers lets the model learn that a short query and a long passage do not play interchangeable roles even if they share vocabulary.
def format_pair(query, document):
# Token names mirror the role-marker idea, not a required tokenizer API.
query_text = "<bos><unused0>" + query + "<eos>"
doc_text = "<bos><unused1>" + document + "<eos>"
return query_text, doc_text
def mean_pool(hidden, attention_mask, role_token_mask):
# Mask padding and ordinary boundaries, but keep role markers when desired.
usable = attention_mask * role_token_mask
return (hidden * usable[..., None]).sum(axis=1) / usable.sum(axis=1, keepdims=True)The details vary by tokenizer and model family, but the design principle is stable: make the role explicit before tokenization, and keep pooling behavior consistent between training, evaluation, and inference.
Teams with private documents can fine-tune a retrieval adapter without sending training examples to an external GPU service. This is useful for legal, medical, enterprise, or internal-code corpora where the hard examples are sensitive. Local training does not remove the need for governance, but it reduces the number of systems that receive raw query-document pairs.
General embedding models often underperform when queries use local jargon, product names, abbreviations, or support-ticket phrasing. A retrieval adapter can learn that domain vocabulary by training on clicked documents, expert-labeled pairs, synthetic question-answer pairs, or curated hard negatives. The result should be measured on the target retrieval task, not only on a broad leaderboard.
Local MLX training is well suited to short experiment cycles: try a negative-mining recipe, run several hundred or thousand steps, evaluate, inspect failures, and repeat. That loop is harder when every trial needs a cloud training job, artifact upload, and conversion back into a local format.
The source project can stream training data from local JSONL files or Elasticsearch indices. Remote streaming is convenient for large corpora and live mining, but the README notes that network latency can make token throughput vary. For stable benchmarking, freeze a local JSONL shard and measure throughput without remote I/O in the critical path.
{"query": "how do I rotate an api key?", "positive": "Rotate keys from the account security page...", "negatives": ["Create a new project from the dashboard...", "Billing invoices are exported monthly..."]}A useful data loader should validate that every row has a query, at least one positive, and a negative set that is not accidentally identical to the positive. It should also preserve source IDs so failures can be traced back to the corpus.
MTEB integration in the source repository is important because embedding loss alone can be misleading. A lower contrastive loss may reflect easier batches rather than better retrieval. Running a fixed retrieval task every few checkpoints gives a stable signal for early stopping, adapter selection, and regression detection.
for step in training_steps:
loss = train_one_step(batch)
if step % eval_steps == 0:
metrics = run_retrieval_eval(model, task="NanoMSMARCORetrieval")
log({"step": step, "loss": loss, "ndcg_at_10": metrics["ndcg_at_10"]})