Best way to compute the moving average of word vectors in JAX

user_6084245 0
Viewed 87

Let say I have a matrix W of shape (n_words, model_dim) where n_words is the number of words in a sentence and model_dim is the dimension of the space where the word vectors are represented. What is the fastest way to compute the moving average of these vectors ?

For example, with a window size of 2 (window length = 5), I could have something like this (which raises an error TypeError: JAX 'Tracer' objects do not support item assignment):

from jax import random
import jax.numpy as jnp

# Fake word vectors (17 words vectors of dimension 32)
W = random.normal(random.PRNGKey(0), shape=(17, 32)) 

ws = 2          # window size
N = W.shape[0]  # number of words

new_W = jnp.zeros(W.shape)

for i in range(N):
    window = W[max(0, i-ws):min(N, i+ws+1)]
    n = window.shape[0]
    for j in range(n):
        new_W[i] += W[j] / n

I guess there is a faster solution with jnp.convolve but I'm not familiar with it.

jax python
1 Answers

This looks like you're trying to do a convolution, so jnp.convolve or similar would likely be a more performant approach.

That said, your example is a bit strange because n is never larger than 4, so you never access any but the first four elements of W. Also, you overwrite the previous value in each iteration of the inner loop, so each row of new_W just contained a scaled copy of one of the first four rows of W.

Changing your code to what I think you meant and using index_update to make it compatible with JAX's immutable arrays gives this:

from jax import random
import jax.numpy as jnp

# Fake word vectors (17 words vectors of dimension 32)
W = random.normal(random.PRNGKey(0), shape=(17, 32)) 

ws = 2          # window size
N = W.shape[0]  # number of words

new_W = jnp.zeros(W.shape)

for i in range(N):
    window = W[max(0, i-ws):min(N, i+ws)]
    n = window.shape[0]
    for j in range(n):
      new_W = new_W.at[i].add(window[j] / n)

and here is the equivalent in terms of a much more efficient convolution:

from jax.scipy.signal import convolve
kernel = jnp.ones((4, 1))
new_W_2 = convolve(W, kernel, mode='same') / convolve(jnp.ones_like(W), kernel, mode='same')

jnp.allclose(new_W, new_W_2)
# True
user_29378310
Related