Bias Initializer for unbiased classification

[LucaCappelletti94]

Hi! I have been using the following initialization for the biases output layer of my classifiers, and it seems to work very well. In practice, I see it converge much faster (for the latest task I am using this for, it is taking 50 epochs on average instead of about 300, averaged across 10 holdouts). If you deem it useful, I can make a pull request to add it in.

Please do let me know whether I should do so.

"""Bias inizialer for unbalanced prediction tasks."""

from keras.api.initializers import Initializer
from keras.api.backend import epsilon
import numpy as np


class LogitBiasInitializer(Initializer):
    """Initializer that sets the bias based on the prior distribution of the training labels."""

    def __init__(self, labels: np.ndarray):
        """
        Initialize the bias based on the prior distribution of the training labels.
        The labels should be one-hot encoded (multilabel classification).

        :param labels: numpy array of one-hot encoded labels
        """
        # We check that labels is a numpy array
        if not isinstance(labels, np.ndarray):
            raise ValueError("Labels must be a numpy array.")

        # We check that labels is a 2D array
        if len(labels.shape) != 2:
            raise ValueError("Labels must be a 2D array.")

        # We check that labels is a binary array
        if not np.array_equal(np.unique(labels), np.array([0, 1])):
            raise ValueError("Labels must be a binary array.")

        # Calculate the prior probability for each label
        label_sums = np.sum(labels, axis=0)
        total_samples = labels.shape[0]
        self.prior_probs = label_sums / total_samples

    def __call__(self, shape, dtype=None):
        """
        Compute the bias values based on the log-odds of each class.

        :param shape: shape of the bias tensor (should match the number of output classes)
        :param dtype: the dtype of the tensor (optional)
        :return: initialized bias tensor
        """
        # Prevent divide by zero or log of zero issues
        log_odds = np.log(epsilon() + self.prior_probs / (1 - self.prior_probs + epsilon()))

        # In some unfortunate cases, such as when no sample has a certain label,
        # the log-odds will have as value -inf. We replace these values with 0.
        log_odds[np.isneginf(log_odds)] = 0

        # We check that there is no NaN in the log-odds
        if np.isnan(log_odds).any():
            raise ValueError("NaN values found in the log-odds.")

        # Ensure the shape matches the expected number of output classes
        if len(log_odds) != shape[0]:
            raise ValueError(
                f"Shape mismatch: expected {shape[0]} biases, but found {len(log_odds)} log-odds."
            )

        # Return the log-odds as the bias initialization
        return log_odds.astype(dtype)

    def get_config(self):
        """
        Return the configuration of the initializer for serialization.
        """
        return {"prior_probs": self.prior_probs.tolist()}