Struct dfdx::nn::MultiHeadAttention

pub struct MultiHeadAttention<const M: usize, const N: usize, const K: usize, const V: usize, const H: usize> {
    pub w_q: Linear<M, K>,
    pub w_k: Linear<N, K>,
    pub w_v: Linear<N, V>,
    pub w_o: Linear<V, M>,
}

Requires Nightly A multi-head attention layer.

Generics

• M The embedding size of token vectors from decoder.
• N The embedding size of token vectors from encoder.
• K The size of the keys in self attention.
• V The size of the values.
• H The number of attention heads.

Examples

MultiHeadAttention<8, 10, 10, 10, 2> is an attention layer with 2 heads and 10 token, key and value dims. TODO: Doctests fail for some reason

Fields

w_q: Linear<M, K>

w_k: Linear<N, K>

w_v: Linear<N, V>

w_o: Linear<V, M>

Trait Implementations

impl<const M: usize, const N: usize, const K: usize, const V: usize, const H: usize> CanUpdateWithGradients for MultiHeadAttention<M, N, K, V, H>

fn update<G: GradientProvider>(
    &mut self,
    grads: &mut G,
    unused: &mut UnusedTensors
)

Updates self given the GradientProvider. When any parameters that are NOT present in G, then this function should add the tensor’s UniqueId to UnusedTensors.

impl<const M: usize, const N: usize, const K: usize, const V: usize, const H: usize> Clone for MultiHeadAttention<M, N, K, V, H>

fn clone(&self) -> MultiHeadAttention<M, N, K, V, H>

Returns a copy of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl<const M: usize, const N: usize, const K: usize, const V: usize, const H: usize> Debug for MultiHeadAttention<M, N, K, V, H>

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl<const M: usize, const N: usize, const K: usize, const V: usize, const H: usize> Default for MultiHeadAttention<M, N, K, V, H>

fn default() -> MultiHeadAttention<M, N, K, V, H>

Returns the “default value” for a type.

impl<const M: usize, const N: usize, const K: usize, const V: usize, const S1: usize, const S2: usize, const H: usize, T: 'static + Tape> Module<(Tensor2D<S1, M, T>, Tensor2D<S2, N, NoneTape>)> for MultiHeadAttention<M, N, K, V, H> where
    Assert<{ _ }>: ConstTrue,
    Assert<{ _ }>: ConstTrue,
    Assert<{ _ }>: ConstTrue,
    Assert<{ _ }>: ConstTrue, 

fn forward(
    &self,
    (input, from_enc): (Tensor2D<S1, M, T>, Tensor2D<S2, N>)
) -> Self::Output

Encoder-Decoder style self attention where one set of tensors is used for values and keys, and another is used for queries

type Output = Tensor2D<S1, M, T>

The type that this unit produces given Input.

fn forward_mut(&mut self, input: Input) -> Self::Output

Pass an Input through the unit and produce Self::Output. Can be implemented for multiple Input types.

impl<const B: usize, const M: usize, const N: usize, const K: usize, const V: usize, const S1: usize, const S2: usize, const H: usize, T: 'static + Tape> Module<(Tensor3D<B, S1, M, T>, Tensor3D<B, S2, N, NoneTape>)> for MultiHeadAttention<M, N, K, V, H> where
    Assert<{ _ }>: ConstTrue,
    Assert<{ _ }>: ConstTrue,
    Assert<{ _ }>: ConstTrue,
    Assert<{ _ }>: ConstTrue, 

fn forward(
    &self,
    (input, from_enc): (Tensor3D<B, S1, M, T>, Tensor3D<B, S2, N>)
) -> Self::Output

Batched Encoder-Decoder style self attention where one set of tensors is used for values and keys, and another is used for queries

type Output = Tensor3D<B, S1, M, T>

The type that this unit produces given Input.

fn forward_mut(&mut self, input: Input) -> Self::Output

Pass an Input through the unit and produce Self::Output. Can be implemented for multiple Input types.

impl<const M: usize, const K: usize, const V: usize, const S: usize, const H: usize, T: 'static + Tape> Module<Tensor2D<S, M, T>> for MultiHeadAttention<M, M, K, V, H> where
    Assert<{ _ }>: ConstTrue,
    Assert<{ _ }>: ConstTrue,
    Assert<{ _ }>: ConstTrue, 

fn forward(&self, input: Tensor2D<S, M, T>) -> Self::Output

Normal self attention (where same tensors are used for keys, queries and values)

type Output = Tensor2D<S, M, T>

The type that this unit produces given Input.

fn forward_mut(&mut self, input: Input) -> Self::Output

Pass an Input through the unit and produce Self::Output. Can be implemented for multiple Input types.

impl<const B: usize, const M: usize, const K: usize, const V: usize, const S: usize, const H: usize, T: 'static + Tape> Module<Tensor3D<B, S, M, T>> for MultiHeadAttention<M, M, K, V, H> where
    Assert<{ _ }>: ConstTrue,
    Assert<{ _ }>: ConstTrue,
    Assert<{ _ }>: ConstTrue, 

fn forward(&self, input: Tensor3D<B, S, M, T>) -> Self::Output

Batched normal self attention (where same tensors are used for keys, queries and values)

type Output = Tensor3D<B, S, M, T>

The type that this unit produces given Input.

fn forward_mut(&mut self, input: Input) -> Self::Output

Pass an Input through the unit and produce Self::Output. Can be implemented for multiple Input types.

impl<const M: usize, const N: usize, const K: usize, const V: usize, const H: usize> ResetParams for MultiHeadAttention<M, N, K, V, H>

fn reset_params<R: Rng>(&mut self, rng: &mut R)

Mutate the unit’s parameters using rand::Rng. Each implementor of this trait decides how the parameters are initialized. In fact, some impls may not even use the rng.