Overcoming the quadratic scaling in transformer architecture

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. 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.