npc0/qwen3-embedding-rerank-projection

Qwen3-Embedding-0.6B + Soft-Label Rerank Projection

Experimental v2 is now available: `npc0/qwen3-embedding-rerank-projection-v2` (long run, early-stopped, BEIR-nq nDCG@10 0.4380 vs 0.4354 here).

A small projection head trained on frozen `Qwen/Qwen3-Embedding-0.6B` embeddings that turns a general-purpose bi-encoder into a stronger reranker, adding "embed documents once, rank with a dot product" property.

Update: replaced hard Y/N relevance labels with a SMOOTH soft target (kernel-smoothed over the similarity axis) beats the hard multi-similarity loss on out-of-domain BEIR-nq — nDCG@10 0.4354 vs 0.4257 (+0.0097), MRR@10 0.4021 vs 0.3919 (+0.0103), at identical recall@100.

Model details

The smooth label (this model)

Instead of q = onehot(positive), the target is a soft distribution over the candidate list that gives a small amount of mass to negatives sitting near the positive in similarity space — the ranking analogue of DreamSmooth's temporal reward smoothing:

m_k  = exp(-0.5 * ((s_neg_k - s_pos) / sigma)^2)      # gaussian kernel, frozen sims
w    = m / m.sum()                                     # distribution over negatives
q    = [1 - eps,  eps * w]                             # positive keeps 1-eps
loss = -Σ q_i · log_softmax(sim_i / T) · T^2           # Hinton-scaled soft CE

Two findings that matter:

1. 1.The temperature is not cosmetic. With cosine-similarity logits bounded in [-1,1] and T=1.0, softmax over ~30 candidates can place at most ~0.2 mass on the positive — so a target of q_pos = 0.7 is unreachable, the loss pins at its ceiling (2.9 vs a 1.6 entropy floor) and never converges. At T=0.05 the target is reachable and the smooth label overtakes the hard one.
2. 2.A smooth target alone cannot beat raw if it distills the evaluation space. Using the frozen embedding similarity as the teacher means the target is derived from exactly the ordering the eval measures; through a 1024→256 bottleneck the head can at best reproduce it. The gain here comes from the soft supervision regularizing the ranking, not from new information.

Held-out ablation (identical subset, seed, eval surface)

100k anchors (80k train / 20k val), 3 epochs, bs 512. Δ = projected − raw Hits@5. The hard-label control is multisim (rows mode); the rest are grouped.

kernel_smooth at T=1.0 scores −0.1354 — the temperature bug above, not a property of smooth labels. ema_teacher collapses (self-distillation with no supervised anchor).

Full-set held-out (864k train / 216k val anchors): kernel_smooth −0.0058, wsls −0.0070. These are held-out; the earlier hard-label MS figure of +0.039 was measured in-sample, so the two are not comparable.

BEIR-nq (BM25 top-100 candidates, 3,452 queries)

Every row scored on the identical candidate set with pytrec_eval:

Clean attribution of the smooth-label gain

The comparison above spans two training-data views, so we ran the missing control — the hard-label multi-positive MS loss on the same grouped view as the winner:

On an identical data view the smooth label wins by +0.0965 nDCG@10. The grouped hard-MS control is weak because each anchor contrasts against only its own ~24–100 negatives (no cross-batch pooling), and a listwise soft target copes with that far better than a hard pair-based objective. The smooth head also beats the strongest hard-label model (+0.0097), so both comparisons favour it.

Identical recall@100 confirms this is a pure ranking improvement, not a candidate-selection effect. The cross-encoder remains stronger (+0.147 nDCG@10) because it reads query and document jointly — a frozen-encoder projection cannot recover that signal.

Usage

import torch, torch.nn as nn, numpy as np
from safetensors.torch import load_file

class MLPProjection(nn.Module):
    def __init__(self, input_dim=1024, output_dim=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 512), nn.LayerNorm(512), nn.ReLU(), nn.Dropout(0.1),
            nn.Linear(512, 256), nn.LayerNorm(256), nn.ReLU(), nn.Dropout(0.1),
            nn.Linear(256, output_dim), nn.LayerNorm(output_dim),
        )
    def forward(self, x):
        z = self.net(x)
        return z / (z.norm(dim=-1, keepdim=True) + 1e-8)

head = MLPProjection().eval()
head.load_state_dict(load_file("model.safetensors"))

with torch.no_grad():
    doc_vecs   = head(torch.from_numpy(doc_embs)).numpy()
    query_vecs = head(torch.from_numpy(query_embs)).numpy()

doc_vecs   /= np.linalg.norm(doc_vecs, axis=1, keepdims=True) + 1e-8
query_vecs /= np.linalg.norm(query_vecs, axis=1, keepdims=True) + 1e-8
scores = query_vecs @ doc_vecs.T          # rerank with a single matmul

Base embeddings are int8-quantized with scale = 432.4123025768911 (dequantize: int8 / scale); config.json records it.

Limitations

• —The projection is a lossy 1024→256 bottleneck; it improves the bi-encoder but does not approach a cross-encoder.
• —The soft target is derived from the frozen encoder's own similarity, so it can regularize the ranking but cannot inject new relevance information.
• —The T² factor in the soft cross-entropy is a Hinton convention that assumes a temperature-scaled teacher; here the target is temperature-independent, so at T=0.05 it also shrinks the gradient 400× (an implicit learning-rate cut). Part of the T sweep is therefore a learning-rate effect, not only target reachability.
• —The 518,227 recovered extra positive pairs (extra_pos.npz) were built and tested but not fed to the winning model — the winner used only the multi_pos.npz grouped positives.
• —Full-set held-out validation uses a contiguous prefix of the anchor array (not a random split), and the rows-view control does not exclude val anchors from training; treat in-domain deltas as indicative, not precise.
• —sigma/eps/temperature were tuned on a 100k-anchor subset; a full sweep is future work.
• —Evaluation is English BEIR-nq with BM25 candidates; training data is mostly Chinese + MS MARCO + NQ.
• —Dataset padding convention: neg_id is zero-padded and 0 is a valid id — mask with neg_count.

License

Apache-2.0 for the projection head. Base model and datasets retain their original licenses.

Citation

If you used this in your research, please cite:

@misc{xu2026qwen3rerank,
  title        = {Qwen3-Embedding-0.6B Soft-Label Rerank Projection},
  author       = {Yuan Xu},
  year         = {2026},
  howpublished = {\url{https://huggingface.co/npc0/qwen3-embedding-rerank-projection}},
  note         = {Soft-label (kernel-smoothed) projection head on frozen Qwen3-Embedding-0.6B}
}

Experimental v2 is now available: `npc0/qwen3-embedding-rerank-projection-v2` (long run, early-stopped, BEIR-nq nDCG@10 0.4380 vs 0.4354 here).