![]() |
Gaussian process-based pooling head for sentence classification.
Inherits From: ClassificationHead
tfm.nlp.layers.GaussianProcessClassificationHead(
inner_dim,
num_classes,
cls_token_idx=0,
activation='tanh',
dropout_rate=0.0,
initializer='glorot_uniform',
use_spec_norm=True,
use_gp_layer=True,
temperature=None,
**kwargs
)
This class implements a classifier head for BERT encoder that is based on the spectral-normalized neural Gaussian process (SNGP) [1]. SNGP is a simple method to improve a neural network's uncertainty quantification ability without sacrificing accuracy or lantency. It applies spectral normalization to the hidden pooler layer, and then replaces the dense output layer with a Gaussian process.
[1]: Jeremiah Liu et al. Simple and Principled Uncertainty Estimation with Deterministic Deep Learning via Distance Awareness. In Neural Information Processing Systems, 2020. https://arxiv.org/abs/2006.10108
Attributes | |
---|---|
checkpoint_items
|
Methods
call
call(
features, training=False, return_covmat=False
)
Returns model output.
Dring training, the model returns raw logits. During evaluation, the model returns uncertainty adjusted logits, and (optionally) the covariance matrix.
Arguments | |
---|---|
features
|
A tensor of input features, shape (batch_size, feature_dim). |
training
|
Whether the model is in training mode. |
return_covmat
|
Whether the model should also return covariance matrix if
use_gp_layer=True . During training, it is recommended to set
return_covmat=False to be compatible with the standard Keras pipelines
(e.g., model.fit() ).
|
Returns | |
---|---|
logits
|
Uncertainty-adjusted predictive logits, shape (batch_size, num_classes). |
covmat
|
(Optional) Covariance matrix, shape (batch_size, batch_size). Returned only when return_covmat=True. |
reset_covariance_matrix
reset_covariance_matrix()
Resets covariance matrix of the Gaussian process layer.