# How is the complexity of the chunked attention computation in "Self Attention Does Not Need O(n2) Memory" independent from the query chunks size?

In Self-attention Does Not Need O(n{2}) Memory the authors present a say to have a constant memory complexity attention algorithm that is sequential in nature and also present an implementation that leverages the parallelisation capabilities of modern hardware by computing the attention on chunks of the queries, keys and values. This is a trade-off between the pure sequential algorithm (over each query, key and value) and having to compute, thus to store, all of the attention weights.

Without looking at the algorithm I'd say this trade-off leads to a memory complexity of $$O(query\_chunk\_size \times key\_chunk\_size)$$, at least to store the $$QK^{T}$$. The chunks here being chunks over the input/sentence length.

But the authors say this:

Assuming a chunk size of $$\sqrt{n}$$ for the keys and values, we hence obtain $$\sqrt{n}$$ summaries, giving rise to the $$O(\sqrt{n})$$ memory complexity.

I'm failing to understand why isn't the query chunk size also taking into account in the memory complexity.

Please find below their algorithm implementation:

1 import functools, jax, math
2 from jax import numpy as jnp
3
4 def _query_chunk_attention(query, key, value, precision, key_chunk_size=4096):
5  """Multi-head dot product attention with a limited number of queries."""
6  num_kv, num_heads, k_features = key.shape
7  v_features = value.shape[-1]
8  key_chunk_size = min(key_chunk_size, num_kv)
9  query = query / jnp.sqrt(k_features)
10
11  @functools.partial(jax.checkpoint, prevent_cse=False)
12  def summarize_chunk(query, key, value):
13    attn_weights = jnp.einsum(qhd,khd->qhk', query, key, precision=precision)
14    max_score = jnp.max(attn_weights, axis=-1, keepdims=True)
16    exp_weights = jnp.exp(attn_weights - max_score)
17    exp_values = jnp.einsum('vhf,qhv->qhf', value, exp_weights, precision=precision)
18    return (exp_values, exp_weights.sum(axis=-1),
20
21  def chunk_scanner(chunk_idx):
22    key_chunk = lax.dynamic_slice(
23        key, (chunk_idx, 0, 0),
25    value_chunk = lax.dynamic_slice(
26        value, (chunk_idx, 0, 0),
28    return summarize_chunk(query, key_chunk, value_chunk)
29
30  chunk_values, chunk_weights, chunk_max = lax.map(
31      chunk_scanner, xs=jnp.arange(0, num_kv, key_chunk_size))
32
33  global_max = jnp.max(chunk_max, axis=0, keepdims=True)
34  max_diffs = jnp.exp(chunk_max - global_max)
35  chunk_values *= jnp.expand_dims(max_diffs, axis=-1)
36  chunk_weights *= max_diffs
37
38  all_values = chunk_values.sum(axis=0)
39  all_weights = jnp.expand_dims(chunk_weights, -1).sum(axis=0)
40  return all_values / all_weights
41
42 def mefficient_attention_jax(query, key, value, precision=jax.lax.Precision.HIGHEST,
43                             query_chunk_size=1024):
44  """Memory-efficient multi-head dot product attention."""
45  num_q, num_heads, q_features = query.shape
46
47  def chunk_scanner(chunk_idx, _):
48    query_chunk = lax.dynamic_slice(
49        query, (chunk_idx, 0, 0),

As you can see in line 13, we have the $$qhk$$ dimensions as the output shape of the attention weights, so isn't it reasonable to expect the query chunk size to figure in the memory complexity?
To exploit the parallelism available in modern hardware, we split the computation into chunks at the cost of some additional memory. In the outer loop (lines 54-55), we split the queries in to chunks of constant size, resulting in a linear number of iterations. In each iteration of the outer loop, we call _query_chunk_attention, which itself processes the keys and values in chunks (lines 30-31). The chunks are processed sequentially and each chunk is summarized independently (lines 12 to 19). Assuming a chunk size of $$\sqrt{n}$$ for the keys and values, we hence obtain $$\sqrt{n}$$ summaries, giving rise to the $$O(\sqrt{n})$$ memory complexity.