Transformer encoder

[yushiyangk]

TransformerPositionalEncoding

$$ \mathrm{PE}_{i, 2z} = \sin \left( \frac{i}{10000^{2z/d}} \right) $$

$$ \mathrm{PE}_{i, 2z + 1} = \cos \left( \frac{i}{10000^{2z/d}} \right) $$

where $i$ is the sequence position, $2z$ and $2z+1$ are the dimensions of the input embedding, and $d$ is the dimensionality of the input embedding.

The multiplicative factors $\frac{1}{10000^{2z/d}}$ are precomputed during object creation as they are constant for all $i$.

The full PE is initially precomputed for all $i$ up to 256 (configurable). This is then extended and stored if the module is called with a sequence length larger than the initial value.

Returns a 2D tensor matching the last two dimensions of the input tensor to TransformerEncoder.

FeedForward

Simple 2-layer fully connected neural network with relu activation. This is kept as a private class for now. If we want to make this to be exported it should probably be in a separate file.

TransformerEncoderLayer

As described in Vaswani et al.

TransformerEncoder

The full encoder half of the Transformer, using a Sequential containing arbitrary number of TransformerEncoderLayers.

This includes the positional encoding, but does not include any initial embedding of an input sequence into vectors (which would be separately done by e.g. word2vec)

[yushiyangk]

The test TransformerEncoder > calculates gradient is causing bun wiptest to fail silently with exit code 139 (segfault) or 138. This occurs in the result.backward() call.

Reproduced by me locally but could not be reproduced by @cryptodeal locally. We're both using macOS 12.6 and bun 0.2.2. Same error occurred for me when I used bun 0.1.13.

[bwasti]

do you have a local file called libflashlight.0.dylib anywhere? if not, maybe try bun install --force to update that dylib locally

oops never mind, I see it repros in the CI

[cryptodeal]

I know GitHub Actions Darwin runners are x86_64, so wanted to see if I could repro on my 2015 MBP (since it's AMD64), but damn, all 344 tests pass locally on that device as well (running Bun v0.2.2 and OSx 12.6).

Additional things I've tried in an attempt to induce a reproduction of this error (none of which have yet worked to repro):

• clear bun dev cache rm -rf ~/.bun/install/cache (in case there was some weird caching of shumai dependency)

[yushiyangk]

Tests are now passing for me too, after upgrading ArrayFire from 3.8.1 to 3.8.2.

As discussed on discord, the CI is still running version 3.8.1; the issue will probably be resolved once that's upgraded to 3.8.2 as well.

[yushiyangk]

Expect this to be fixed by #105. Please merge this after #103

[bwasti]

after rebase, would you mind putting the comments in this PR as comments in the transformer module?

we're using typedoc for comment formats https://typedoc.org/example/

[yushiyangk]

after rebase, would you mind putting the comments in this PR as comments in the transformer module?

we're using typedoc for comment formats https://typedoc.org/example/

Oh yes, I should definitely do that.

Edit: Added to all the transformer modules

[bwasti]

awesome work!