apaam/mlpack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update multihead attention layer

Browse Source
This commit is contained in:
Mrityunjay Tripathi
2021-01-13 09:19:09 +05:30
parent dd6b70c2f2
commit 09e2995346
4 changed files with 276 additions and 275 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/mlpack/methods/ann/layer/layer_types.hpp
+1 -1
View File
@@ -33,7 +33,7 @@
#include <mlpack/methods/ann/layer/linear3d.hpp>
#include <mlpack/methods/ann/layer/log_softmax.hpp>
#include <mlpack/methods/ann/layer/lookup.hpp>
// #include <mlpack/methods/ann/layer/multihead_attention.hpp>
#include <mlpack/methods/ann/layer/multihead_attention.hpp>
// #include <mlpack/methods/ann/layer/multiply_constant.hpp>
// #include <mlpack/methods/ann/layer/max_pooling.hpp>
// #include <mlpack/methods/ann/layer/mean_pooling.hpp>
src/mlpack/methods/ann/layer/multihead_attention.hpp
+56 -50
View File
@@ -48,24 +48,24 @@ namespace ann /** Artificial Neural Network. */ {
* of shape `(embedDim * tgtSeqLen, batchSize)`. The embeddings are stored
* consequently.
*
* @tparam InputDataType Type of the input data (arma::colvec, arma::mat,
* @tparam InputType Type of the input data (arma::colvec, arma::mat,
* arma::sp_mat or arma::cube).
* @tparam OutputDataType Type of the output data (arma::colvec, arma::mat,
* @tparam OutputType Type of the output data (arma::colvec, arma::mat,
* arma::sp_mat or arma::cube).
* @tparam RegularizerType Type of the regularizer to be used.
*/
template <
typename InputDataType = arma::mat,
typename OutputDataType = arma::mat,
typename InputType = arma::mat,
typename OutputType = arma::mat,
typename RegularizerType = NoRegularizer
>
class MultiheadAttention
class MultiheadAttentionType : public Layer<InputType, OutputType>
{
public:
/**
* Default constructor.
*/
MultiheadAttention();
MultiheadAttentionType();
/**
* Create the MultiheadAttention object using the specified modules.
@@ -74,11 +74,15 @@ class MultiheadAttention
* @param srcSeqLen Source sequence length.
* @param embedDim Total dimension of the model.
* @param numHeads Number of parallel attention heads.
* @param attnMask Two dimensional Attention Mask.
* @param keyPaddingMask Key Padding Mask.
*/
MultiheadAttention(const size_t tgtSeqLen,
const size_t srcSeqLen,
const size_t embedDim,
const size_t numHeads);
MultiheadAttentionType(const size_t tgtSeqLen,
const size_t srcSeqLen,
const size_t embedDim,
const size_t numHeads,
const InputType& attnmask = InputType(),
const InputType& keyPaddingMask = InputType());
/**
* Reset the layer parameters.
@@ -92,8 +96,7 @@ class MultiheadAttention
* @param input The query matrix.
* @param output Resulting output activation.
*/
template<typename eT>
void Forward(const arma::Mat<eT>& input, arma::Mat<eT>& output);
void Forward(const InputType& input, OutputType& output);
/**
* Ordinary feed backward pass of a neural network, calculating the function
@@ -103,10 +106,9 @@ class MultiheadAttention
* @param gy The backpropagated error.
* @param g The calculated gradient.
*/
template<typename eT>
void Backward(const arma::Mat<eT>& /* input */,
const arma::Mat<eT>& gy,
arma::Mat<eT>& g);
void Backward(const InputType& /* input */,
const OutputType& gy,
OutputType& g);
/**
* Calculate the gradient using the output delta and the input activation.
@@ -115,10 +117,9 @@ class MultiheadAttention
* @param error The calculated error.
* @param gradient The calculated gradient.
*/
template<typename eT>
void Gradient(const arma::Mat<eT>& input,
const arma::Mat<eT>& error,
arma::Mat<eT>& gradient);
void Gradient(const InputType& input,
const OutputType& error,
OutputType& gradient);
/**
* Serialize the layer.
@@ -147,46 +148,46 @@ class MultiheadAttention
size_t& NumHeads() { return numHeads; }
//! Get the two dimensional Attention Mask.
OutputDataType const& AttentionMask() const { return attnMask; }
OutputType const& AttentionMask() const { return attnMask; }
//! Modify the two dimensional Attention Mask.
OutputDataType& AttentionMask() { return attnMask; }
OutputType& AttentionMask() { return attnMask; }
//! Get Key Padding Mask.
OutputDataType const& KeyPaddingMask() const { return keyPaddingMask; }
OutputType const& KeyPaddingMask() const { return keyPaddingMask; }
//! Modify the Key Padding Mask.
OutputDataType& KeyPaddingMask() { return keyPaddingMask; }
OutputType& KeyPaddingMask() { return keyPaddingMask; }
//! Get the output parameter.
OutputDataType const& OutputParameter() const { return outputParameter; }
OutputType const& OutputParameter() const { return outputParameter; }
//! Modify the output parameter.
OutputDataType& OutputParameter() { return outputParameter; }
OutputType& OutputParameter() { return outputParameter; }
//! Get the delta.
OutputDataType const& Delta() const { return delta; }
OutputType const& Delta() const { return delta; }
//! Modify the delta.
OutputDataType& Delta() { return delta; }
OutputType& Delta() { return delta; }
//! Get the gradient.
OutputDataType const& Gradient() const { return grad; }
OutputType const& Gradient() const { return grad; }
//! Modify the gradient.
OutputDataType& Gradient() { return grad; }
OutputType& Gradient() { return grad; }
//! Get the parameters.
OutputDataType const& Parameters() const { return weights; }
OutputType const& Parameters() const { return weights; }
//! Modify the parameters.
OutputDataType& Parameters() { return weights; }
OutputType& Parameters() { return weights; }
private:
//! Element Type of the input.
typedef typename OutputDataType::elem_type ElemType;
typedef typename OutputType::elem_type ElemType;
//! Target sequence length.
size_t tgtSeqLen;
//! Source sequence lenght.
//! Source sequence length.
size_t srcSeqLen;
//! Locally-stored module output size.
//! Locally-stored dimensionality of each embedding vector.
size_t embedDim;
//! Locally-stored number of parallel attention heads.
@@ -196,37 +197,37 @@ class MultiheadAttention
size_t headDim;
//! Two dimensional Attention Mask of shape (tgtSeqLen, srcSeqLen).
OutputDataType attnMask;
OutputType attnMask;
//! Key Padding Mask.
OutputDataType keyPaddingMask;
OutputType keyPaddingMask;
//! Locally-stored weight matrix associated with query.
OutputDataType queryWt;
OutputType queryWt;
//! Locally-stored weight matrix associated with key.
OutputDataType keyWt;
OutputType keyWt;
//! Locally-stored weight matrix associated with value.
OutputDataType valueWt;
OutputType valueWt;
//! Locally-stored weight matrix associated with attnWt.
OutputDataType outWt;
OutputType outWt;
//! Locally-stored bias associated with query.
OutputDataType qBias;
OutputType qBias;
//! Locally-stored bias associated with key.
OutputDataType kBias;
OutputType kBias;
//! Locall-stored bias associated with value.
OutputDataType vBias;
OutputType vBias;
//! Locally-stored bias associated with attnWt.
OutputDataType outBias;
OutputType outBias;
//! Locally-stored weights parameter.
OutputDataType weights;
OutputType weights;
//! Locally-stored projected query matrix over linear layer.
arma::Cube<ElemType> qProj;
@@ -244,20 +245,25 @@ class MultiheadAttention
arma::Cube<ElemType> attnOut;
//! Softmax layer to represent the probabilities of next sequence.
Softmax<InputDataType, OutputDataType> softmax;
Softmax softmax;
//! Locally-stored delta object.
OutputDataType delta;
OutputType delta;
//! Locally-stored gradient.
OutputDataType grad;
OutputType grad;
//! Locally-stored output parameter.
OutputDataType outputParameter;
OutputType outputParameter;
//! Locally-stored regularizer object.
RegularizerType regularizer;
}; // class MultiheadAttention
// Standard MultiheadAttention layer using no regularization.
typedef MultiheadAttentionType<arma::mat, arma::mat, NoRegularizer>
MultiheadAttention;
} // namespace ann
} // namespace mlpack
src/mlpack/methods/ann/layer/multihead_attention_impl.hpp
+51 -58
View File
@@ -21,31 +21,35 @@
namespace mlpack {
namespace ann /** Artificial Neural Network. */ {
template <typename InputDataType, typename OutputDataType,
typename RegularizerType>
MultiheadAttention<InputDataType, OutputDataType, RegularizerType>::
MultiheadAttention() :
template <typename InputType, typename OutputType, typename RegularizerType>
MultiheadAttentionType<InputType, OutputType, RegularizerType>::
MultiheadAttentionType() :
tgtSeqLen(0),
srcSeqLen(0),
embedDim(0),
numHeads(0),
headDim(0)
headDim(0),
attnMask(InputType()),
keyPaddingMask(InputType())
{
// Nothing to do here.
}
template <typename InputDataType, typename OutputDataType,
typename RegularizerType>
MultiheadAttention<InputDataType, OutputDataType, RegularizerType>::
MultiheadAttention(
template <typename InputType, typename OutputType, typename RegularizerType>
MultiheadAttentionType<InputType, OutputType, RegularizerType>::
MultiheadAttentionType(
const size_t tgtSeqLen,
const size_t srcSeqLen,
const size_t embedDim,
const size_t numHeads) :
const size_t numHeads,
const InputType& attnMask,
const InputType& keyPaddingMask) :
tgtSeqLen(tgtSeqLen),
srcSeqLen(srcSeqLen),
embedDim(embedDim),
numHeads(numHeads)
numHeads(numHeads),
attnMask(attnMask),
keyPaddingMask(keyPaddingMask)
{
if (embedDim % numHeads != 0)
{
@@ -57,38 +61,33 @@ MultiheadAttention(
weights.set_size(4 * (embedDim + 1) * embedDim, 1);
}
template <typename InputDataType, typename OutputDataType,
typename RegularizerType>
void MultiheadAttention<InputDataType, OutputDataType, RegularizerType>::
template <typename InputType, typename OutputType, typename RegularizerType>
void MultiheadAttentionType<InputType, OutputType, RegularizerType>::
Reset()
{
typedef typename arma::Mat<typename OutputDataType::elem_type> MatType;
queryWt = MatType(weights.memptr(), embedDim, embedDim, false, false);
keyWt = MatType(weights.memptr() + embedDim * embedDim,
queryWt = OutputType(weights.memptr(), embedDim, embedDim, false, false);
keyWt = OutputType(weights.memptr() + embedDim * embedDim,
embedDim, embedDim, false, false);
valueWt = MatType(weights.memptr() + 2 * embedDim * embedDim,
valueWt = OutputType(weights.memptr() + 2 * embedDim * embedDim,
embedDim, embedDim, false, false);
outWt = MatType(weights.memptr() + 3 * embedDim * embedDim,
outWt = OutputType(weights.memptr() + 3 * embedDim * embedDim,
embedDim, embedDim, false, false);
qBias = MatType(weights.memptr()
qBias = OutputType(weights.memptr()
+ 4 * embedDim * embedDim, embedDim, 1, false, false);
kBias = MatType(weights.memptr()
kBias = OutputType(weights.memptr()
+ (4 * embedDim + 1) * embedDim, embedDim, 1, false, false);
vBias = MatType(weights.memptr()
vBias = OutputType(weights.memptr()
+ (4 * embedDim + 2) * embedDim, embedDim, 1, false, false);
outBias = MatType(weights.memptr()
outBias = OutputType(weights.memptr()
+ (4 * embedDim + 3) * embedDim, 1, embedDim, false, false);
}
template <typename InputDataType, typename OutputDataType,
typename RegularizerType>
template <typename eT>
void MultiheadAttention<InputDataType, OutputDataType, RegularizerType>::
Forward(const arma::Mat<eT>& input, arma::Mat<eT>& output)
template <typename InputType, typename OutputType, typename RegularizerType>
void MultiheadAttentionType<InputType, OutputType, RegularizerType>::
Forward(const InputType& input, OutputType& output)
{
typedef typename arma::Cube<eT> CubeType;
typedef typename arma::Cube<typename InputType::elem_type> CubeType;
if (input.n_rows != embedDim * (tgtSeqLen + 2 * srcSeqLen))
{
@@ -104,12 +103,12 @@ Forward(const arma::Mat<eT>& input, arma::Mat<eT>& output)
// The shape of q : (embedDim, tgtSeqLen, batchSize).
// The shape of k : (embedDim, srcSeqLen, batchSize).
// The shape of v : (embedDim, srcSeqLen, batchSize).
const CubeType q(const_cast<arma::Mat<eT>&>(input).memptr(),
const CubeType q(const_cast<InputType&>(input).memptr(),
embedDim, tgtSeqLen, batchSize, false, false);
const CubeType k(const_cast<arma::Mat<eT>&>(input).memptr() +
const CubeType k(const_cast<InputType&>(input).memptr() +
embedDim * tgtSeqLen * batchSize,
embedDim, srcSeqLen, batchSize, false, false);
const CubeType v(const_cast<arma::Mat<eT>&>(input).memptr() +
const CubeType v(const_cast<InputType&>(input).memptr() +
embedDim * (tgtSeqLen + srcSeqLen) * batchSize,
embedDim, srcSeqLen, batchSize, false, false);
@@ -188,15 +187,13 @@ Forward(const arma::Mat<eT>& input, arma::Mat<eT>& output)
}
}
template <typename InputDataType, typename OutputDataType,
typename RegularizerType>
template <typename eT>
void MultiheadAttention<InputDataType, OutputDataType, RegularizerType>::
Backward(const arma::Mat<eT>& /* input */,
const arma::Mat<eT>& gy,
arma::Mat<eT>& g)
template <typename InputType, typename OutputType, typename RegularizerType>
void MultiheadAttentionType<InputType, OutputType, RegularizerType>::
Backward(const InputType& /* input */,
const OutputType& gy,
OutputType& g)
{
typedef typename arma::Cube<eT> CubeType;
typedef typename arma::Cube<typename OutputType::elem_type> CubeType;
if (gy.n_rows != tgtSeqLen * embedDim)
{
@@ -210,7 +207,7 @@ Backward(const arma::Mat<eT>& /* input */,
// The shape of gyTemp : (tgtSeqLen, embedDim, batchSize).
// We need not split it into n heads now because this is the part when
// output were concatenated from n heads.
CubeType gyTemp(const_cast<arma::Mat<eT>&>(gy).memptr(), embedDim,
CubeType gyTemp(const_cast<OutputType&>(gy).memptr(), embedDim,
tgtSeqLen, batchSize, true, false);
// The shape of gyTemp : (embedDim, tgtSeqLen, batchSize).
@@ -280,16 +277,13 @@ Backward(const arma::Mat<eT>& /* input */,
}
}
template <typename InputDataType, typename OutputDataType,
typename RegularizerType>
template <typename eT>
void MultiheadAttention<InputDataType, OutputDataType, RegularizerType>::
Gradient(const arma::Mat<eT>& input,
const arma::Mat<eT>& error,
arma::Mat<eT>& gradient)
template <typename InputType, typename OutputType, typename RegularizerType>
void MultiheadAttentionType<InputType, OutputType, RegularizerType>::
Gradient(const InputType& input,
const OutputType& error,
OutputType& gradient)
{
typedef typename arma::Cube<eT> CubeType;
typedef typename arma::Mat<eT> MatType;
typedef typename arma::Cube<typename InputType::elem_type> CubeType;
if (input.n_rows != embedDim * (tgtSeqLen + 2 * srcSeqLen))
{
@@ -307,16 +301,16 @@ Gradient(const arma::Mat<eT>& input,
// The shape of gradient : (4 * embedDim * embedDim + 4 * embedDim, 1).
gradient.set_size(arma::size(weights));
const CubeType q(const_cast<MatType&>(input).memptr(),
const CubeType q(const_cast<InputType&>(input).memptr(),
embedDim, tgtSeqLen, batchSize, false, false);
const CubeType k(const_cast<MatType&>(input).memptr() + q.n_elem,
const CubeType k(const_cast<InputType&>(input).memptr() + q.n_elem,
embedDim, srcSeqLen, batchSize, false, false);
const CubeType v(const_cast<MatType&>(input).memptr() + q.n_elem + k.n_elem,
const CubeType v(const_cast<InputType&>(input).memptr() + q.n_elem + k.n_elem,
embedDim, srcSeqLen, batchSize, false, false);
// Reshape the propagated error into a cube.
// The shape of errorTemp : (embedDim, tgtSeqLen, batchSize).
CubeType errorTemp(const_cast<arma::Mat<eT>&>(error).memptr(), embedDim,
CubeType errorTemp(const_cast<OutputType&>(error).memptr(), embedDim,
tgtSeqLen, batchSize, true, false);
// Gradient wrt. outBias, i.e. dL/d(outBias).
@@ -430,10 +424,9 @@ Gradient(const arma::Mat<eT>& input,
regularizer.Evaluate(weights, gradient);
}
template <typename InputDataType, typename OutputDataType,
typename RegularizerType>
template <typename InputType, typename OutputType, typename RegularizerType>
template <typename Archive>
void MultiheadAttention<InputDataType, OutputDataType, RegularizerType>::
void MultiheadAttentionType<InputType, OutputType, RegularizerType>::
serialize(Archive& ar, const uint32_t /* version */)
{
ar(CEREAL_NVP(tgtSeqLen));
src/mlpack/tests/ann_layer_test.cpp
+168 -166
View File
@@ -4650,199 +4650,201 @@ TEST_CASE("AdaptiveMeanPoolingTestCase", "[ANNLayerTest]")
// }
// }
// /**
// * Simple Multihead Attention test.
// */
// TEST_CASE("SimpleMultiheadAttentionTest", "[ANNLayerTest]")
// {
// size_t tLen = 5;
// size_t sLen = tLen;
// size_t embedDim = 4;
// size_t numHeads = 2;
// size_t bsz = 3;
/**
* Simple Multihead Attention test.
*/
TEST_CASE("SimpleMultiheadAttentionTest", "[ANNLayerTest]")
{
size_t tLen = 5;
size_t sLen = tLen;
size_t embedDim = 4;
size_t numHeads = 2;
size_t bsz = 3;
// arma::mat query = 0.1 * arma::randu(embedDim * tLen, bsz);
// arma::mat output;
arma::mat query = 0.1 * arma::randu(embedDim * tLen, bsz);
arma::mat output;
// arma::mat attnMask = arma::zeros(tLen, sLen);
// for (size_t i = 0; i < tLen; ++i)
// {
// for (size_t j = 0; j < sLen; ++j)
// {
// if (i < j)
// attnMask(i, j) = std::numeric_limits<double>::lowest();
// }
// }
arma::mat attnMask = arma::zeros(tLen, sLen);
for (size_t i = 0; i < tLen; ++i)
{
for (size_t j = 0; j < sLen; ++j)
{
if (i < j)
attnMask(i, j) = std::numeric_limits<double>::lowest();
}
}
// arma::mat keyPaddingMask = arma::zeros(1, sLen);
// keyPaddingMask(sLen - 1) = std::numeric_limits<double>::lowest();
arma::mat keyPaddingMask = arma::zeros(1, sLen);
keyPaddingMask(sLen - 1) = std::numeric_limits<double>::lowest();
// MultiheadAttention<> module(tLen, sLen, embedDim, numHeads);
// module.AttentionMask() = attnMask;
// module.KeyPaddingMask() = keyPaddingMask;
// module.Reset();
// module.Parameters().randu();
MultiheadAttention module(tLen, sLen, embedDim, numHeads);
module.AttentionMask() = attnMask;
module.KeyPaddingMask() = keyPaddingMask;
module.Reset();
module.Parameters().randu();
// // Forward test.
// arma::mat input = arma::join_cols(arma::join_cols(query, query), query);
// Forward test.
arma::mat input = arma::join_cols(arma::join_cols(query, query), query);
// module.Forward(input, output);
// REQUIRE(output.n_rows == embedDim * tLen);
// REQUIRE(output.n_cols == bsz);
module.Forward(input, output);
REQUIRE(output.n_rows == embedDim * tLen);
REQUIRE(output.n_cols == bsz);
// // Backward test.
// arma::mat gy = 0.01 * arma::randu(embedDim * tLen, bsz);
// arma::mat g;
// module.Backward(input, gy, g);
// REQUIRE(g.n_rows == input.n_rows);
// REQUIRE(g.n_cols == input.n_cols);
// Backward test.
arma::mat gy = 0.01 * arma::randu(embedDim * tLen, bsz);
arma::mat g;
module.Backward(input, gy, g);
REQUIRE(g.n_rows == input.n_rows);
REQUIRE(g.n_cols == input.n_cols);
// // Gradient test.
// arma::mat error = 0.05 * arma::randu(embedDim * tLen, bsz);
// arma::mat gradient;
// module.Gradient(input, error, gradient);
// REQUIRE(gradient.n_rows == module.Parameters().n_rows);
// REQUIRE(gradient.n_cols == module.Parameters().n_cols);
// }
// Gradient test.
arma::mat error = 0.05 * arma::randu(embedDim * tLen, bsz);
arma::mat gradient;
module.Gradient(input, error, gradient);
REQUIRE(gradient.n_rows == module.Parameters().n_rows);
REQUIRE(gradient.n_cols == module.Parameters().n_cols);
}
// /**
// * Jacobian MultiheadAttention module test.
// */
// TEST_CASE("JacobianMultiheadAttentionTest", "[ANNLayerTest]")
// {
// // Check when query = key = value.
// for (size_t i = 0; i < 5; ++i)
// {
// const size_t tgtSeqLen = 2;
// const size_t embedDim = 4;
// const size_t nHeads = 2;
// const size_t batchSize = 1;
/**
* Jacobian MultiheadAttention module test.
*/
TEST_CASE("JacobianMultiheadAttentionTest", "[ANNLayerTest]")
{
// Check when query = key = value.
for (size_t i = 0; i < 5; ++i)
{
const size_t tgtSeqLen = 2;
const size_t embedDim = 4;
const size_t nHeads = 2;
const size_t batchSize = 1;
// arma::mat query = arma::randu(embedDim * tgtSeqLen, batchSize);
// arma::mat input = arma::join_cols(arma::join_cols(query, query), query);
arma::mat query = arma::randu(embedDim * tgtSeqLen, batchSize);
arma::mat input = arma::join_cols(arma::join_cols(query, query), query);
// MultiheadAttention<> module(tgtSeqLen, tgtSeqLen, embedDim, nHeads);
// module.Parameters().randu();
MultiheadAttention module(tgtSeqLen, tgtSeqLen, embedDim, nHeads);
module.Parameters().randu();
// double error = CustomJacobianTest(module, input);
// REQUIRE(error <= 1e-5);
// }
double error = CustomJacobianTest(module, input);
REQUIRE(error <= 1e-5);
}
// // Check when key = value.
// for (size_t i = 0; i < 5; ++i)
// {
// const size_t tgtSeqLen = 2;
// const size_t srcSeqLen = math::RandInt(2, 5);
// const size_t embedDim = 4;
// const size_t nHeads = 2;
// const size_t batchSize = 1;
// Check when key = value.
for (size_t i = 0; i < 5; ++i)
{
const size_t tgtSeqLen = 2;
const size_t srcSeqLen = math::RandInt(2, 5);
const size_t embedDim = 4;
const size_t nHeads = 2;
const size_t batchSize = 1;
// arma::mat query = arma::randu(embedDim * tgtSeqLen, batchSize);
// arma::mat key = 0.091 * arma::randu(embedDim * srcSeqLen, batchSize);
// arma::mat input = arma::join_cols(arma::join_cols(query, key), key);
arma::mat query = arma::randu(embedDim * tgtSeqLen, batchSize);
arma::mat key = 0.091 * arma::randu(embedDim * srcSeqLen, batchSize);
arma::mat input = arma::join_cols(arma::join_cols(query, key), key);
// MultiheadAttention<> module(tgtSeqLen, srcSeqLen, embedDim, nHeads);
// module.Parameters().randu();
MultiheadAttention module(tgtSeqLen, srcSeqLen, embedDim, nHeads);
module.Parameters().randu();
// double error = CustomJacobianTest(module, input);
// REQUIRE(error <= 1e-5);
// }
double error = CustomJacobianTest(module, input);
REQUIRE(error <= 1e-5);
}
// // Check when query, key and value are not same.
// for (size_t i = 0; i < 5; ++i)
// {
// const size_t tgtSeqLen = 2;
// const size_t srcSeqLen = math::RandInt(2, 5);
// const size_t embedDim = 4;
// const size_t nHeads = 2;
// const size_t batchSize = 1;
// Check when query, key and value are not same.
for (size_t i = 0; i < 5; ++i)
{
const size_t tgtSeqLen = 2;
const size_t srcSeqLen = math::RandInt(2, 5);
const size_t embedDim = 4;
const size_t nHeads = 2;
const size_t batchSize = 1;
// arma::mat query = arma::randu(embedDim * tgtSeqLen, batchSize);
// arma::mat key = 0.091 * arma::randu(embedDim * srcSeqLen, batchSize);
// arma::mat value = 0.045 * arma::randu(embedDim * srcSeqLen, batchSize);
// arma::mat input = arma::join_cols(arma::join_cols(query, key), value);
arma::mat query = arma::randu(embedDim * tgtSeqLen, batchSize);
arma::mat key = 0.091 * arma::randu(embedDim * srcSeqLen, batchSize);
arma::mat value = 0.045 * arma::randu(embedDim * srcSeqLen, batchSize);
arma::mat input = arma::join_cols(arma::join_cols(query, key), value);
// MultiheadAttention<> module(tgtSeqLen, srcSeqLen, embedDim, nHeads);
// module.Parameters().randu();
MultiheadAttention module(tgtSeqLen, srcSeqLen, embedDim, nHeads);
module.Parameters().randu();
// double error = JacobianTest(module, input);
// REQUIRE(error <= 1e-5);
// }
// }
double error = JacobianTest(module, input);
REQUIRE(error <= 1e-5);
}
}
// /**
// * Numerical gradient test for MultiheadAttention layer.
// */
// TEST_CASE("GradientMultiheadAttentionTest", "[ANNLayerTest]")
// {
// struct GradientFunction
// {
// GradientFunction() :
// tgtSeqLen(2),
// srcSeqLen(2),
// embedDim(4),
// nHeads(2),
// vocabSize(5),
// batchSize(2)
// {
// input = arma::randu(embedDim * (tgtSeqLen + 2 * srcSeqLen), batchSize);
// target = arma::zeros(vocabSize, batchSize);
// for (size_t i = 0; i < target.n_elem; ++i)
// {
// const size_t label = mlpack::math::RandInt(1, vocabSize);
// target(i) = label;
// }
/**
* Numerical gradient test for MultiheadAttention layer.
*/
TEST_CASE("GradientMultiheadAttentionTest", "[ANNLayerTest]")
{
struct GradientFunction
{
GradientFunction() :
tgtSeqLen(2),
srcSeqLen(2),
embedDim(4),
nHeads(2),
vocabSize(5),
batchSize(2)
{
input = arma::randu(embedDim * (tgtSeqLen + 2 * srcSeqLen), batchSize);
target = arma::zeros(vocabSize, batchSize);
for (size_t i = 0; i < target.n_elem; ++i)
{
const size_t label = mlpack::math::RandInt(1, vocabSize);
target(i) = label;
}
// attnMask = arma::zeros(tgtSeqLen, srcSeqLen);
// for (size_t i = 0; i < tgtSeqLen; ++i)
// {
// for (size_t j = 0; j < srcSeqLen; ++j)
// {
// if (i < j)
// attnMask(i, j) = std::numeric_limits<double>::lowest();
// }
// }
attnMask = arma::zeros(tgtSeqLen, srcSeqLen);
for (size_t i = 0; i < tgtSeqLen; ++i)
{
for (size_t j = 0; j < srcSeqLen; ++j)
{
if (i < j)
attnMask(i, j) = std::numeric_limits<double>::lowest();
}
}
// keyPaddingMask = arma::zeros(1, srcSeqLen);
// keyPaddingMask(srcSeqLen - 1) = std::numeric_limits<double>::lowest();
keyPaddingMask = arma::zeros(1, srcSeqLen);
keyPaddingMask(srcSeqLen - 1) = std::numeric_limits<double>::lowest();
// model = new FFN<NegativeLogLikelihood<>, XavierInitialization>();
// model->Predictors() = input;
// model->Responses() = target;
// attnModule = new MultiheadAttention<>(tgtSeqLen, srcSeqLen,
// embedDim, nHeads);
// attnModule->AttentionMask() = attnMask;
// attnModule->KeyPaddingMask() = keyPaddingMask;
// model->Add(attnModule);
// model->Add<Linear<>>(embedDim * tgtSeqLen, vocabSize);
// model->Add<LogSoftMax<>>();
// }
model = new FFN<NegativeLogLikelihood<>, XavierInitialization>();
model->Predictors() = input;
model->Responses() = target;
// attnModule = new MultiheadAttention(tgtSeqLen, srcSeqLen, embedDim,
// nHeads);
// attnModule->AttentionMask() = attnMask;
// attnModule->KeyPaddingMask() = keyPaddingMask;
// model->Add(attnModule);
model->Add<MultiheadAttention>(tgtSeqLen, srcSeqLen, embedDim, nHeads,
attnMask, keyPaddingMask);
model->Add<Linear>(embedDim * tgtSeqLen, vocabSize);
model->Add<LogSoftMax>();
}
// ~GradientFunction()
// {
// delete model;
// }
~GradientFunction()
{
delete model;
}
// double Gradient(arma::mat& gradient) const
// {
// double error = model->Evaluate(model->Parameters(), 0, batchSize);
// model->Gradient(model->Parameters(), 0, gradient, batchSize);
// return error;
// }
double Gradient(arma::mat& gradient) const
{
double error = model->Evaluate(model->Parameters(), 0, batchSize);
model->Gradient(model->Parameters(), 0, gradient, batchSize);
return error;
}
// arma::mat& Parameters() { return model->Parameters(); }
arma::mat& Parameters() { return model->Parameters(); }
// FFN<NegativeLogLikelihood<>, XavierInitialization>* model;
// MultiheadAttention<>* attnModule;
FFN<NegativeLogLikelihood<>, XavierInitialization>* model;
// MultiheadAttention* attnModule;
// arma::mat input, target, attnMask, keyPaddingMask;
// const size_t tgtSeqLen;
// const size_t srcSeqLen;
// const size_t embedDim;
// const size_t nHeads;
// const size_t vocabSize;
// const size_t batchSize;
// } function;
arma::mat input, target, attnMask, keyPaddingMask;
const size_t tgtSeqLen;
const size_t srcSeqLen;
const size_t embedDim;
const size_t nHeads;
const size_t vocabSize;
const size_t batchSize;
} function;
// REQUIRE(CheckGradient(function) <= 3e-06);
// }
REQUIRE(CheckGradient(function) <= 3e-06);
}