A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 06:00 UTC
feat: rename logit_bias/logit_scale to logit_mean/logit_sigma for affine score calibration (#39530)
Signed-off-by: Jesus Federico <jefp@amazon.com> Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>
Jesus Federico committed
Apr 13, 2026 at 00:43 UTC
fa6ae3117762f12ce9beebf77e255cd1370b6276
7 files changed
+86
-35
docs/models/pooling_models/classify.md
+10
-8
@@ -273,27 +273,25 @@ Affine Score Calibration, also known as [Platt Scaling](https://en.wikipedia.org
273
274
The calibration follows the transformation:
275
276
-`activation(logit_scale * (logit - logit_bias))`
276
+`activation((logit - logit_mean) / logit_sigma)`
277
278
| Parameter | Default | Description |
279
| --------- | ------- | ----------- |
280
-| `logit_bias` | `None` | Bias subtracted from logits before activation |
281
-| `logit_scale` | `None` | Scale factor applied to logits after bias subtraction |
282
-
283
-Note: `logit_bias` is **subtracted** from the logits (not added), consistent with the `sigmoid_normalize` convention where `sigmoid(x - bias)` centers the sigmoid around the bias value.
280
+| `logit_mean` | `None` | Mean subtracted from logits (centers scores) |
281
+| `logit_sigma` | `None` | Standard deviation used to scale logits after mean subtraction |
282
283
The computation order is as follows:
284
285
```python
288
-logits -= logit_bias # subtract bias (center scores)
289
-logits *= logit_scale # scale logits
286
+logits -= logit_mean # subtract mean (center scores)
287
+logits /= logit_sigma # divide by sigma (scale)
288
logits = activation(logits) # e.g. sigmoid
289
```
290
291
Example configuration:
292
293
```bash
296
---pooler-config '{"use_activation": true, "logit_bias": 4.5, "logit_scale": 1.0}'
294
+--pooler-config '{"use_activation": true, "logit_mean": 4.5, "logit_sigma": 1.0}'
295
```
296
297
## Removed Features
@@ -301,3 +299,7 @@ Example configuration:
299
### Remove softmax from PoolingParams
300
301
We have already removed `softmax` and `activation` from PoolingParams. Instead, use `use_activation`, since we allow `classify` and `token_classify` to use any activation function.
302
+
303
+### Remove `logit_bias` and `logit_scale`
304
+
305
+`logit_bias` and `logit_scale` are deprecated aliases for `logit_mean` and `logit_sigma` respectively. When using `logit_scale`, it is automatically converted to `logit_sigma = 1/logit_scale`. These deprecated parameters will be removed in v0.21.
vllm/config/pooler.py
+52
-5
@@ -77,17 +77,31 @@ class PoolerConfig:
77
Defaults to None (i.e. set to max_model_len).
78
"""
79
80
- ## for classification models
80
+ ## for classification models — affine score calibration
81
+ logit_mean: float | None = None
82
+ """
83
+ If provided, subtract this value from classification logits before
84
+ activation. Used for affine score calibration (Platt scaling):
85
+ activation((logit - logit_mean) / logit_sigma). Defaults to None.
86
+ """
87
+
88
+ logit_sigma: float | None = None
89
+ """
90
+ If provided, divide the classification logits by this value after
91
+ mean subtraction. Used for affine score calibration (Platt scaling):
92
+ activation((logit - logit_mean) / logit_sigma). Defaults to None.
93
+ """
94
+
95
+ # Deprecated aliases — will be removed in v0.21
96
logit_bias: float | None = None
97
"""
83
- If provided, apply classification logit biases. Defaults to None.
98
+ Deprecated: Use logit_mean instead. Will be removed in v0.21.
99
"""
100
101
logit_scale: float | None = None
102
"""
88
- If provided, scale the classification logits by this factor before
89
- activation. Combined with logit_bias, enables affine score calibration:
90
- activation(logit_scale * (score - logit_bias)). Defaults to None.
103
+ Deprecated: Use logit_sigma instead (note: logit_sigma = 1/logit_scale).
104
+ Will be removed in v0.21.
105
"""
106
107
## for reward models
@@ -105,6 +119,39 @@ class PoolerConfig:
119
"""
120
121
def __post_init__(self) -> None:
122
+ # Handle deprecated logit_bias → logit_mean
123
+ if self.logit_bias is not None:
124
+ if self.logit_mean is not None:
125
+ raise ValueError(
126
+ "Cannot set both `logit_bias` and `logit_mean`. "
127
+ "`logit_bias` is deprecated, use `logit_mean` instead."
128
+ )
129
+ logger.warning(
130
+ "`logit_bias` is deprecated and will be removed in v0.21. "
131
+ "Use `logit_mean` instead."
132
+ )
133
+ self.logit_mean = self.logit_bias
134
+ self.logit_bias = None
135
+
136
+ # Handle deprecated logit_scale → logit_sigma
137
+ if self.logit_scale is not None:
138
+ if self.logit_sigma is not None:
139
+ raise ValueError(
140
+ "Cannot set both `logit_scale` and `logit_sigma`. "
141
+ "`logit_scale` is deprecated, use `logit_sigma` instead."
142
+ )
143
+ logger.warning(
144
+ "`logit_scale` is deprecated and will be removed in v0.21. "
145
+ "Use `logit_sigma` instead (logit_sigma = 1/logit_scale)."
146
+ )
147
+ if self.logit_scale == 0:
148
+ raise ValueError("logit_scale cannot be 0 (division by zero)")
149
+ self.logit_sigma = 1.0 / self.logit_scale
150
+ self.logit_scale = None
151
+
152
+ if self.logit_sigma is not None and self.logit_sigma == 0:
153
+ raise ValueError("logit_sigma cannot be 0 (division by zero)")
154
+
155
if pooling_type := self.pooling_type:
156
if self.seq_pooling_type is not None:
157
raise ValueError(
vllm/model_executor/layers/pooler/seqwise/heads.py
+9
-8
@@ -103,16 +103,16 @@ class ClassifierPoolerHead(SequencePoolerHead):
103
def __init__(
104
self,
105
classifier: ClassifierFn | None = None,
106
- logit_bias: float | None = None,
107
- logit_scale: float | None = None,
106
+ logit_mean: float | None = None,
107
+ logit_sigma: float | None = None,
108
head_dtype: torch.dtype | str | None = None,
109
activation: ActivationFn | None = None,
110
) -> None:
111
super().__init__()
112
113
self.classifier = classifier
114
- self.logit_bias = logit_bias
115
- self.logit_scale = logit_scale
114
+ self.logit_mean = logit_mean
115
+ self.logit_sigma = logit_sigma
116
self.head_dtype = head_dtype
117
self.activation = activation
118
@@ -140,10 +140,11 @@ class ClassifierPoolerHead(SequencePoolerHead):
140
logits = pooled_data
141
142
# logits shape: [batchsize, num_labels]
143
- if self.logit_bias is not None:
144
- logits -= self.logit_bias
145
- if self.logit_scale is not None:
146
- logits *= self.logit_scale
143
+ # Affine score calibration: activation((logit - mean) / sigma)
144
+ if self.logit_mean is not None:
145
+ logits = logits - self.logit_mean
146
+ if self.logit_sigma is not None:
147
+ logits = logits / self.logit_sigma
148
149
if self.activation is not None:
150
flags = [p.use_activation for p in pooling_params]
vllm/model_executor/layers/pooler/seqwise/poolers.py
+2
-2
@@ -118,8 +118,8 @@ def pooler_for_classify(
118
head = ClassifierPoolerHead(
119
head_dtype=model_config.head_dtype,
120
classifier=classifier,
121
- logit_bias=model_config.pooler_config.logit_bias,
122
- logit_scale=model_config.pooler_config.logit_scale,
121
+ logit_mean=model_config.pooler_config.logit_mean,
122
+ logit_sigma=model_config.pooler_config.logit_sigma,
123
activation=resolve_classifier_act_fn(
124
model_config, static_num_labels=True, act_fn=act_fn
125
),
vllm/model_executor/layers/pooler/tokwise/heads.py
+9
-8
@@ -92,16 +92,16 @@ class TokenClassifierPoolerHead(TokenPoolerHead):
92
def __init__(
93
self,
94
classifier: ClassifierFn | None = None,
95
- logit_bias: float | None = None,
96
- logit_scale: float | None = None,
95
+ logit_mean: float | None = None,
96
+ logit_sigma: float | None = None,
97
head_dtype: torch.dtype | str | None = None,
98
activation: ActivationFn | None = None,
99
) -> None:
100
super().__init__()
101
102
self.classifier = classifier
103
- self.logit_bias = logit_bias
104
- self.logit_scale = logit_scale
103
+ self.logit_mean = logit_mean
104
+ self.logit_sigma = logit_sigma
105
self.head_dtype = head_dtype
106
self.activation = activation
107
@@ -127,10 +127,11 @@ class TokenClassifierPoolerHead(TokenPoolerHead):
127
logits = pooled_data
128
# logits shape: [n_token, num_labels]
129
130
- if self.logit_bias is not None:
131
- logits -= self.logit_bias
132
- if self.logit_scale is not None:
133
- logits *= self.logit_scale
130
+ # Affine score calibration: activation((logit - mean) / sigma)
131
+ if self.logit_mean is not None:
132
+ logits = logits - self.logit_mean
133
+ if self.logit_sigma is not None:
134
+ logits = logits / self.logit_sigma
135
136
if self.activation is not None and pooling_param.use_activation:
137
logits = self.activation(logits)
vllm/model_executor/layers/pooler/tokwise/poolers.py
+2
-2
@@ -127,8 +127,8 @@ def pooler_for_token_classify(
127
head = TokenClassifierPoolerHead(
128
head_dtype=model_config.head_dtype,
129
classifier=classifier,
130
- logit_bias=model_config.pooler_config.logit_bias,
131
- logit_scale=model_config.pooler_config.logit_scale,
130
+ logit_mean=model_config.pooler_config.logit_mean,
131
+ logit_sigma=model_config.pooler_config.logit_sigma,
132
activation=resolve_classifier_act_fn(
133
model_config, static_num_labels=False, act_fn=act_fn
134
),
vllm/model_executor/models/config.py
+2
-2
@@ -232,8 +232,8 @@ class JinaVLForSequenceClassificationConfig(VerifyAndUpdateConfig):
232
config = model_config.hf_config
233
config.num_labels = 1
234
pooler_config = model_config.pooler_config
235
- if pooler_config.logit_bias is None:
236
- pooler_config.logit_bias = 2.65
235
+ if pooler_config.logit_mean is None:
236
+ pooler_config.logit_mean = 2.65
237
238
239
class LlamaBidirectionalConfig(VerifyAndUpdateConfig):