Struct dfdx::nn::MultiHeadAttention

pub struct MultiHeadAttention<const EMBED_DIM: usize, const NUM_HEADS: usize, const K_DIM: usize = EMBED_DIM, const V_DIM: usize = EMBED_DIM> {
    pub w_q: Linear<EMBED_DIM, K_DIM>,
    pub w_k: Linear<EMBED_DIM, K_DIM>,
    pub w_v: Linear<EMBED_DIM, V_DIM>,
    pub w_o: Linear<V_DIM, EMBED_DIM>,
}

Requires Nightly A multi-head attention layer.

Generics:

• EMBED_DIM: The size of query vectors.
• NUM_HEADS The number of heads to split query/key/value into.
• Optional K_DIM: The size of key vectors. Defaults to EMBED_DIM
• Optional V_DIM The size of value vectors. Defaults to EMBED_DIM

Pytorch equivalent: torch.nn.MultiheadAttention(EMBED_DIM, NUM_HEADS, batch_first=True)

Examples

• MultiHeadAttention<8, 2> is an attention layer with 2 heads and 8 token, key and value dims.
• MultiHeadAttention<8, 2, 6, 4> is an attention layer with the key and value dimension different than the embed dimension TODO: Doctests fail for some reason

Fields

w_q: Linear<EMBED_DIM, K_DIM>

w_k: Linear<EMBED_DIM, K_DIM>

w_v: Linear<EMBED_DIM, V_DIM>

w_o: Linear<V_DIM, EMBED_DIM>

Trait Implementations

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

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 EMBED_DIM: usize, const NUM_HEADS: usize, const K_DIM: usize, const V_DIM: usize> Clone for MultiHeadAttention<EMBED_DIM, NUM_HEADS, K_DIM, V_DIM>

fn clone(&self) -> MultiHeadAttention<EMBED_DIM, NUM_HEADS, K_DIM, V_DIM>

Returns a copy of the value.

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

Performs copy-assignment from source.

impl<const EMBED_DIM: usize, const NUM_HEADS: usize, const K_DIM: usize, const V_DIM: usize> Debug for MultiHeadAttention<EMBED_DIM, NUM_HEADS, K_DIM, V_DIM>

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

Formats the value using the given formatter.

impl<const EMBED_DIM: usize, const NUM_HEADS: usize, const K_DIM: usize, const V_DIM: usize> Default for MultiHeadAttention<EMBED_DIM, NUM_HEADS, K_DIM, V_DIM>

fn default() -> MultiHeadAttention<EMBED_DIM, NUM_HEADS, K_DIM, V_DIM>

Returns the “default value” for a type.

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

fn read<R: Read + Seek>(
    &mut self,
    p: &str,
    r: &mut ZipArchive<R>
) -> Result<(), NpzError>

Reads this object from a ZipArchive. r with a base filename of filename_prefix.

fn load<P: AsRef<Path>>(&mut self, path: P) -> Result<(), NpzError>

Loads data from a .npz zip archive at the specified path.

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

fn forward(
    &self,
    (q, k, v): (Tensor2D<S1, M, TAPE>, Tensor2D<S2, M>, Tensor2D<S2, M>)
) -> 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, TAPE>

The type that this unit produces given Input.

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

fn forward(
    &self,
    (q, k, v): (Tensor3D<B, S1, M, TAPE>, Tensor3D<B, S2, M>, Tensor3D<B, S2, M>)
) -> 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, TAPE>

The type that this unit produces given Input.

impl<const M: usize, const H: usize, const K: usize, const V: usize, T> ModuleMut<T> for MultiHeadAttention<M, H, K, V>where
    Self: Module<T>,

type Output = <MultiHeadAttention<M, H, K, V> as Module<T>>::Output

The type that this unit produces given Input.

fn forward_mut(&mut self, t: T) -> Self::Output

Forward Input through the module and produce ModuleMut::Output.

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

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.

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

fn write<W: Write + Seek>(&self, p: &str, w: &mut ZipWriter<W>) -> ZipResult<()>

Write this object into ZipWriter w with a base filename of filename_prefix.

fn save<P: AsRef<Path>>(&self, path: P) -> ZipResult<()>

Save this object into the .npz file determined located at path.