Calibration

bensemble.calibration.scaling

TemperatureScaling

TemperatureScaling(init_temp: float = 1.5)

Bases: Module

Temperature Scaling for model calibration.

Divides logits by a single learnable scalar parameter T (temperature). This softens probabilities and calibrates confidence without changing classification accuracy (argmax remains identical).

Initializes the TemperatureScaling module.

Parameters:

Name	Type	Description	Default
`init_temp`	`float`	Initial value for the temperature scalar. Defaults to 1.5.	`1.5`

Source code in bensemble/calibration/scaling.py

def __init__(self, init_temp: float = 1.5):
    """Initializes the TemperatureScaling module.

    Args:
        init_temp: Initial value for the temperature scalar. Defaults to 1.5.
    """
    super().__init__()
    self.temperature = nn.Parameter(torch.ones(1) * init_temp)

fit

fit(
    logits: Tensor, labels: Tensor, max_iter: int = 50
) -> TemperatureScaling

Finds the optimal temperature T using a validation set.

Optimizes the negative log-likelihood via L-BFGS.

Parameters:

Name	Type	Description	Default
`logits`	`Tensor`	Unscaled logits from a hold-out validation set of shape (N, num_classes).	required
`labels`	`Tensor`	Ground truth class indices of shape (N,).	required
`max_iter`	`int`	Maximum number of L-BFGS iterations. Defaults to 50.	`50`

Returns:

Name	Type	Description
`TemperatureScaling`	`TemperatureScaling`	The fitted instance itself.

Source code in bensemble/calibration/scaling.py

def fit(
    self, logits: torch.Tensor, labels: torch.Tensor, max_iter: int = 50
) -> "TemperatureScaling":
    """Finds the optimal temperature T using a validation set.

    Optimizes the negative log-likelihood via L-BFGS.

    Args:
        logits: Unscaled logits from a hold-out validation set of shape
            (N, num_classes).
        labels: Ground truth class indices of shape (N,).
        max_iter: Maximum number of L-BFGS iterations. Defaults to 50.

    Returns:
        TemperatureScaling: The fitted instance itself.
    """
    logits = logits.detach()
    optimizer = optim.LBFGS([self.temperature], lr=0.01, max_iter=max_iter)

    def eval_loss():
        optimizer.zero_grad()
        scaled_logits = self.forward(logits)
        loss = F.cross_entropy(scaled_logits, labels)
        loss.backward()
        return loss

    optimizer.step(eval_loss)
    return self

forward

forward(logits: Tensor) -> torch.Tensor

Applies temperature scaling to the input logits.

Parameters:

Name	Type	Description	Default
`logits`	`Tensor`	Raw uncalibrated logits of shape (batch_size, num_classes).	required

Returns:

Type	Description
`Tensor`	torch.Tensor: Scaled logits of shape (batch_size, num_classes).

Source code in bensemble/calibration/scaling.py

def forward(self, logits: torch.Tensor) -> torch.Tensor:
    """Applies temperature scaling to the input logits.

    Args:
        logits: Raw uncalibrated logits of shape (batch_size, num_classes).

    Returns:
        torch.Tensor: Scaled logits of shape (batch_size, num_classes).
    """
    return logits / self.temperature

VectorScaling

VectorScaling(num_classes: int)

Bases: Module

Vector Scaling for multi-class calibration (extension of Platt Scaling).

Applies a per-class affine transformation to uncalibrated logits: calibrated_logits = logits * a + b

Initializes the VectorScaling module.

Parameters:

Name	Type	Description	Default
`num_classes`	`int`	Number of classes in the classification task.	required

Source code in bensemble/calibration/scaling.py

def __init__(self, num_classes: int):
    """Initializes the VectorScaling module.

    Args:
        num_classes: Number of classes in the classification task.
    """
    super().__init__()
    self.a = nn.Parameter(torch.ones(num_classes))
    self.b = nn.Parameter(torch.zeros(num_classes))

fit

fit(
    logits: Tensor, labels: Tensor, max_iter: int = 50
) -> VectorScaling

Finds optimal scaling vectors 'a' and 'b' using a validation set.

Optimizes the negative log-likelihood via L-BFGS.

Parameters:

Name	Type	Description	Default
`logits`	`Tensor`	Unscaled logits from a hold-out validation set of shape (N, num_classes).	required
`labels`	`Tensor`	Ground truth class indices of shape (N,).	required
`max_iter`	`int`	Maximum number of L-BFGS iterations. Defaults to 50.	`50`

Returns:

Name	Type	Description
`VectorScaling`	`VectorScaling`	The fitted instance itself.

Source code in bensemble/calibration/scaling.py

def fit(
    self, logits: torch.Tensor, labels: torch.Tensor, max_iter: int = 50
) -> "VectorScaling":
    """Finds optimal scaling vectors 'a' and 'b' using a validation set.

    Optimizes the negative log-likelihood via L-BFGS.

    Args:
        logits: Unscaled logits from a hold-out validation set of shape
            (N, num_classes).
        labels: Ground truth class indices of shape (N,).
        max_iter: Maximum number of L-BFGS iterations. Defaults to 50.

    Returns:
        VectorScaling: The fitted instance itself.
    """
    logits = logits.detach()
    optimizer = optim.LBFGS([self.a, self.b], lr=0.01, max_iter=max_iter)

    def eval_loss():
        optimizer.zero_grad()
        scaled_logits = self.forward(logits)
        loss = F.cross_entropy(scaled_logits, labels)
        loss.backward()
        return loss

    optimizer.step(eval_loss)
    return self

forward

forward(logits: Tensor) -> torch.Tensor

Applies learned affine transformation to the input logits.

Parameters:

Name	Type	Description	Default
`logits`	`Tensor`	Raw uncalibrated logits of shape (batch_size, num_classes).	required

Returns:

Type	Description
`Tensor`	torch.Tensor: Calibrated logits of shape (batch_size, num_classes).

Source code in bensemble/calibration/scaling.py

def forward(self, logits: torch.Tensor) -> torch.Tensor:
    """Applies learned affine transformation to the input logits.

    Args:
        logits: Raw uncalibrated logits of shape (batch_size, num_classes).

    Returns:
        torch.Tensor: Calibrated logits of shape (batch_size, num_classes).
    """
    return logits * self.a + self.b