Permutation language-model loss under one factorization order — morie_permutation_lm

Equation (3) of the XLNet paper, negated: $$-\sum_{t} \log p_\theta(x_{z_t} \mid x_{z_{<t}}).$$

Usage

morie_permutation_lm_loss(
  logits,
  targets,
  permutation,
  num_predict = NULL,
  reduction = c("mean", "sum", "none")
)

Arguments

logits: Numeric matrix, T rows by V vocabulary columns.
targets: Integer vector of true token ids in 0:(V-1).
permutation: Factorization order, a permutation of 0:(T-1).
num_predict: Optional count of trailing positions to score.
reduction: One of "mean", "sum" or "none".

Value

A list with loss, token_nll, scored_positions, perplexity, order_invariant, mean_context_length.

Details

With the logits held fixed the full-sequence loss does NOT depend on the permutation, because reordering the terms of a sum leaves the sum alone. The permutation acts only on what the network was allowed to condition on when producing those logits – that is, on the masks – never on this arithmetic. order_invariant records that, and is TRUE by construction whenever the whole sequence is scored.

Equation (5) is partial prediction: score only the trailing num_predict positions of the factorization order, whose contexts are the longest available. The scored subset then depends on the permutation, so the loss does too, and order_invariant goes FALSE.

References

Yang et al (2019) arXiv:1906.08237, eq (3) and (5).

Examples

set.seed(10)
logits <- matrix(rnorm(4 * 5), 4, 5)
morie_permutation_lm_loss(logits, targets = c(0, 3, 1, 4), permutation = c(2, 0, 3, 1))$loss
#> [1] 1.707376