text_expansion export missing SPLADE max-pooling — third-party SPLADE models produce severely degraded vectors

[ah89]

When a third-party SPLADE model (e.g. naver/splade-v3) is exported via eland_import_hub_model --task-type text_expansion, the resulting sparse vectors contain roughly 20% of the tokens that the equivalent SentenceTransformers.SparseEncoder produces, causing severe recall collapse in sparse search.

Root cause

Eland exports AutoModelForMaskedLM.forward() unchanged:

# eland/ml/pytorch/transformers.py
if self._task_type == "text_expansion":
    model = transformers.AutoModelForMaskedLM.from_pretrained(...)
    return _TraceableTextExpansionModel(self._tokenizer, model)

_TraceableTextExpansionModel traces the model as-is with torch.jit.trace. The exported model outputs raw MLM logits of shape [batch, seq_len, vocab_size].

The Elasticsearch inference pipeline then:

1. ml-cpp runs the model and passes the tensor through unchanged.
2. Java's TextExpansionProcessor reads result[0][0] — sequence position 0, i.e. the CLS token — and applies a score > 0 threshold.

No max-pooling is performed anywhere in the ES inference pipeline — it always reads position 0, so the exported model must place the full sparse vector there. Standard SPLADE models instead distribute vocabulary expansion across all sequence positions and rely on max-pooling (w_j = max_i log(1 + ReLU(logit_{i,j}))) to combine them. By reading only position 0, Elasticsearch gets a strict mathematical subset of the correct vector (since max(positions) >= position[0] always holds), causing the token collapse observed.

Notably, Eland already handles this for the text_embedding task — it wraps the model with the appropriate pooling/normalization layer before tracing (_SentenceTransformerWrapperModule). The text_expansion path has no equivalent: it traces the bare AutoModelForMaskedLM, so the SPLADE pooling step is simply omitted. (ELSER and Elastic's own packaged sparse models work through this same pipeline because they are built and validated by Elastic specifically for it; a generic third-party HF import has no such guarantee.)

Evidence

The gap is reproducible directly against naver/splade-v3 — what ES reads (position 0) vs. what SPLADE intends (max-pool over positions):

from transformers import AutoModelForMaskedLM, AutoTokenizer
import torch, torch.nn.functional as F

tok = AutoTokenizer.from_pretrained("naver/splade-v3")
model = AutoModelForMaskedLM.from_pretrained("naver/splade-v3")

enc = tok("an example document", return_tensors="pt")
logits = model(**enc).logits[0]                      # [seq_len, vocab]

cls     = (F.relu(logits[0]) > 0).sum()              # what ES keeps (position 0 / CLS)
maxpool = (F.relu(logits).max(0).values > 0).sum()   # what SPLADE intends

print(cls.item(), maxpool.item())                    # cls is always a subset of maxpool

Customer's real-world observation with naver/splade-v3:

• Local SparseEncoder: ~387 non-zero tokens per document
• ES stored vector: ~82 non-zero tokens — a strict subset of the local set (100% overlap), the exact signature of taking position 0 instead of max-pooling.

Possible fix

_TraceableTextExpansionModel can wrap the model to apply the SPLADE activation and max-pool before tracing, so the exported graph emits an already-pooled vector:

class _SpladeWrapper(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, input_ids, attention_mask, token_type_ids, position_ids):
        logits = self.model(input_ids, attention_mask, token_type_ids, position_ids)[0]
        # SPLADE activation + max-pool across sequence positions
        activated = torch.log(1 + torch.relu(logits))
        pooled = activated.max(dim=1, keepdim=True).values   # [batch, 1, vocab]
        return (pooled,)

This makes the exported model output [batch, 1, vocab_size], so result[0][0] in Java gives the correct max-pooled SPLADE vector. The ops this introduces (aten::max, aten::log, aten::relu) are already in ml-cpp's operation allowlist, so a pooled export still passes model-graph validation. (The forward signature should mirror the inputs Eland already passes for the model's tokenizer type — e.g. 2 args for RoBERTa/DistilBERT-based models.)