[Feature] LinearAttention layer for Keras

Summary

Add a backend-agnostic, serializable keras.layers.LinearAttention layer for kernelized self-attention with linear sequence complexity and constant-size state during causal streaming.

The layer should support:

• Causal and non-causal self-attention
• 2D sequence masks
• Mixed precision
• Full-sequence, chunked, and token-by-token causal inference
• Dynamic sequence lengths
• Keras serialization and Functional models
• TensorFlow, JAX, and PyTorch through keras.ops

The v1 implementation should provide a portable reference path using keras.ops, including ops.scan and prefix accumulation. Optimized backend kernels can be added later without changing the public API or checkpoint layout.

This proposal is limited to kernelized linear attention. Gated delta-rule networks, selective state-space models, cross-attention, model conversion, and complete transformer blocks are out of scope for v1.

Motivation

Keras already provides quadratic attention layers such as:

• keras.layers.MultiHeadAttention
• keras.layers.GroupQueryAttention
• Causal and sliding-window masks
• keras.ops.scan
• keras.ops.associative_scan

What is missing is a reusable linear-attention layer with a stable contract for projections, masking, streaming state, serialization and cross-backend execution.

Proposed API

layer = keras.layers.LinearAttention(
    num_heads,
    key_dim,
    value_dim=None,
    output_shape=None,
    feature_map="elu_plus_one",
    use_normalizer=True,
    normalization_epsilon=1e-6,
    causal=False,
    dropout=0.0,
    use_bias=True,
    execution_mode="auto",
    kernel_initializer="glorot_uniform",
    bias_initializer="zeros",
    kernel_regularizer=None,
    bias_regularizer=None,
    activity_regularizer=None,
    kernel_constraint=None,
    bias_constraint=None,
    seed=None,
    **kwargs,
)

Call signature:

outputs = layer(
    inputs,
    mask=None,
    initial_state=None,
    return_state=False,
    training=None,
)

v1 is self-attention only. Cross-attention can be considered separately once its state and masking semantics are defined.

Constructor

• num_heads: Number of attention heads.
• key_dim: Query/key dimension per head.
• value_dim: Value dimension per head; defaults to key_dim.
• output_shape: Output dimension; defaults to the input feature dimension.
• feature_map: "elu_plus_one" or "relu" in v1. Registered serializable callables may be supported.
• use_normalizer: Whether to normalize the output.
• normalization_epsilon: Epsilon used for numerical stability.
• causal: Enables causal attention and recurrent state.
• dropout: Applied to the attention output before the output projection.
• execution_mode: "auto", "parallel", or "recurrent".

Input / Output

inputs:  (batch_size, sequence_length, input_dim)
outputs: (batch_size, sequence_length, output_dim)

The layer accepts rank-3 inputs with a statically known final dimension. Batch size and sequence length may be dynamic.

For causal attention:

outputs, final_state = layer(
    inputs,
    return_state=True,
)

Non-causal attention does not expose streaming state.

Mathematical definition

After projection:

Q, K ∈ R^(B × T × H × Dk)
V    ∈ R^(B × T × H × Dv)

Apply a non-negative feature map φ to Q and K.

Non-causal

S = Σj φ(Kj)^T Vj
Z = Σj φ(Kj)

Yi = φ(Qi) S / (φ(Qi) Z + ε)

When use_normalizer=False:

Yi = φ(Qi) S

No (T, T) attention matrix is materialized.

Causal

St = S(t-1) + φ(Kt)^T Vt
Zt = Z(t-1) + φ(Kt)

Yt = φ(Qt) St / (φ(Qt) Zt + ε)

The current token is included in its own context.

Feature maps

elu_plus_one

φ(x) = elu(x) + 1

Recommended default because it is deterministic, non-negative, and backend-independent.

relu

φ(x) = relu(x) + normalization_epsilon

Random feature maps are deferred from v1.

Streaming State

For normalized causal attention:

state = (key_value_state, key_state)

key_value_state: (batch, heads, key_dim, value_dim)
key_state:       (batch, heads, key_dim)

For use_normalizer=False, only key_value_state is required.

State is runtime data, not layer weights, and is excluded from serialization.

state = layer.get_initial_state(batch_size=batch_size)

y1, state = layer(
    chunk_1,
    initial_state=state,
    return_state=True,
)

y2, state = layer(
    chunk_2,
    initial_state=state,
    return_state=True,
)

With dropout disabled, chunked causal execution should match a single full-sequence call.

Masking

mask is a boolean tensor broadcastable to:

(batch_size, sequence_length)

Non-causal

• Invalid keys/values do not contribute to the accumulated statistics.
• Invalid query positions produce zero output.
• The input mask is propagated.

Causal

Invalid timesteps must not update the recurrent state:

candidate_S = S_previous + φ(Kt)^T Vt
candidate_Z = Z_previous + φ(Kt)

St = where(mask_t, candidate_S, S_previous)
Zt = where(mask_t, candidate_Z, Z_previous)
Yt = where(mask_t, readout(St, Zt), 0)

v1 supports only 2D sequence-validity masks. Arbitrary 3D pairwise masks are out of scope.

Execution Modes

"recurrent"

Uses keras.ops.scan and provides the canonical streaming/reference implementation.

"parallel"

• Non-causal: sequence reductions.
• Causal: prefix accumulation using ops.cumsum or ops.associative_scan.

"auto"

default:

• Non-causal → parallel
• Causal with state or single-token input → recurrent
• Causal sequence input → parallel when supported
• Unsupported optimized paths → portable fallback

Execution mode must not change weights or serialization.

Dropout

Because linear attention does not materialize attention probabilities, v1 dropout applies to the concatenated attention output before the output projection.

With training=False or dropout=0, full-sequence and chunked outputs must be equivalent.

Numerical Stability

For float16/bfloat16 computation:

• Accumulate state and normalization statistics in at least float32.
• Apply feature maps in accumulator dtype.
• Use normalization_epsilon.
• Return outputs in the layer compute dtype.
• Return recurrent state in accumulator dtype.
• Reject non-floating inputs.
• Test long sequences for finite outputs and gradients.

Float64 should be preserved where supported.

Implementation

Files:

keras/src/layers/attention/
    linear_attention.py
    linear_attention_test.py

Export with:

@keras_export("keras.layers.LinearAttention")
class LinearAttention(Layer):
    ...

Use existing Keras projection patterns such as EinsumDense where appropriate.

The implementation should use public keras.ops primitives:

• ops.einsum / ops.matmul
• ops.sum
• ops.cumsum / ops.associative_scan
• ops.scan
• ops.where
• ops.transpose
• ops.reshape
• ops.cast

No Python loop should iterate over sequence positions inside call().

Optimized backend implementations can be introduced later behind internal dispatch while retaining the portable implementation as the correctness reference.

Complexity

MethodTimePairwise memoryStreaming state
Softmax attention	O(T² H Dk)	O(T² H)	Grows with T
Linear attention	O(T H Dk Dv)	No T × T matrix	O(H Dk Dv + H Dk)

Linear complexity does not guarantee lower wall-clock latency for short sequences; benchmarks should determine the practical crossover points.

Test Plan

Layer behavior

• Shapes and dynamic sequence lengths
• Constructor/call validation
• Functional and subclassed models
• Serialization and cloning
• .keras save/load
• Initializers, regularizers, and constraints
• Mixed precision
• Dropout and seeds
• Explicit and implicit masks
• JIT/compiled execution

Correctness

• Match a direct linear-attention reference implementation.
• Match an explicitly unrolled causal recurrence.
• Match full, chunked, and token-by-token causal execution.
• Match recurrent and parallel execution.
• Verify current-token inclusion.
• Verify no (T, T) tensor is created.

Mask and state behavior

• Right padding
• Left padding and interior gaps
• All-false/all-true masks
• Invalid state structures and shapes
• State dtype under mixed precision
• State excluded from serialization

Stability

• Finite outputs and gradients on long sequences
• Near-zero denominators
• Float64 numerical-gradient checks where supported

Acceptance only if

• keras.layers.LinearAttention is exported and documented.
• Supports causal and non-causal self-attention.
• Supports dynamic sequence lengths.
• Causal mode supports initial_state, return_state, and get_initial_state().
• No (T, T) attention tensor is materialized.
• Full, chunked, and token-level causal outputs match when dropout is disabled.
• Recurrent and parallel implementations have forward/gradient parity within documented tolerances.
• Masks preserve state and zero invalid outputs.
• TensorFlow, JAX, and PyTorch portable paths work.
• Mixed-precision state uses float32 accumulation.
• .keras save/load and cloning preserve behavior.
• Streaming state remains constant in sequence length.

Rollout

1. Core layer: non-causal and recurrent causal paths, state, masking, serialization, and tests.
2. Parallel causal path: prefix accumulation, automatic execution selection, parity tests, and benchmarks.
3. Documentation and optimization: guides/examples followed by optional backend-specific kernels.

My questions to maintainers:

1. Should execution_mode be public in v1, or only automatic selection with a private test override?
2. Should output_shape initially accept only an integer?
3. Should custom feature maps be supported in v1, or only named deterministic maps?
4. Should state always use accumulator dtype, or should state_dtype be configurable?
5. Should dropout remain the name for output dropout, or should it be output_dropout to avoid confusion with MultiHeadAttention?
6. Should causal remain constructor-only for clearer serialization and tracing?
7. Should get_initial_state() require the layer to be built?
8. Should "elu_plus_one" be the default feature map?
9. Should cross-attention become a follow-up to this layer or a separate API?