[RFC] Support additional output formats for sparse models

[yuye-aws]

Background

The current sparse encoding model and sparse tokenizer model outputs the sparse embedding in the following format.

{
  "world": 3.4208686,
  "hello": 6.9377565
}

While this format is human-readable, it has the following limitations:

• Increased storage requirements due to storing string tokens
• Potential performance overhead when processing string keys
• Limited compatibility with some vector operations that expect numerical indices

We propose supporting two additional output formats. The first one is Integer-index format. This format replaces word tokens with their corresponding integer IDs from tokenizer, reducing storage requirements while maintaining the map structure.

{
  "2088": 6.9377565
  "7592": 3.4208686,
}

The second one is array-based format, representing a sparse vector with two lists: indices and values. This format represents the sparse vector with parallel arrays of indices and values.

{
  "indices": [2088, 7592],
  "values": [6.9377565, 3.4208686]
}

Options

Option 1: Runtime Format Selection via Request Parameters (Recommended)

Users can specify the desired output format at inference time through the sparse_encoding_format parameter:

POST _plugins/_ml/models/Lkjp8ZYBTIHkDc6TJ4q6/_predict
{
  "text_docs": ["hello world"],
  "parameters": {
    "sparse_encoding_format": "int" // Options: "word" (default), "int" and "array"
  }
} 

Advantages

• Flexibility to choose output format per request
• Backward compatible with existing models and implementation

Next Steps

I have implemented this PR as POC. If this option is chosen, I am going to continue working on the array format and sparse tokenizer model.

Option 2: Model Config (with technical challenge)

User can specify the output format during model registration:

POST /_plugins/_ml/models/_register?deploy=true
{
  "name": "amazon/neural-sparse/opensearch-neural-sparse-encoding-doc-v2-distill",
  "version": "1.0.0",
  "model_format": "TORCH_SCRIPT",
  "model_config": {
    "sparse_encoding_format": "int"
  }
}

Limitation

This approach has a technical challenge. For pre-trained models, there is no use specifying the model_config during model registration. The model config received by the TextEmbeddingSparseEncodingModel class will always be null. Besides, the model configuration of pre-trained models cannot be modified after registration.

Open questions

1. Should we support a numeric key format without quotes for the integer-index format?

{
  2088: 6.9377565
  7592: 3.4208686
}

2. For the array-based format, should we guarantee that indices are sorted in ascending order?
3. How should we handle the case where a user is using remote model?

[yuye-aws]

Hi maintainers! Please label this issue as RFC, thanks!

[brianf-aws]

Hi @yuye-aws wanted to share my thoughts for your RFC

For open question 1.
I think well be able to save more memory if we keep it as a int so 4 bytes every time. If we decide to use a string then we use as many bytes as the string contains characters plus the null terminated char (src). e.g. tokenId is "12345" then it takes 5 + 1 bytes. which is not bad if the word token is a complicated word like "ambidextrous".

However It depends on the search capability handling numbers instead of strings. Do you know if Neural Sparse Search can search with tokenIds?

src

My only concern is that the use for using neural sparse search with numbers as keys will be blocked when the ML Inference processor uses this capability because when it parses it uses a Map of String to Object

ml-commons/plugin/src/main/java/org/opensearch/ml/processor/ModelExecutor.java

Line 192 in b22e61a

Map<String, Object> modelTensorOutputMap = gson.fromJson(modelOutputJsonStr, Map.class);

This method ^ is used in the MLInferenceIngestProcessor, MLInferenceSearchRequestProcessor, MLInferenceSearchResponseProcessor. So to make this work it would use keys of strings but with the int representation.

[yuye-aws]

However It depends on the search capability handling numbers instead of strings. Do you know if Neural Sparse Search can search with tokenIds?

This is a valid concern, so I think the first step is to support string token id, instead of the integer token id.

[yuye-aws]

After discussion with @model-collapse @zhichao-aws @xinyual , the conclusion is to choose Option1 and implement two types: "array" and "map". I am going to work on the PR this week.

[zhichao-aws]

We also need to provide sample SageMaker deployment scripts for this new format.

There are also 2 options

• First option is using 2 dedicated scripts for 2 format, and users should pick the correct script based on the index
• Second option is using 1 script, and pass an extra parameter to specify the embedding format.

For me option 2 is perferred, but we need to first check whether we can pass parameters when calling remote connectors from neural-search.