class

ML::NN::MultiHeadAttention

Inherits Reference / Object

Multi-Head Attention Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V

Constructors

new(embed_dim : Int32, num_heads : Int32, dropout : Float32 = 0.0_f32, bias : Bool = true, device : Tensor::Device = Tensor.default_device)
Source

Instance methods

call(query : Autograd::Variable, key : Autograd::Variable, value : Autograd::Variable, attn_mask : Tensor | Nil = nil) : Autograd::Variable
Source
dropout
Source
embed_dim
Source
forward(query : Autograd::Variable, key : Autograd::Variable, value : Autograd::Variable, attn_mask : Tensor | Nil = nil, need_weights : Bool = false) : Autograd::Variable

Forward pass query, key, value: [batch, seq_len, embed_dim] attn_mask: optional [batch, seq_len, seq_len] or [seq_len, seq_len] Returns: [batch, seq_len, embed_dim]

Source
head_dim
Source
k_proj
Source
num_heads
Source
out_proj
Source
parameters

Get all trainable parameters

Source
q_proj

Projections

Source
self_attention(x : Autograd::Variable, attn_mask : Tensor | Nil = nil) : Autograd::Variable

Self-attention convenience method

Source
v_proj
Source