SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Listing 2: Memory attention: Let  $\ell$  be a memory attention layer. During the training, we make  $\ell$  attend to the (key, value) pairs from the local context, previous local context, and  $d - 1$  contexts coming from other documents. During the inference, queries from  $\ell$  attend to the local context and  $k$  nearest neighbors retrieved from memory. For simplicity, we provide the code for one head and assume that the crossbatch dimension  $d$  is equal to the batch size.
def mem_attn_layer(Q1, K1, V1, C1, Km, Vm, Kp, Vp, attn_scf, mode):
    '''Attention mechanism for crossbatch and memory attention
Args:
    Q1, K1, V1: tensors of shape [batch, ctx_len, dim]
    with queries, keys and values from the local context
    Km, Vm: tensors of shape [batch, ctx_len, k, dim]
    with k most matching memory keys for
        each query from Q1 along with associated values
    Kp, Vp: tensors of shape [batch, ctx_len, dim] with
        keys and values from the previous context
    attn_scf: a scale factor used before softmax
    mode: either training or inference
Returns:
    y: a vector with shape [batch, ctx_len, dim]
    '''#
    # attention to the local context
    localattention = jnp.einsum("bqd,bkd ->bqk", Q1, K1)
    localattention *= attn_scf
    localattention = apply_causal_mask(localattention)
if mode == "train":
    # In train mode, we additionally use previous context
    # and batch-1 contexts from other documents.
    prev attention = jnp.einsum("bqd,ckd ->bqck", Q1, Kp)
    shape = prevattention.shape
    additional attention = prevattention.reshape(shape[-2] + (-1,))
elif mode == "inference":
    # In the inference mode, we additionally use nearest neighbors
    # retrieved from memory. We retrieve k (key, value)
    # pairs for each query.
    memoryattention = jnp.einsum("bqd,bqnd ->bqn", Q1, Km)
    additional attention = memoryattention
else:
    raise Exception("Only train and inference modes are supported")
additionalattention *= attn_scf
# We merge the raw attention scores and calculate the softmax
combined attention = jnpconcatenate([localattention,
                          additionalattention],
                          axis=-1)
combinedweights = jax(nn柔软(max(combined attention, axis=-1)
ctx_len = Q1.shape[1]
localweights = combinedweights[,..., :ctx_len]
additionalweights = combinedweights[,..., ctx_len])
y = jnp.einsum("bqk,bkd ->bqd", localweights, V1)
if mode == "train":
    prevweights = additionalweights
shape = prevweights.shape
prevweights = prevweights.reshape(shape[-1] + (-1, ctx_len))
y += jnp.einsum("bqck,ckd ->bqd", prevweights, Vp)
else:
    memoryweights = additionalweights
y += jnp.einsum("bqn,bqnd ->bqd", memoryweights, Vm)
return y