| 1 | # Pooling Models |
| 2 | |
| 3 | !!! note |
| 4 | We currently support pooling models primarily for convenience. This is not guaranteed to provide any performance |
| 5 | improvements over using Hugging Face Transformers or Sentence Transformers directly. |
| 6 | We plan to optimize pooling models in vLLM. Please comment on <https://github.com/vllm-project/vllm/issues/21796> if you have any suggestions! |
| 7 | |
| 8 | ## What are pooling models? |
| 9 | |
| 10 | Natural Language Processing (NLP) can be primarily divided into the following two types of tasks: |
| 11 | |
| 12 | - Natural Language Understanding (NLU) |
| 13 | - Natural Language Generation (NLG) |
| 14 | |
| 15 | The generative models supported by vLLM cover a variety of task types, such as the large language models (LLMs) we are |
| 16 | familiar with, multimodal models (VLM) that handle multimodal inputs like images, videos, and audio, speech-to-text |
| 17 | transcription models, and real-time models that support streaming input. Their common feature is the ability to generate |
| 18 | text. Taking it a step further, vLLM-Omni supports the generation of multimodal content, including images, videos, and audio. |
| 19 | |
| 20 | As the capabilities of generative models continue to improve, the boundaries of these models are also constantly expanding. |
| 21 | However, certain application scenarios still require specialized small language models to efficiently complete specific tasks. |
| 22 | These models typically have the following characteristics: |
| 23 | |
| 24 | - They do not require content generation. |
| 25 | - They only need to perform very limited functions, without requiring strong generalization, creativity, or high intelligence. |
| 26 | - They demand extremely low latency and may operate on cost-constrained hardware. |
| 27 | - Text-only models typically have fewer than 1 billion parameters, while multimodal models generally have fewer than 10 billion parameters. |
| 28 | |
| 29 | Although these models are relatively small in scale, they are still based on the Transformer architecture, similar or |
| 30 | even identical to the most advanced large language models today. Many recently released pooling models are also fine-tuned |
| 31 | from large language models, allowing them to benefit from the continuous improvements in large models. This architecture |
| 32 | similarity enables them to reuse much of vLLM’s infrastructure. If compatible, we would be happy to help them leverage |
| 33 | the latest features of vLLM as well. |
| 34 | |
| 35 | ### Cheat Sheet |
| 36 | |
| 37 | As illustrated in the figure below, we have summarized the relationships among the key elements of pooling models as a takeaway. |
| 38 | |
| 39 |  |
| 40 | |
| 41 | ### Sequence-wise Task and Token-wise Task |
| 42 | |
| 43 | The key distinction between sequence-wise task and token-wise task lies in their output granularity: sequence-wise task |
| 44 | produces a single result for an entire input sequence, whereas token-wise task yields a result for each individual token |
| 45 | within the sequence. |
| 46 | |
| 47 | Many Pooling models support both (sequence) task and token task. When the default pooling task (e.g. a sequence-wise task) |
| 48 | is not what you want, you need to manually specify (e.g. a token-wise task) via `PoolerConfig(task=<task>)` offline or |
| 49 | `--pooler-config.task <task>` online. |
| 50 | |
| 51 | Of course, we also have "plugin" tasks that allow users to customize input and output processors. For more information, |
| 52 | please refer to [IO Processor Plugins](../../design/io_processor_plugins.md). |
| 53 | |
| 54 | ### Pooling Tasks |
| 55 | |
| 56 | | Pooling Tasks | Granularity | Outputs | |
| 57 | |-----------------------|---------------|-------------------------------------------------| |
| 58 | | `classify` (see note) | Sequence-wise | probability vector of classes for each sequence | |
| 59 | | `embed` | Sequence-wise | vector representations for each sequence | |
| 60 | | `token_classify` | Token-wise | probability vector of classes for each token | |
| 61 | | `token_embed` | Token-wise | vector representations for each token | |
| 62 | |
| 63 | !!! note |
| 64 | Within classification tasks, there is a specialized subcategory: Cross-encoder (aka reranker) models. These models |
| 65 | are a subset of classification models that accept two prompts as input and output num_labels equal to 1. |
| 66 | |
| 67 | ### Pooling Types |
| 68 | |
| 69 |  |
| 70 | |
| 71 | | Pooling Tasks | Granularity | Description | |
| 72 | |----------------|---------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| |
| 73 | | `CLS` pooling | Sequence-wise | For BERT‑like (bidirectional self‑attention) models, CLS pooling is used by default. This means the last_hidden_states corresponding to the first token (the [CLS] token) is taken as the output. | |
| 74 | | `LAST` pooling | Sequence-wise | For GPT‑like (causal self‑attention) models, LAST pooling is used by default. This means the last_hidden_states corresponding to the last token is taken as the output. | |
| 75 | | `MEAN` pooling | Sequence-wise | Many studies have shown that averaging the last_hidden_states over all input tokens performs better on certain downstream tasks. Therefore, more and more models are using MEAN pooling. | |
| 76 | | `ALL` pooling | Token-wise | Outputs the last_hidden_states for all input tokens. | |
| 77 | | `STEP` pooling | Token-wise | Filters and outputs the last_hidden_states corresponding to the token IDs returned by returned_token_ids. | |
| 78 | |
| 79 | ### Score Types |
| 80 | |
| 81 |  |
| 82 | |
| 83 | The scoring models is designed to compute similarity scores between two input prompts. It supports three model types |
| 84 | (aka `score_type`): `cross-encoder`, `late-interaction`, and `bi-encoder`. |
| 85 | |
| 86 | | Pooling Tasks | Granularity | Outputs | Score Types | scoring function | |
| 87 | |-----------------------|---------------|----------------------------------------------|--------------------|--------------------------| |
| 88 | | `classify` (see note) | Sequence-wise | reranker score for each sequence | `cross-encoder` | linear classifier | |
| 89 | | `embed` | Sequence-wise | vector representations for each sequence | `bi-encoder` | cosine similarity | |
| 90 | | `token_classify` | Token-wise | probability vector of classes for each token | N/A | N/A | |
| 91 | | `token_embed` | Token-wise | vector representations for each token | `late-interaction` | late interaction(MaxSim) | |
| 92 | |
| 93 | !!! note |
| 94 | Only when a classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled. |
| 95 | |
| 96 | ### Pooling Usages |
| 97 | |
| 98 | | Pooling Usages | Description | |
| 99 | |-----------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------| |
| 100 | | Classification Usages | Predicting which predefined category, class, or label best corresponds to a given input. | |
| 101 | | Embedding Usages | Converts unstructured data (text, images, audio, etc.) into structured numerical vectors (embeddings). | |
| 102 | | Token Classification Usages | Token-wise classification | |
| 103 | | Token Embedding Usages | Token-wise embedding | |
| 104 | | Reward Usages | Evaluates the quality of outputs generated by a language model, acting as a proxy for human preferences. | |
| 105 | | Scoring Usages | Computes similarity scores between two inputs. It supports three model types (aka `score_type`): `cross-encoder`, `late-interaction`, and `bi-encoder`. | |
| 106 | | Plugins Usages | Allow users to customize input and output processors. For more information, please refer to [IO Processor Plugins](../../design/io_processor_plugins.md). | |
| 107 | |
| 108 | We also have some special models that support multiple pooling tasks, or have specific usage scenarios, or support special inputs and outputs. |
| 109 | |
| 110 | For more detailed information, please refer to the link below. |
| 111 | |
| 112 | - [Classification Usages](classify.md) |
| 113 | - [Embedding Usages](embed.md) |
| 114 | - [Token Classification Usages](token_classify.md) |
| 115 | - [Token Embedding Usages](token_embed.md) |
| 116 | - [Reward Usages](reward.md) |
| 117 | - [Scoring Usages](scoring.md) |
| 118 | - [Specific Model Examples](specific_models.md) |
| 119 | |
| 120 | ## Offline Inference |
| 121 | |
| 122 | Each pooling model in vLLM supports one or more of these tasks according to |
| 123 | [Pooler.get_supported_tasks][vllm.model_executor.layers.pooler.Pooler.get_supported_tasks], |
| 124 | enabling the corresponding APIs. |
| 125 | |
| 126 | ### Offline APIs corresponding to pooling usages |
| 127 | |
| 128 | | Pooling Usages | Dedicated API | Pooling task for `LLM.encode` API | Score Types | scoring function | |
| 129 | |-----------------------------|---------------------|-----------------------------------|----------------------------|--------------------------| |
| 130 | | Classification Usages | `LLM.classify(...)` | `classify` | `cross-encoder` (see note) | linear classifier | |
| 131 | | Embedding Usages | `LLM.embed(...)` | `embed` | `bi-encoder` | cosine similarity | |
| 132 | | Token Classification Usages | N/A | `token_classify` | N/A | N/A | |
| 133 | | Token Embedding Usages | N/A | `token_embed` | `late-interaction` | late interaction(MaxSim) | |
| 134 | | Reward Usages | N/A | `classify` & `token_classify` | N/A | N/A | |
| 135 | | Scoring Usages | `LLM.score(...)` | N/A | N/A | N/A | |
| 136 | | Plugins Usages | N/A | `plugin` | N/A | N/A | |
| 137 | |
| 138 | !!! note |
| 139 | Only when a classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled. |
| 140 | |
| 141 | ### `LLM.classify` |
| 142 | |
| 143 | The [classify][vllm.LLM.classify] method outputs a probability vector for each prompt. |
| 144 | It is primarily designed for [classification models](classify.md). |
| 145 | For more information about `LLM.classify`, see [this page](classify.md#offline-inference). |
| 146 | |
| 147 | ### `LLM.embed` |
| 148 | |
| 149 | The [embed][vllm.LLM.embed] method outputs an embedding vector for each prompt. |
| 150 | It is primarily designed for [embedding models](embed.md). |
| 151 | For more information about `LLM.embed`, see [this page](embed.md#offline-inference). |
| 152 | |
| 153 | ### `LLM.score` |
| 154 | |
| 155 | The [score][vllm.LLM.score] method outputs similarity scores between sentence pairs. |
| 156 | It is primarily designed for [score models](scoring.md). |
| 157 | |
| 158 | ### `LLM.encode` |
| 159 | |
| 160 | The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM. |
| 161 | |
| 162 | Please use one of the more specific methods or set the task directly when using `LLM.encode`, refer to the [table above](#offline-apis-corresponding-to-pooling-usages). |
| 163 | |
| 164 | ### Examples |
| 165 | |
| 166 | ```python |
| 167 | from vllm import LLM |
| 168 | |
| 169 | llm = LLM(model="intfloat/e5-small", runner="pooling") |
| 170 | (output,) = llm.encode("Hello, my name is", pooling_task="embed") |
| 171 | |
| 172 | data = output.outputs.data |
| 173 | print(f"Data: {data!r}") |
| 174 | ``` |
| 175 | |
| 176 | ## Online Serving |
| 177 | |
| 178 | Our online Server provides endpoints that correspond to the offline APIs: |
| 179 | |
| 180 | - Corresponding to `LLM.embed`: |
| 181 | - [Cohere Embed API](embed.md#cohere-embed-api) (`/v2/embed`) |
| 182 | - [OpenAI-compatible Embeddings API](embed.md#openai-compatible-embeddings-api) (`/v1/embeddings`) |
| 183 | - Corresponding to `LLM.classify`: |
| 184 | - [Classification API](classify.md#online-serving)(`/classify`) |
| 185 | - Corresponding to `LLM.score`: |
| 186 | - [Score API](scoring.md#score-api) (`/score`, `/v1/score`) |
| 187 | - [Cohere Rerank API](scoring.md#cohere-rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`) |
| 188 | - Pooling API (`/pooling`) is similar to `LLM.encode`, being applicable to all types of pooling models. |
| 189 | |
| 190 | The following introduces the Pooling API. For other APIs, please refer to the link above. |
| 191 | |
| 192 | ### Pooling API |
| 193 | |
| 194 | Our Pooling API (`/pooling`) is similar to `LLM.encode`, being applicable to all types of pooling models. |
| 195 | |
| 196 | The input format is the same as [Embeddings API](embed.md#openai-compatible-embeddings-api), but the output data can contain an arbitrary nested list, not just a 1-D list of floats. |
| 197 | |
| 198 | Please use one of the more specific APIs or set the task directly when using the Pooling API, refer to the [table above](#offline-apis-corresponding-to-pooling-usages). |
| 199 | |
| 200 | Code examples: |
| 201 | |
| 202 | - [Online example](../../../examples/pooling/reward/token_reward_online.py) |
| 203 | - [Offline example](../../../examples/pooling/reward/token_reward_offline.py) |
| 204 | |
| 205 | ### Examples |
| 206 | |
| 207 | ```python |
| 208 | # start a supported embeddings model server with `vllm serve`, e.g. |
| 209 | # vllm serve intfloat/e5-small |
| 210 | import requests |
| 211 | |
| 212 | host = "localhost" |
| 213 | port = "8000" |
| 214 | model_name = "intfloat/e5-small" |
| 215 | |
| 216 | api_url = f"http://{host}:{port}/pooling" |
| 217 | |
| 218 | prompts = [ |
| 219 | "Hello, my name is", |
| 220 | "The president of the United States is", |
| 221 | "The capital of France is", |
| 222 | "The future of AI is", |
| 223 | ] |
| 224 | prompt = {"model": model_name, "input": prompts, "task": "embed"} |
| 225 | |
| 226 | response = requests.post(api_url, json=prompt) |
| 227 | |
| 228 | for output in response.json()["data"]: |
| 229 | data = output["data"] |
| 230 | print(f"Data: {data!r} (size={len(data)})") |
| 231 | ``` |
| 232 | |
| 233 | ## Configuration |
| 234 | |
| 235 | In vLLM, pooling models implement the [VllmModelForPooling][vllm.model_executor.models.VllmModelForPooling] interface. |
| 236 | These models use a [Pooler][vllm.model_executor.layers.pooler.Pooler] to extract the final hidden states of the input |
| 237 | before returning them. |
| 238 | |
| 239 | ### Model Runner |
| 240 | |
| 241 | Run a model in pooling mode via the option `--runner pooling`. |
| 242 | |
| 243 | !!! tip |
| 244 | There is no need to set this option in the vast majority of cases as vLLM can automatically |
| 245 | detect the appropriate model runner via `--runner auto`. |
| 246 | |
| 247 | ### Model Conversion |
| 248 | |
| 249 | vLLM can adapt models for various pooling tasks via the option `--convert <type>`. |
| 250 | |
| 251 | If `--runner pooling` has been set (manually or automatically) but the model does not implement the |
| 252 | [VllmModelForPooling][vllm.model_executor.models.VllmModelForPooling] interface, |
| 253 | vLLM will attempt to automatically convert the model according to the architecture names |
| 254 | shown in the table below. |
| 255 | |
| 256 | | Architecture | `--convert` | Supported pooling tasks | |
| 257 | |-------------------------------------------------|-------------|------------------------------| |
| 258 | | `*ForTextEncoding`, `*EmbeddingModel`, `*Model` | `embed` | `token_embed`, `embed` | |
| 259 | | `*ForRewardModeling`, `*RewardModel` | `embed` | `token_embed`, `embed` | |
| 260 | | `*For*Classification`, `*ClassificationModel` | `classify` | `token_classify`, `classify` | |
| 261 | |
| 262 | !!! tip |
| 263 | You can explicitly set `--convert <type>` to specify how to convert the model. |
| 264 | |
| 265 | ### Pooler Configuration |
| 266 | |
| 267 | #### Predefined models |
| 268 | |
| 269 | If the [Pooler][vllm.model_executor.layers.pooler.Pooler] defined by the model accepts `pooler_config`, |
| 270 | you can override some of its attributes via the `--pooler-config` option. |
| 271 | |
| 272 | #### Converted models |
| 273 | |
| 274 | If the model has been converted via `--convert` (see above), |
| 275 | the pooler assigned to each task has the following attributes by default: |
| 276 | |
| 277 | | Task | Pooling Type | Normalization | Softmax | |
| 278 | | ---------- | ------------ | ------------- | ------- | |
| 279 | | `embed` | `LAST` | ✅︎ | ❌ | |
| 280 | | `classify` | `LAST` | ❌ | ✅︎ | |
| 281 | |
| 282 | #### Resolution precedence |
| 283 | |
| 284 | The pooling method and `use_activation` are resolved per field. An explicitly |
| 285 | set field in `--pooler-config` takes precedence over Sentence Transformers |
| 286 | metadata, which in turn takes precedence over the model architecture or task |
| 287 | default. Fields left unset continue through the chain independently. |
| 288 | |
| 289 | The current `PoolerConfig` has no `normalize` or `activation` field. |
| 290 | `use_activation` controls whether the task's constructed normalization or |
| 291 | classification activation is applied. |
| 292 | |
| 293 | | Field | Source precedence | How to override | |
| 294 | | ----- | ----------------- | --------------- | |
| 295 | | Pooling method (`pooling_type`) | `--pooler-config` > compact `pooling_mode` or legacy boolean `pooling_mode_*` fields in the Pooling module referenced by Sentence Transformers `modules.json` > architecture default (`LAST` for sequence pooling and `ALL` for token pooling unless the architecture overrides it) | Set `{"pooling_type": "CLS"}`, or set `seq_pooling_type` / `tok_pooling_type` explicitly. | |
| 296 | | Embedding normalization (`use_activation`) | `--pooler-config` > Sentence Transformers modules (`true` when a Normalize module is present, otherwise `false`) > pooling-task default (`true`) when no Sentence Transformers Pooling module is found | Set `{"use_activation": false}` to return unnormalized embeddings. | |
| 297 | | Classification activation function | Hugging Face `problem_type` > Sentence Transformers activation metadata > sigmoid or softmax selected from the label count | The function cannot be selected through `--pooler-config`; set `{"use_activation": false}` to return logits instead. | |
| 298 | |
| 299 | Both compact `pooling_mode` values written by Sentence Transformers 5.4+ and |
| 300 | legacy boolean `pooling_mode_*` fields are supported. |
| 301 | |
| 302 | For converted models and predefined models using the standard DispatchPooler |
| 303 | adapters, `embed` and `token_embed` construct an L2-normalization head, while |
| 304 | `classify` and `token_classify` construct the selected classification activation. |
| 305 | In both cases, `use_activation` controls whether that head is applied. Models with |
| 306 | custom poolers can implement different behavior. |
| 307 | |
| 308 | To inspect the resolved fields without loading model weights: |
| 309 | |
| 310 | ```python |
| 311 | from vllm.config import ModelConfig, PoolerConfig |
| 312 | from vllm.model_executor.layers.pooler.activations import get_act_fn |
| 313 | |
| 314 | |
| 315 | def inspect(requested: PoolerConfig) -> None: |
| 316 | model_config = ModelConfig( |
| 317 | "intfloat/e5-small", |
| 318 | runner="pooling", |
| 319 | pooler_config=requested, |
| 320 | ) |
| 321 | resolved = model_config.pooler_config |
| 322 | assert resolved is not None |
| 323 | print( |
| 324 | { |
| 325 | "seq_pooling_type": resolved.seq_pooling_type, |
| 326 | "tok_pooling_type": resolved.tok_pooling_type, |
| 327 | "use_activation": resolved.use_activation, |
| 328 | "sequence_classification_activation": type( |
| 329 | get_act_fn(model_config.hf_config) |
| 330 | ).__name__, |
| 331 | } |
| 332 | ) |
| 333 | |
| 334 | |
| 335 | inspect(PoolerConfig()) |
| 336 | inspect(PoolerConfig(pooling_type="CLS", use_activation=False)) |
| 337 | ``` |
| 338 | |
| 339 | For `intfloat/e5-small`, the first result contains `MEAN`, `ALL`, and `True`. |
| 340 | The second contains `CLS`, `ALL`, and `False`. Both report the classification |
| 341 | activation that the standard sequence-classification adapter would construct. |
| 342 | |
| 343 | ## Removed Features |
| 344 | |
| 345 | ### Encode task |
| 346 | |
| 347 | We have split the `encode` task into two more specific token-wise tasks: `token_embed` and `token_classify`: |
| 348 | |
| 349 | - `token_embed` is the same as `embed`, using normalization as the activation. |
| 350 | - `token_classify` is the same as `classify`, by default using softmax as the activation. |
| 351 | |
| 352 | Pooling models now support token-wise task. |
| 353 | |
| 354 | - Extracting hidden states prefers using `token_embed` task. |
| 355 | - Named Entity Recognition (NER) and reward models prefers using `token_classify` task. |
| 356 | |
| 357 | ### Score task |
| 358 | |
| 359 | `score` task has been removed in v0.21, use `classify` instead. Only when a classification model outputs num_labels |
| 360 | equal to 1 can it be used as a scoring model and have its scoring API enabled. |
| 361 | |
| 362 | ### Pooling multitask support |
| 363 | |
| 364 | Pooling multitask support has been removed in v0.21. When the default pooling task is not what you want, |
| 365 | you need to manually specify it via `PoolerConfig(task=<task>)` offline or `--pooler-config.task <task>` online. |