topos 0.12.0

An autodiff compiler stack in Rust.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67

//! Loss functions: named formulas that compose recorded operations into a
//! scalar training objective.
//!
//! A loss is a composition, not a primitive: it records existing graph
//! operations, so its gradient falls out of the chain rule with no
//! dedicated backward rule. A formula earns a fused `Op` variant
//! only where composition cannot express it — [`Value::log_softmax`] is
//! that fused core here, because its stability demands a max shift no
//! recorded composition can perform — while everything around the core
//! stays plain, readable composition.

use crate::{Element, Value};

/// Records the cross-entropy loss of `logits` against `targets` on their
/// network and returns the rank-0 loss value.
///
/// It composes the mean negative log-likelihood in the expanded form
/// `((targets.sum_along(1) * logsumexp(logits)).sum() -
/// (targets * logits).sum()) / targets.sum()`. The expansion is exact
/// mathematics — each row's `-t . (x - lse)` distributed — and it is the
/// stable spelling: the fused [`Value::logsumexp`] is finite for every
/// finite logit, and no term ever multiplies a zero target by an
/// infinite log-probability (the `0 * -inf = NaN` a
/// `targets * log_softmax` product produces once finite logits differ
/// by more than the representable range). The normalizer is the
/// targets' total mass: the batch size for one-hot targets — the
/// standard mean reduction — while soft or weighted targets normalize
/// by their own weight.
///
/// # Parameters
/// - `logits`: The unnormalized class scores, rank 2 `[batch, classes]`.
/// - `targets`: The target distribution per sample, shaped like `logits`.
///   Feed a one-hot [`Tensor::selection`](crate::Tensor::selection) as a
///   per-run input so one recorded graph serves any batch of labels.
///   Target weights must be finite and nonnegative with a strictly
///   positive total mass; outside that domain the "mean negative
///   log-likelihood" has no interpretation and the result follows IEEE
///   arithmetic (an all-zero target tensor, for instance, divides zero
///   by zero into `NaN`).
///
/// # Panics
/// Panics if the values belong to different networks, `logits` is not
/// rank 2, or the shapes differ.
pub fn cross_entropy<'tape, E: Element>(
    logits: Value<'tape, E>,
    targets: Value<'tape, E>,
) -> Value<'tape, E> {
    let logits_shape = logits.shape();
    assert_eq!(
        logits_shape.rank(),
        2,
        "cross-entropy logits must be rank 2 [batch, classes], got {logits_shape}"
    );
    let targets_shape = targets.shape();
    assert_eq!(
        targets_shape, logits_shape,
        "cross-entropy targets {targets_shape} must be shaped like the logits {logits_shape}"
    );

    let normalizers = (targets.sum_along(1) * logits.logsumexp(1)).sum();
    let alignment = (targets * logits).sum();
    (normalizers - alignment) / targets.sum()
}

#[cfg(test)]
#[path = "tests/loss_tests.rs"]
mod tests;