Do you know any papers that try to overcome quadratic scaling problems by attending lower dimensional representations in the dimension of tokens?

For example, let's say that the input to the transformer is (batch_size, token_number, hidden_dimension).

Do you think it would make some sense to do let's say one Transformer Block (TB) and then map (by some weight matrix) token_number to let's say token_number/2 and do TB again?

The graph might look like this:

input_tokens -> TB -> hidden_tokens -> TB -> ... -> output_tokens (might be equal to input_tokens in dimension)

So, our shapes would be:

(B, T, H) -> (B, T/2, H) -> (B, T/4, H) -> ... -> (B, T/4, H) -> (B, T/2, H) -> (B, T, H)

where B - batch_size, T - token_number, H - hidden_dimension


1 Answer 1


Yes. What you're describing seems to be similar to the Hourglass Transformer, the Funnel Transformer, among others.

Looking at the first paper, they consider two ways of downsampling tokens (upsampling methods are similar).

  1. Linear pooling: reshape (n, d) -> (n/k, k * d), with k being the downsampling factor. Then, use a linear map to project k * d to your desired output dimension.
  2. Attention pooling: First do simple 1D max/mean pooling to downsample across n: (n, d) -> (n', d); n' < n. Then, pass that into an attention layer (with residual connects), where the queries are the downsampled embeddings, and the keys & values are the original embeddings. For example, you may start off with x.shape = (10, 16), then do mean pooling to get x_pooled.shape = (5, 16). Now, each vector in x_pooled gets converted to a weighted sum of vectors in x (attention operation) and added to the original vector in x_pooled (residual). This is described in more detail in the second paper.

There's also approaches that instead do a discrete selection over tokens (e.g., Pyramid-BERT).

  • $\begingroup$ This is exactly what I am looking for. Thank you for your answer! $\endgroup$ Jan 20 at 16:24
  • $\begingroup$ @AndyYermakov No problem! Also, please accept the answer if it sufficiently answers your question. $\endgroup$ Jan 20 at 17:51

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .