[ML] SPLADE embedding support

[daixque]

Overview

Elasticsearch supports sparse embeddings including non-ELSER models which is introduced by #116935.

A significant example of sparse vector model is the SPLADE model, which is a reference model for ELSER. But the format of the output of SPLADE is not same as ELSER, so we need to implement a post processing step for it.

Interface change

This PR introduces a new expansion_type parameter to trained model resource.

GET /_ml/trained_models/hotchpotch__japanese-splade-v2?pretty

{
  "count" : 1,
  "trained_model_configs" : [
    {
      "model_id" : "hotchpotch__japanese-splade-v2",
      "model_type" : "pytorch",
      "created_by" : "api_user",
      "version" : "12.0.0",
      "create_time" : 1753171035706,
      "model_size_bytes" : 0,
      "estimated_operations" : 0,
      "license_level" : "platinum",
      "description" : "Model hotchpotch/japanese-splade-v2 for task type 'text_expansion'",
      "tags" : [ ],
      "metadata" : {
        "per_allocation_memory_bytes" : 392242176,
        "per_deployment_memory_bytes" : 545637376
      },
      "input" : {
        "field_names" : [
          "text_field"
        ]
      },
      "inference_config" : {
        "text_expansion" : {
          "vocabulary" : {
            "index" : ".ml-inference-native-000002"
          },
          "tokenization" : {
            "bert_ja" : {
              "do_lower_case" : false,
              "with_special_tokens" : true,
              "max_sequence_length" : 512,
              "truncate" : "first",
              "span" : -1
            }
          },
          "expansion_type" : "splade"
        }
      },
      "location" : {
        "index" : {
          "name" : ".ml-inference-native-000002"
        }
      }
    }
  ]
}

expansion_type can be one of elser or splade. If it is not specified, it defaults to elser.

To use the SPLADE model, eland needs to be updated to support the expansion_type parameter.
elastic/eland#802

Logic for SPLADE

SPLADE model outputs a embedding in the shape of [1, input_token_size, vocab_size]. The second dimention is different from ELSER, which is [1, chunk_size, vocab_size].
For SPLADE, we need to apply the saturation function to the output, which is log(1 + relu(x)), and then apply max pooling to the second dimension.

To be considered

• Currently, we only implement max pooling for SPLADE. Should we also implement sum and/or other pooling?
• Top-k expansion is not supported. Some models produce a large amount of output from small input. Should we implement it? If yes, we need to have another option, and it can not be determined by eland automatically.
• I observed OOM when I send infer request with a little long text for japanese-splade-v2 model.

Reference

• I referred this for the spladeMaxPooling implementation.

[elasticsearchmachine]

Pinging @elastic/ml-core (Team:ML)

[davidkyle]

Changing the expansionType changes the way the the output is processed and may not be compatible with the different types of model. Switching this value for the ELSER model would break the processing.

Fixing expansionType when the config is created and not allowing it to be overridden at inference is good enough. If the wrong expansionType value is set the user will have to recreate the model

[davidkyle]

Same comment as above regarding updating expansion type. I think the changes in this file can be reverted.

[davidkyle]

In a mixed version cluster older nodes will not expect this new field to be serialised. We protect against this using TransportVersions. Pls add a new version to TransportVersions.java at

elasticsearch/server/src/main/java/org/elasticsearch/TransportVersions.java

Line 350 in b5b636a

public static final TransportVersion ESQL_TOPN_TIMINGS = def(9_128_0_00);

Suggested change

	expansionType = in.readOptionalString();
	if (in.getTransportVersion().onOrAfter(TransportVersions.ML_EXPANSION_TYPE)) {
	expansionType = in.readOptionalString();
	}

[davidkyle]

Suggested change

	out.writeOptionalString(expansionType);
	if (out.getTransportVersion().onOrAfter(TransportVersions. ML_EXPANSION_TYPE)) {
	out.writeOptionalString(expansionType);
	}