ICLR 2024, Part 3 of 4: What Should a Model Remember?

More context can stop helping

In a classification task with many classes and examples, one participant found that adding context stopped improving performance—and sometimes made it worse. No model, dataset or metric accompanies the example, but the problem is clear: a longer context window does not guarantee effective use of its contents.

The response distinguishes extending context from improving reasoning. Data and model capabilities also matter. One participant proposes adding memory or a stack for recursive reasoning; when another suggests a bigger model, they argue that scaling only postpones the difficulty. This is a proposed direction, without a concrete architecture or formal specification of the tasks it would enable.

Asked whether context extension is solved, the participant says it is close, though unfinished. They speculate that around a billion tokens, many applications would face a different question: what is worth putting into context? Video and DNA could justify still larger windows. The conversation also floats putting a language model’s weights into another model’s context and learning over them, without proposing an implementation.

Synthetic data needs selection

The discussion turns to Nous, where a participant describes research spanning language and diffusion models, agents and reasoning, with an openness to promising outside ideas. An interviewer then challenges the synthetic-data story: if GPT-4 generates textbooks for another model to learn from, what is the substantive work beyond transferring the larger model’s knowledge?

Another participant locates that work in choosing questions, selecting examples and filtering outputs. Generation supplies material; the training pipeline determines which material transfers useful knowledge. Three examples illustrate different approaches: - Phi’s textbook-style data organizes generated material for learning. - Orca’s reasoning chains provide intermediate reasoning steps. - WizardLM’s evolving instructions change the tasks used to generate training examples. The participant credits deliberate generation techniques with substantial improvements, without giving controlled comparisons or benchmark figures.

Preference data, annotation and feedback can improve both a model and the synthetic material it subsequently generates. The closing intuition is that humans can evaluate generated material faster than they can produce comparable material themselves, although no timing or cost measurements are offered. The interviewer also suggests synthetic app- and website-generation data as a possible input to WorldSim and WebSim—a proposed connection rather than an existing integration.

FastGen: retain what each attention head uses

FastGen, introduced through the ICLR outstanding-paper honorable mention “Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs,” targets generation’s memory cost. Suyu explains the key-value (KV) cache as historical key and value vectors retained for reuse by each new decoding query. Caching avoids recomputation but occupies memory. Offloading it to CPU or storage can add latency when transfer bandwidth is limited, motivating a smaller GPU-resident cache.

Inspired by the attention patterns studied in “What Does BERT Look At?”, FastGen adapts retention to patterns observed in Llama: - Special-token heads: retain the special tokens on which attention concentrates. - Local-context heads: retain nearby context and discard distant entries. - Broad-attention heads: preserve the full cache when attention spreads across the sequence. The policy is selected per head and per input instance. Using one policy everywhere would ignore the differences that make compression possible.

FastGen searches policies that retain special tokens, punctuation, local context or combinations of those groups. It seeks minimum cache memory subject to an attention-map recovery threshold, T, described using an L1 norm. The selected policy compresses that head’s cache on the fly, requiring no pretraining or fine-tuning.

At 99% attention-map recovery, the presentation reports more than 40% cache compression for the largest evaluated Llama model. Attention recovery does not guarantee task accuracy, so the evaluation also compares generated outputs with an uncompressed full-cache baseline. GPT-4 judges instruction-model outputs, with a 50% win rate representing parity. Base-model evaluations cover code, mathematics, question answering and reading comprehension, including comparisons with non-adaptive cache methods.

Using a judge win rate above 45% as its quality criterion, the presentation reports roughly 40% memory reduction for the largest model. That criterion permits some departure from parity; detailed task scores are not given. Planned work included combining cache eviction with quantization and distillation, and extending attention-mechanism support to newer Llama models.

ZeRO++: reduce communication, then hide its overhead

At cluster scale, communication becomes a bottleneck. In ZeRO-3, model weights and optimizer states are sharded—split across devices—so training requires collective communication among them. The presenter puts communication per step at three times the model size. Motivating measurements showed communication consuming half the training time as training scaled up, even on a cluster with eight InfiniBand connections per node.

ZeRO++ combines three optimizations: - Forward pass: send quantized tensor slices. Block-based quantization reduces numerical precision on individual slices, reportedly halving communication volume at this stage. - Backward pass: keep communication within a node. Heterogeneous partitioning places the relevant model data locally, avoiding inter-node bandwidth limits. - Gradients: communicate at low precision, reduce at full precision. An all-to-all design separates transfer precision from computation precision. The gradient distinction matters because naively reducing quantized gradients accumulates quantization error and can make training diverge. The presenter explicitly distinguishes this from a software crash.

Together, these changes yield a claimed theoretical fourfold reduction in communication volume. Communication savings do not translate directly into throughput gains: quantization, dequantization and hierarchical communication also take time. Customized kernels, fusion, vectorized memory access and overlapping work help recover the savings. Sharding enables a concrete overlap: one shard communicates while another is quantized.

The presenter reports more than 2× throughput speedup on an InfiniBand configuration and 3–4× on Ethernet, without specifying the full model and cluster setup. When an interviewer suggests there is no accuracy tradeoff, the presenter corrects that interpretation: the quantization is lossy. Closely matching loss curves still leave a gap, which the presenter considers worthwhile. The remaining communication fraction is not given and cannot be inferred from the speedup alone.

At the interview, ZeRO++ was available through a flag in DeepSpeed, a framework targeting large-scale model training and expanding into inference optimization. Guanhua describes the next priority as more aggressive overlap while keeping training synchronized. Asked to rank DeepSpeed against another implementation of the ZeRO strategy, he declines: their tradeoffs require users to try them on their workloads.

State-space models: fixed memory for long sequences

The recording then turns to Sasha Rush’s architecture tutorial. Rush considers architecture a tractable route to understanding foundation models, while acknowledging data’s central and harder-to-study role. He presents state-space models (SSMs) as part of a return to recurrent neural networks: models that update a hidden state as inputs arrive and derive outputs from that state.

His comparison assumes extremely long inputs and long-form generation. For the schematic transformer under discussion, two costs grow with sequence length L: - Inference memory grows linearly: the KV cache retains keys and values for each position. - Exact attention requires quadratic training work: practical optimizations and parallel execution do not remove that asymptotic cost. Rush groups alternatives informally as Mamba-like models, emphasizing shared properties rather than treating their architectures as identical.

The recurrent alternative keeps a large but fixed-size hidden state and aims for subquadratic training. Neither property guarantees fast execution at ordinary sequence lengths. Rush’s interest in Mamba rests on combining these properties with competitive language-model loss at roughly comparable floating-point operation budgets to optimized transformers, while executing efficiently on hardware. No numerical benchmark configuration accompanies that comparison.

Selective recurrence: ignore an ad, reset at an article

Rush’s formulation maintains a matrix-valued hidden state H, with an extra dimension for state expansion. Three transformations govern it: B̄ maps input into the state, Ā controls the state’s evolution, and C maps state to output. He introduces S4, a 2021 model, by removing the nonlinearity around a vanilla recurrent network’s hidden-state update. The resulting linear time-invariant (LTI) recurrence uses fixed transformations at every position.

This recurrence is not automatically easy to train: arbitrary random parameters can behave poorly, so stabilizing the system was necessary. S4-style models performed well on long-range tasks, but the loosely grouped LTI-like models Rush discusses lagged transformers in roughly compute-normalized language-model comparisons.

Two web-text examples expose the restriction imposed by fixed transformations: - Filter an interruption. An ad appears inside useful text. Fixed B̄ transforms inputs the same way at every position, so it cannot selectively disable an update based on content. - Reset at a boundary. A page contains several articles separated by titles. The desired behavior is to discard the previous article’s state, but fixed Ā imposes the same state dynamics everywhere.

The remedy is data-dependent recurrence parameters. A neural transformation of the input predicts position-specific Aₖ, Bₖ and Cₖ, analogously to computing transformer queries, keys and values. In the schematic update Hₖ = AₖHₖ₋₁ + BₖXₖ, zeroing Aₖ removes the previous state’s contribution; zeroing Bₖ suppresses the current input’s contribution. This enables filtering and resetting while preserving a fixed-size state.

Associative scans make recurrence parallel

Data-dependent updates leave an implementation problem: how can a recurrence train efficiently on parallel hardware? Rush introduces an associative scan, which computes the combined result for every prefix of a sequence. For cumulative sums, an upward pass through a binary tree combines neighboring values; a downward pass recovers missing intermediate prefixes. His example assumes a power-of-two sequence length, and he credits Guy Blelloch’s report for the treatment.

For recurrence, replace addition with composition of state updates. Represent an update as (A, b), meaning H → AH + b, where b includes the input contribution. Consecutive updates combine as (A₁, b₁) ⊕ (A₂, b₂) = (A₂A₁, A₂b₁ + b₂). Because this operation is associative, the scan can regroup the computation while preserving update order and recover intermediate states in parallel. With zero initial state, the first two states are B₁X₁ and A₂B₁X₁ + B₂X₂.

An audience member asks whether the matrix products make this too expensive. Rush explains that most of these architectures use diagonal transition matrices, reducing updates to scalar operations. Efficient block-diagonal or full-matrix versions remain open directions in his account. Practical GPU scans are more involved than the idealized tree, but Triton’s associative-scan support provides an implementation aid.

The resulting models support parallel training and recurrent generation that retains hidden state instead of the entire history. Promising language-model loss leaves in-context learning unresolved—the ability to learn from examples supplied in context still needs investigation. Rush also stresses that these architectures continue to require large models and substantial training data; efficient context handling does not break scaling laws.

Byte-level modeling: draft tokens, verify bytes

Rush’s first application models bytes directly instead of tokenized text. The motivations include robustness to character-level noise, fewer language-specific tokenization effects and a possible common approach across modalities. But finer granularity lengthens sequences: in his comparison, bytes produce roughly four times as many positions, aggravating a standard transformer’s quadratic cost. Fixed-size byte patches shorten the sequence, but in his experience often require longer training or sacrifice final quality.

Byte-level Mamba applies the model directly to bytes, updating fixed-size state at every position and using parallel scans for training. Rush compares about 10 KB of input in most of the described experiments with about 3.5 MB occupied by the default model’s hidden states. His capacity argument is conditional: the storage is available if the model learns to use it effectively.

Generation still requires more sequential steps—roughly four byte predictions per token in Rush’s example. Speculative decoding across granularities addresses this by having one model propose a continuation and another verify it: 1. A tokenized Mamba drafts a continuation. 2. A byte-level Mamba checks the bytes using a parallel scan, accepting matches until a mismatch. 3. Byte-level generation continues to a token boundary, then token-level drafting resumes. Parallel verification reduces sequential generation and repeated weight loading. Rush reports roughly token-model generation speed, without a numerical throughput benchmark.

The reported experiments test three distinct properties: - Baseline performance: off-the-shelf Mamba outperformed transformer baselines and many patched variants in parameter-matched and compute-matched comparisons, while using less data. Exact scores and data quantities are not given. - Length generalization: training on about 8 KB extended to more than 30 KB with only small degradation in perplexity, the language-model prediction measure used in the talk. - Byte-versus-token performance: parameter-matched Mamba models reportedly converged to roughly the same perplexity, although the normalization across granularities is not specified. Rush treats these as encouraging results; unresolved scaling issues prevent a recommendation for byte-level large language models yet.

Diffusion with less input compression

The final application changes diffusion’s denoising network, which maps Xₜ₊₁ to Xₜ for one step. Rush describes self-attention operating after downsampling in U-Nets or after grouping image inputs into patches in transformer-based systems. He argues that heavier patchification harms image quality. The proposed alternative flattens the image and applies a state-space model, allowing finer-grained processing without first shortening the sequence as aggressively.

The central claim is less compression of the input representation, rather than an intrinsic superiority of state-space models over attention. Rush reports better scaling as reduced patchification produces longer sequences, and approximately state-of-the-art ImageNet conditional generation. That result comes without numerical metrics, model sizes or comparison baselines in this portion of the talk.

Rush closes by proposing video and 3D diffusion as especially promising directions: their longer sequences could make compression penalties more consequential. These remain proposed applications. Part 3 of the ICLR 2024 recording ends with that architectural opportunity—processing finer-grained inputs by changing how the model handles the sequence.