Softmax

The softmax primitive performs softmax along a particular axis on data with arbitrary dimensions. All other axes are treated as independent (batch).

In general form, the operation is defined by the following formulas. The variable names follow the standard Conventions.

Forward

When the specified algorithm is softmax:

\[\dst(\overline{ou}, c, \overline{in}) = \frac {e^{\src(\overline{ou}, c, \overline{in}) - \nu(\overline{ou}, \overline{in})}} { \sum\limits_{ic} e^{\src(\overline{ou}, ic, \overline{in}) - \nu(\overline{ou}, \overline{in})} }.\]

When the specified algorithm is logsoftmax, the following numerically stable formula is used:

\[\dst(\overline{ou}, c, \overline{in}) = \ln\left({\frac { e^{\src(\overline{ou}, c, \overline{in}) - \nu(\overline{ou}, \overline{in})} } { \sum\limits_{ic} e^{\src(\overline{ou}, ic, \overline{in}) - \nu(\overline{ou}, \overline{in})} }}\right) = \left(\src(\overline{ou}, c, \overline{in}) - \nu(\overline{ou}, \overline{in})\right) - \ln\left( \sum\limits_{ic} e^{\src(\overline{ou}, ic, \overline{in}) - \nu(\overline{ou}, \overline{in})} \right)\]

where

• \(c\) axis over which the softmax computation is computed on,

• \(\overline{ou}\) is the outermost index (to the left of softmax axis),

• \(\overline{in}\) is the innermost index (to the right of softmax axis), and

• \(\nu\) is used to produce more accurate results and defined as:

\[\nu(\overline{ou}, \overline{in}) = \max\limits_{ic} \src(\overline{ou}, ic, \overline{in})\]

Difference Between Forward Training and Forward Inference

There is no difference between the forward_training and forward_inference propagation kinds.

Backward

The backward propagation computes \(\diffsrc(ou, c, in)\), based on \(\diffdst(ou, c, in)\) and \(\dst(ou, c, in)\).

Execution Arguments

When executed, the inputs and outputs should be mapped to an execution argument index as specified by the following table.

Primitive input/output	Execution argument index
\(\src\)	DNNL_ARG_SRC
\(\dst\)	DNNL_ARG_DST
\(\diffsrc\)	DNNL_ARG_DIFF_SRC
\(\diffdst\)	DNNL_ARG_DIFF_DST

Operation Details

1. Both forward and backward propagation support in-place operations, meaning that src can be used as input and output for forward propagation, and diff_dst can be used as input and output for backward propagation. In case of in-place operation, the original data will be overwritten.

Post-ops and Attributes

The softmax primitive does not have to support any post-ops or attributes.

Data Types Support

The softmax primitive supports the following combinations of data types.

Note

Here we abbreviate data types names for readability. For example, dnnl::memory::data_type::f32 is abbreviated to f32.

Propagation	Source / Destination
forward / backward	bf16, f32
forward	f16

Data Representation

Source, Destination, and Their Gradients

The softmax primitive works with arbitrary data tensors. There is no special meaning associated with any logical dimensions. However, the softmax axis is typically referred to as channels (hence in formulas we use \(c\)).

API

struct softmax_forward : public dnnl::primitive

Softmax forward propagation primitive.

Public Functions

softmax_forward()

Default constructor. Produces an empty object.

softmax_forward(const primitive_desc &pd)

Constructs a softmax forward propagation primitive.

Parameters:

pd – Primitive descriptor for a softmax forward propagation primitive.

struct primitive_desc : public dnnl::primitive_desc

Primitive descriptor for a softmax forward propagation primitive.

Public Functions

primitive_desc() = default

Default constructor. Produces an empty object.

primitive_desc(const engine &aengine, prop_kind aprop_kind, algorithm aalgorithm, const memory::desc &src_desc, const memory::desc &dst_desc, int axis, const primitive_attr &attr = default_attr(), bool allow_empty = false)

Constructs a primitive descriptor for a softmax forward propagation primitive.

Parameters:

• aengine – Engine to use.

• aprop_kind – Propagation kind. Possible values are dnnl::prop_kind::forward_training, and dnnl::prop_kind::forward_inference.

• aalgorithm – Softmax algorithm kind: either dnnl::algorithm::softmax_accurate, or dnnl::algorithm::softmax_log.

• src_desc – Source memory descriptor.

• dst_desc – Destination memory descriptor.

• axis – Axis over which softmax is computed.

• attr – Primitive attributes to use. Attributes are optional and default to empty attributes.

• allow_empty – A flag signifying whether construction is allowed to fail without throwing an exception. In this case an empty object will be produced. This flag is optional and defaults to false.

memory::desc src_desc() const

Returns a source memory descriptor.

Returns:

Source memory descriptor.

Returns:

A zero memory descriptor if the primitive does not have a source parameter.

memory::desc dst_desc() const

Returns a destination memory descriptor.

Returns:

Destination memory descriptor.

Returns:

A zero memory descriptor if the primitive does not have a destination parameter.

dnnl::algorithm get_algorithm() const

Returns an algorithm kind.

Returns:

An algorithm kind.

Returns:

dnnl::algorithm::undef if the primitive does not have an algorithm parameter.

dnnl::prop_kind get_prop_kind() const

Returns a propagation kind.

Returns:

A propagation kind.

Returns:

dnnl::prop_kind::undef if the primitive does not have a propagation parameter.

int get_axis() const

Returns an axis.

Returns:

An axis.

Returns:

A negative number if the primitive does not have an axis parameter.

struct softmax_backward : public dnnl::primitive

Softmax backward propagation primitive.

Public Functions

softmax_backward()

Default constructor. Produces an empty object.

softmax_backward(const primitive_desc &pd)

Constructs a softmax backward propagation primitive.

Parameters:

pd – Primitive descriptor for a softmax backward propagation primitive.

struct primitive_desc : public dnnl::primitive_desc

Primitive descriptor for a softmax backward propagation primitive.

Public Functions

primitive_desc() = default

Default constructor. Produces an empty object.

primitive_desc(const engine &aengine, algorithm aalgorithm, const memory::desc &diff_src_desc, const memory::desc &diff_dst_desc, const memory::desc &dst_desc, int axis, const softmax_forward::primitive_desc &hint_fwd_pd, const primitive_attr &attr = default_attr(), bool allow_empty = false)

Constructs a primitive descriptor for a softmax backward propagation primitive.

Parameters:

• aengine – Engine to use.

• aalgorithm – Softmax algorithm kind: either dnnl::algorithm::softmax_accurate, or dnnl::algorithm::softmax_log.

• diff_src_desc – Diff source memory descriptor.

• diff_dst_desc – Diff destination memory descriptor.

• dst_desc – Destination memory descriptor.

• axis – Axis over which softmax is computed.

• hint_fwd_pd – Primitive descriptor for a softmax forward propagation primitive. It is used as a hint for deciding which memory format to use.

• attr – Primitive attributes to use. Attributes are optional and default to empty attributes.

• allow_empty – A flag signifying whether construction is allowed to fail without throwing an exception. In this case an empty object will be produced. This flag is optional and defaults to false.

memory::desc dst_desc() const

Returns a destination memory descriptor.

Returns:

Destination memory descriptor.

Returns:

A zero memory descriptor if the primitive does not have a destination parameter.

memory::desc diff_src_desc() const

Returns a diff source memory descriptor.

Returns:

Diff source memory descriptor.

Returns:

A zero memory descriptor if the primitive does not have a diff source memory with.

memory::desc diff_dst_desc() const

Returns a destination memory descriptor.

Returns:

Destination memory descriptor.

Returns:

A zero memory descriptor if the primitive does not have a destination parameter.

dnnl::algorithm get_algorithm() const

Returns an algorithm kind.

Returns:

An algorithm kind.

Returns:

dnnl::algorithm::undef if the primitive does not have an algorithm parameter.

dnnl::prop_kind get_prop_kind() const

Returns a propagation kind.

Returns:

A propagation kind.

Returns:

dnnl::prop_kind::undef if the primitive does not have a propagation parameter.

int get_axis() const

Returns an axis.

Returns:

An axis.

Returns:

A negative number if the primitive does not have an axis parameter.