Reward Model

Reward Model and Pairwise Loss function

Reward Model

RewardModel

 RewardModel (model_name:str, dropout:float=0.1, device:str='cuda')

Reward model.

	Type	Default	Details
model_name	str		`transformers`’s model name
dropout	float	0.1
device	str	cuda

source

RewardModel.forward

 RewardModel.forward (input_ids:typing.Annotated[torch.Tensor,{'__torchtyp
                      ing__':True,'details':('batch_size','seq_len',),'cls
                      _name':'TensorType'}], attention_mask:typing.Annotat
                      ed[torch.Tensor,{'__torchtyping__':True,'details':('
                      batch_size','seq_len',),'cls_name':'TensorType'}]=No
                      ne)

Calculate reward for each item in a batch.

	Type	Default	Details
input_ids	typing.Annotated[torch.Tensor, {‘torchtyping’: True, ‘details’: (‘batch_size’, ‘seq_len’,), ‘cls_name’: ‘TensorType’}]
attention_mask	typing.Annotated[torch.Tensor, {‘torchtyping’: True, ‘details’: (‘batch_size’, ‘seq_len’,), ‘cls_name’: ‘TensorType’}]	None
Returns	typing.Annotated[torch.Tensor, {‘torchtyping’: True, ‘details’: (‘batch_size’, 1,), ‘cls_name’: ‘TensorType’}]		A reward scalar for each item in a batch

Pairwise Loss

\(\operatorname{loss}(\theta)=-\frac{1}{\left(\begin{array}{c} K \\ 2 \end{array}\right)} E_{\left(x, y_w, y_l\right) \sim D}\left[\log \left(\sigma\left(r_\theta\left(x, y_w\right)-r_\theta\left(x, y_l\right)\right)\right)\right]\)

source

PairwiseLoss

 PairwiseLoss (*args, **kwargs)

Pairwise loss function.

source

PairwiseLoss.forward

 PairwiseLoss.forward (chosen_rewards:typing.Annotated[torch.Tensor,{'__to
                       rchtyping__':True,'details':('batch_size',1,),'cls_
                       name':'TensorType'}], rejected_rewards:typing.Annot
                       ated[torch.Tensor,{'__torchtyping__':True,'details'
                       :('batch_size',1,),'cls_name':'TensorType'}])

Compute the loss value.

	Type	Details
chosen_rewards	typing.Annotated[torch.Tensor, {‘torchtyping’: True, ‘details’: (‘batch_size’, 1,), ‘cls_name’: ‘TensorType’}]	The reward of the chosen prompt
rejected_rewards	typing.Annotated[torch.Tensor, {‘torchtyping’: True, ‘details’: (‘batch_size’, 1,), ‘cls_name’: ‘TensorType’}]	The reward of the rejected prompt
Returns	typing.Annotated[torch.Tensor, {‘torchtyping’: True, ‘details’: (1,), ‘cls_name’: ‘TensorType’}]	A scalar loss