Lattice: Learning to Efficiently Compress the Memory

• Attention’s quadratic time and space complexity and unbounded cache limit its scalability for long sequences and resource-constrained settings.Source: §1 · p. 1
• Existing recurrent alternatives trade off expressivity, accuracy, or memory capacity, while some linear updates cannot selectively discard redundant information.Source: §1 · p. 1
• Lattice formulates cache compression as online optimization, deriving state- and input-conditioned updates that write only information orthogonal to each slot’s current state.The paper also derives a chunk-wise parallel form and reports superior memory efficiency and language modeling capabilities.Source: §1 · p. 2

• The key-value cache grows linearly with sequence length, making retrieval linear per token and full-sequence computation O(T^2).Source: §2 · pp. 2–3
• Linear recurrent updates support sub-quadratic parallel training and constant-time per-token inference.Training can use chunkwise computation or parallel scan.Source: §2 · p. 3
• A sliding window bounds cache size by retaining the latest m tokens, but its limited receptive field restricts long-range dependencies and global context.Observed low-rank structure in key-value matrices motivates selectively retaining essential context instead of simply truncating memory.Source: §2 · p. 3
• Linear-attention updates add an input-dependent outer product independently of the current state, which can cause interference when tokens compete for limited slots.Source: §2 · p. 3
• State-dependent gating could use each slot’s contents to decide what to add or discard, but naive state conditioning can impede parallelization.Source: §2 · p. 3

3 Compression Layer

• Lattice stores contextual history in a compact state matrix and uses that compressed state for readout, trading memory size and computational cost against query precision.The state is updated as tokens arrive; the compression objective reconstructs the current token embedding from its latent code.Source: §3 · p. 4

3.1 Decoding Layer

• Each token’s latent code and embedding define a reconstruction objective, and a single online gradient-descent step updates the internal state while outer training learns the projection weights.The inner loop adapts state within each sequence; the outer loop trains model parameters across sequences.Source: §3.1 · pp. 4–5

3.1.1 State Normalization

• The normalized-state gradient projects each slot’s input update onto the space orthogonal to that slot, discarding aligned content so updates add non-redundant information.The writing intensity controls each token’s contribution to each slot; OSR names the unified orthogonal update formulation.Source: §3.1.1 · p. 6

3.2 Stabilizing Memory Updates via Normalization

• Orthogonal updates increase each slot’s norm monotonically, so Lattice projects updated states back onto the unit sphere to constrain their magnitude.The scaling factor preserves the steepest-descent direction and acts analogously to a forgetting gate by adaptively normalizing the effective step size.Source: §3.2 · p. 6

3.3 Forgetting by State Regularization

• L2 state regularization yields a recurrence with a forget gate controlling retained past memory, while an alternative L1 objective promotes sparse memory representations.The paper focuses on L2 regularization because it yields scaling weight decay in the recurrence.Source: §3.3 · p. 7

3.4 Efficient Computation

• Writing each orthogonal projection as a dot product and scalar-vector multiplication avoids a full matrix-vector product, reducing the recurrence cost to O(dm).The explicit form also eliminates vector-Jacobian-product computations.Source: §3.4 · p. 8

3.4.1 Parallel and Hardware Efficient Form

• Chunk-wise parallelization approximates gradients using each chunk’s starting state, linearizing the nonlinear recurrence; an alternative linearizes only the projected-input term for efficient parallel evaluation.Figure 2 plots perplexity against context length for 110M Books models and against model size on FineWeb-Edu; Transformer++ results are capped at T≤8k.Source: §3.4.1 · p. 8, Figure 2 · p. 9

• Language-modeling and commonsense models use FineWeb-Edu at 4k context, while Books3 scaling experiments span training contexts from 512 to 16k tokens.Training batch size remains 0.5M tokens across sequence lengths.Source: §4 · p. 9
• Lattice achieves lower perplexity than all baselines across context lengths and model sizes, with gains over other linear RNNs growing as sequences lengthen.The comparison is reported for Figures 2 and 5.Source: §4 · p. 9
• Lattice outperforms baseline models on most zero-shot commonsense tasks and achieves the highest average accuracy.Table 2 compares language-modeling and zero-shot commonsense performance for models trained on FineWeb-Edu.Source: §4 · p. 9, Table 2 · p. 10
• On Multi-Query Associative Recall, Lattice improves recall accuracy, particularly as sequence length increases.The evaluation assesses memory capacity and in-context learning.Source: §4 · p. 9
• Orthogonal state recurrence and normalized projection significantly contribute to overall performance, while the explicit forget gate has negligible impact.The authors suggest normalized projection inherently acts as an effective forgetting mechanism.Source: §4 · p. 9
• With m=d=64, the proposed orthogonal update mechanism is reported to utilize memory capacity significantly more efficiently; extended ablation studies examine scaling memory size and per-token versus per-chunk normalization.Source: §4 · p. 10

5 Conclusion

• Lattice compresses information into a matrix-valued state with limited slots, using a single gradient-descent step to derive state- and input-dependent updates that write orthogonal, non-redundant information.The update includes an input- and state-dependent writing intensity and has sub-quadratic complexity.Source: §5 · p. 10

A Discussion and Related Works

• The compression layer acts as a fast network whose memory states adapt to context, while the outer loop learns generalizable parameters; the state update also implements test-time training.The discussion connects this framing to Fast Weight Programmers [Schmidhuber, 1992, Schlag et al., 2021] and test-time training [Sun et al., 2020, 2024, Behrouz et al., 2024b, von Oswald et al., 2025].Source: A Discussion and Related Works · p. 17
• Lattice’s memory adaptation resembles LMS updates based on instantaneous squared-error gradients, but uses a nonlinear rule that retains non-redundant components of new tokens.The discussion also draws parallels between normalization and NLMS, and between state decay and Leaky LMS.Source: A Discussion and Related Works · p. 17
• Unlike online matrix factorization methods that learn latent codes with block coordinate descent, Lattice uses an encoding layer within a larger neural-network training procedure.The encoding layer generates the latent representation k_t.Source: A Discussion and Related Works · p. 17
• Removing the nonlinearity recovers the delta rule, associated with higher memory capacity and associative recall [Schlag et al., 2021, Yang et al., 2024c]; Lattice instead uses nonlinear interactions with non-redundant new-key information.Source: A Discussion and Related Works · p. 18
• An alternative normalization applies to the decoded output rather than state columns, producing a nonlinear recurrence with an error correction projected onto the orthogonal complement of the decoded output.The discussion relates state-vector normalization and output normalization to weight and layer normalization [Salimans and Kingma, 2016; Ba, 2016].Source: A Discussion and Related Works · p. 18

• For each chunk, the gradient is approximated using the state at its beginning, which is the preceding chunk’s final state, to linearize the nonlinear recurrence.The approximation evaluates the chunk’s gradients in parallel.Source: §B.1 · p. 18
• The chunk input X_b stacks the tokens in that chunk, with x_bτ denoting the input at its τ-th step.Source: §B.1 · p. 18
• The chunk-level recurrence uses an indicator for i≤j and element-wise division, with its intermediate quantities computed by tensor contraction or Einstein summation.Source: §B.1 · p. 19
• An alternative linearizes only the h-hat term; its rank-one G-hat parameterization lets intra-chunk updates use parallel gated linear attention.Source: §B.1 · p. 19
• Normalization factors β can be computed with matrix multiplication, avoiding materialization of each ΔS_τ and its associated memory and compute inefficiency.Source: §B.1 · p. 19
• A matrix form computes states only at block ends, reducing computation and I/O overhead while using accelerator-friendly matrix multiplications.The appendix describes chunk-wide matrix multiplication for dS_τ.Source: §B.1 · p. 19

• The encoding layer maps value v_t to latent key k-hat_t using the transposed feature map of the state, ϕ(S_t)^T v_t, and pairs this encoder with an ℓ2 loss.The formulation is motivated by viewing PCA as linear regression onto a lower-dimensional latent space.Source: §B.2 · p. 19
• A closed-form gradient expression yields the encoding layer’s recurrent update.Source: §B.2 · p. 19
• The paper calls the general update framework Orthogonal State Recurrence (OSR), with h_t and c_t defined differently for each compression layer.The framework unifies the update formulations presented in equations 6, 25, and 8.Source: §B.2 · p. 20

B.3.1 Encoder Layer with State Normalization

• Each memory slot is normalized to form a basis, enabling reconstruction error to be decomposed and differentiated separately by column.The matrix gradient uses the projection onto each slot’s orthogonal complement.Source: §B.3.1 · p. 20
• The update adds the incoming token’s orthogonal component to the current slot, unlike linear attention’s superposition update, which adds the full token representation.Figure 3 assumes unit writing intensity for both updates.Source: Figure 3 · p. 21
• Projecting the updated slot back onto the unit sphere is equivalent to Riemannian gradient descent with a retraction step.The retraction replaces the sphere’s exponential map with a first-order approximation.Source: §B.3.1 · p. 21

Training setup

• Language-modeling and commonsense-reasoning models use FineWeb-Edu at 4k context, with a default Lattice chunk size of C=4; exact-compression evaluations use C=1.All experiments use batches of 0.5 million tokens, while sub-quadratic models follow the Mamba-style Conv1D and GeLU post-gate architecture.Source: C Experiment Details · pp. 21–23, Figure 4 · p. 22

Context-length evaluation

• The 110M-parameter context-length evaluation tests Books3 from 512 to 16k tokens and The Pile at 2k and 8k; Transformer results stop at 8k.The caption notes that Transformers pretrained from scratch often perform poorly at very long contexts and are commonly finetuned from shorter-context models.Source: C Experiment Details · p. 22, Figure 5 · p. 23

Ablation: Scaling Memory Size

• With 16 slots (m=d/4), Lattice outperforms TTT with 64 slots, demonstrating a 4× compression gain.Increasing m continues to improve performance, but m>d increases parameters in the Wq and Wk embedding projections.Source: C Experiment Details · p. 23

Ablation: Per-token vs. Per-chunk Normalization

• Per-token normalization better preserves performance and stability as chunk size grows; per-chunk normalization diverges more and destabilizes the C=16 simplified approximation.The authors attribute this to differences between intra-chunk and inter-chunk update dynamics.Source: C Experiment Details · p. 24, Table 5 · p. 24

C.1 Recall-Intensive Task: Multi-Query Associative Recall

• MQAR requires multiple recalls at varying positions in one pass, with difficulty increased across 48, 96, and 256 key-value pairs.Evaluations use context lengths 1024 and 2048; models use two single-head blocks and set m=d for Lattice, TTT, DeltaNet, and GLA.Source: §C.1 · pp. 24–25
• Lattice achieves the best MQAR accuracy in most settings, significantly improving over TTT and DeltaNet while maintaining accuracy across context lengths.The task tests storing and retrieving many key-value associations in fixed recurrent memory.Source: §C.1 · p. 25, Figure 7 · p. 25

C.2 Throughput Comparison.

• Lattice’s throughput scales better than Transformer++ as sequence length increases and remains competitive with TTT.The comparison uses 110M-parameter models on 2×2 TPU Trillium, with chunk sizes 64 and 16; Lattice and TTT use pure JAX without specialized kernels.Source: §C.2 · p. 25, Figure 8 · p. 26

D Implementation

• The chunk implementation reshapes keys, values, queries, and gates by chunk, then scans over chunks and returns the final state and outputs.The listing also provides a recurrent single-step compression operation.Source: D Implementation · p. 27
• The chunk computation forms normalized token projections and gates before applying a parallel intra-chunk operation based on linear attention.The listing identifies the projection as Equation 14 and the gate as Equation 17.Source: D Implementation · p. 28
• The nonlinear test-time-training update prevents inter-chunk associative scans, limiting hardware utilization despite parallelizable matrix operations within chunks.The orthogonal projection also adds O(dm) computation relative to simpler linear RNNs.Source: D Implementation · p. 28