Author: Ruchir Tewari

Chronological list of known learned representations (increasing date)

· Leave a comment ·

Chronological list of known learned representations that were explicitly identified, named, and evidenced in a paper/post with reproducible analysis.

The representation basis answers “what algebra the model chooses to live in. The circuit answers “how the transformer computes in that algebra.”

First reported (approx)Representation (what it is)Where it shows upCanonical referenceImportance & generality (researcher comment)
1996Sparse / wavelet-like (Gabor-like) receptive-field basesUnsupervised vision models learning efficient codes for natural imagesOlshausen & Field, Nature 1996 (Courses at Washington University)This is one of the earliest clean demonstrations that optimizing a simple objective (sparsity/efficient coding) yields structured bases resembling classical signal representations. It is highly general for natural-image statistics and still conceptually underlies why “edge-like” first-layer features are so universal.
2013 (Jan)Linear semantic substructure in word-vector spaces (directions encode relations; analogies ≈ parallelograms)Word embeddings from neural objectivesMikolov et al. 2013 (word2vec) (arXiv) and Pennington et al. 2014 (GloVe explicitly discusses the analogy geometry) (Stanford NLP)This made “distributed representations” operational: relations become approximately linear operators/directions. Generality is high across corpora and embedding methods, though the reliability of specific analogies varies and is not guaranteed by training.
2013–2014 (Nov → ECCV)Early CNN layers learn oriented edge / color-opponency filters (Gabor-like)Supervised convnets on natural imagesZeiler & Fergus visualization work (arXiv)Important because it empirically tied deep vision features to classical linear-systems intuition: even with end-to-end supervision, the network “chooses” a near-optimal front-end basis for images. Very general across CNN families trained on natural images.
2014 (Oct)Differentiable addressing representations (content- and location-based “attention” over external memory)Memory-augmented networksGraves et al., Neural Turing Machines (arXiv)This is a representation of state and retrieval rather than of sensory input: key/value-like addressing emerges as a learnable interface between computation and storage. Generality is moderate: powerful, but most mainstream models replaced explicit external memory with transformer attention over context.
2015 (Nov)Convolutional algorithmic state representations (Neural GPU learns internal states that generalize addition/multiplication to long lengths)Algorithm learning on sequencesKaiser & Sutskever, Neural GPUs Learn Algorithms (arXiv)This is a landmark for “nets can learn algorithmic latent states,” not just pattern matching. Generality is medium: it works well for certain algorithmic tasks with the right inductive bias, but is not a universal recipe for systematic generalization.
2017 (Oct)Capsule pose-vector representations (entity presence + instantiation parameters; routing groups parts into wholes)Vision architectures emphasizing part–whole structureSabour et al., Dynamic Routing Between Capsules (arXiv)Conceptually important: it proposes a factorized internal code (pose/part structure) rather than “bags of features.” Generality is debated in mainstream practice, but the representational idea is crisp and has influenced later equivariant and compositional approaches.
2018 (Mar)Grid-like spatial codes (grid/border/band-cell-like units)RNNs trained for path integration / navigationCueva & Wei 2018 (arXiv)Very important scientifically: it shows a strong convergence between trained artificial networks and biological coding hypotheses. Generality is high within navigation/path-integration objectives; less directly portable to arbitrary domains.
2018 (Aug)Explicit arithmetic representations via specialized units (linear codes + gated primitive ops)Neural arithmetic modulesTrask et al., NALU (arXiv)This line is important because it cleanly separates “representation of quantity” from “operators on quantities,” targeting extrapolation. Generality is medium: works best when the task truly factors into arithmetic primitives and the architecture is used appropriately.
2020 (Jun)Fourier-feature positional encodings / spectral reparameterizations (map inputs through sinusoidal features to defeat spectral bias)Implicit neural representations; MLPs for signals/scenesTancik et al., Fourier Features… (NeurIPS Papers)Important as a unifying explanation for why plain MLPs underfit high frequencies and how a spectral basis fixes it. Generality is high for continuous regression/INR tasks; it is partly “designed,” but it formalizes the representational need very clearly.
2022 (Sep)Induction-head representations (“copy-from-previous-match” algorithm; pointer-like behavior)Transformers doing in-context learning / pattern completionOlsson et al., In-context Learning and Induction Heads (arXiv)This is one of the most important circuit-level representational discoveries in transformers: it identifies a reusable mechanism that looks like learned algorithmic pointer-chasing. Generality is high across autoregressive transformers and many ICL-like behaviors.
2022 (Sep)Superposition of features (many sparse features packed into fewer dimensions; polysemanticity as a geometric tradeoff)ReLU nets and plausibly large modelsElhage et al., Toy Models of Superposition (arXiv)Foundational for interpretability: it reframes “neurons are messy” as “the representation is compressed and distributed by necessity.” Generality is extremely high—this is an architectural/optimization-level phenomenon, not a task-specific trick.
2023 (Jan)Discrete Fourier Transform (DFT) / trig-identity representation for modular additionSmall transformers that grok modular arithmeticNanda et al., Progress measures for grokking via mechanistic interpretability (arXiv) (plus walkthrough (Neel Nanda))The model represents elements in a Fourier basis where modular addition becomes phase addition/rotation. Importance is high as a proof-of-mechanism (nets rediscover classic algebraic representations). Generality is moderate: strongest for tasks with group structure (cyclic groups, convolutions, periodicity).
2023 (Mar–Sep)Linear “world-state” representations in sequence models (latent state corresponds to board state; controllable by vector arithmetic)Othello-GPT-style modelsNanda’s exposition (Neel Nanda) and the associated paper on emergent linear representations (arXiv)Important because it shows a model trained only to predict tokens can learn an explicit internal state (a “world model”) that is linearly recoverable and causally editable. Generality is promising but not universal; it likely emerges when the task forces consistent latent state tracking.
2023 (Oct)Feature dictionaries / “monosemantic” features via sparse autoencoders (dictionary learning on activations)Mechanistic interpretability for transformersAnthropic’s “Towards Monosemanticity” line (Anthropic)This is less “the model’s native representation” and more “a recovered basis that better matches it,” but it’s crucial: it suggests models are organized around a large set of sparse features even when neurons are polysemantic mixtures. Generality is likely high, and it directly shapes practical interpretability workflows.
2024 (Feb, community analysis)Chess/Othello-like linear world representations (extensions/replications)Board-game GPTs; “world model” probing and interventionsExample community writeup (LessWrong)This is a continuation/expansion of the 2023 world-representation finding. Importance depends on replication rigor, but it is part of the emerging picture that “latent-state tracking” is a common representational strategy in sequence models under the right data/task constraints.

Update: Some more interesting representations

1) Finite-state / automaton-like representations (regular languages)

Transformers trained on formal languages can end up simulating automata, and recent work explicitly extracts finite state machines from trained transformers to characterize what they learned. This is close to “boolean/bitmap logic” in that the latent state is discrete and transitions are rule-like. https://arxiv.org/pdf/2410.06045

2) Stack-like representations for parentheses / Dyck-style tasks

Balanced bracket classification tasks are widely used in mech-interp pedagogy because they pressure the model toward a latent “depth” or stack surrogate. In practice, small transformers often learn a distributed state that tracks nesting structure, sometimes in a way that can be probed linearly.  https://arena-chapter1-transformer-interp.streamlit.app/%5B1.5.1%5D_Balanced_Bracket_Classifier

3) “World-state bitmaps” (board-state as a linear code)

In Othello-GPT-style settings, the residual stream contains a linearly recoverable encoding of the board. This is arguably a learned bitmap-like representation (one direction per square / feature), embedded in a continuous space.  https://www.neelnanda.io/mechanistic-interpretability/othello

4) Group-operation representations beyond modular addition

A closely related line studies how small nets learn group composition more broadly (a “universality” testbed). This generalizes the “DFT for cyclic groups” story into a broader family of algebraic representations and circuits.  https://openreview.net/pdf?id=jCOrkuUpss

5) Boolean satisfiability style reasoning (logical structure)

There is mechanistic-interpretability work on transformer-based models trained to solve 2-SAT, which is a canonical boolean-logic problem. This is a direct example of boolean structure expressed in transformer activations and circuits.  https://arxiv.org/html/2407.13594v1

6) Induction / copy (pointer-style algorithm)

Not boolean algebra per se, but it is a very simple learned algorithmic representation: a head learns to represent and retrieve repeated patterns (“copy what followed last time”). This often coexists with more symbolic-feeling representations in toy tasks.  https://arxiv.org/abs/2312.03002

Learned Representations in Neural Networks

· Leave a comment ·


Neural networks transform raw inputs — pixels, text, audio — into internal descriptions built layer by layer through learned weights and nonlinearities. The core mechanism is hierarchical composition: early layers detect local patterns like edges or n-gram features, while deeper layers combine these into abstract structures like object parts, semantic concepts, or reasoning patterns. Rather than relying on hand-engineered features, the network discovers whatever internal geometry best serves its training objective.

Representation spaces are not mere lookup tables; they are high-dimensional manifolds with structure that can be analyzed with the tools of differential geometry and information geometry. The Fisher information metric, for instance, naturally measures distances between probability distributions that a network implicitly encodes, connecting the curvature of representation space to the model’s sensitivity and generalization behavior.

More visibly, semantic relationships in language models manifest as linear directions in activation space, enabling vector arithmetic over meaning. This regularity reflects the network solving a smooth optimization problem in which nearby inputs on the data manifold are mapped to nearby points in representation space.

A critical consequence of this structure is transferability. Representations learned on large datasets tend to capture the intrinsic geometry of the data distribution itself, making them reusable across tasks. This underpins the modern pretrain-and-adapt paradigm: a foundation model distills general representational structure from vast data, and fine-tuning merely redirects it.
Interpretability research has complicated this picture. Networks appear to use superposition, encoding more features than they have dimensions by distributing concepts across overlapping, near-orthogonal directions rather than isolated neurons. This is geometrically efficient — nearly orthogonal vectors in high dimensions allow exponentially many features to coexist — but it makes the representation space harder to read.

Understanding a model now requires studying directions, circuits, and geodesics in activation space, not individual units. This is the project of mechanistic interpretability: recovering the internal computational geometry that produces a model’s behavior.

Three frontiers concentrate current research.

– First, what geometric properties of a representation predict its generalizability — smoothness, dimensionality, curvature of the learned manifold?

– Second, how do large language models encode causal relations, abstractions, and multi-step reasoning, and does this reflect genuine geometric structure or brittle surface statistics?

– Third, can training objectives be designed to produce representations that are sparse, disentangled, or causally structured by construction — making the geometry legible from the start rather than reverse-engineered after the fact? This last question connects representation learning directly to AI safety: systems whose internal geometry can be inspected and tested are systems whose behavior can actually be understood.

Examples of these three frontiers.

1) Generalization of representations
The clearest example is CLIP, which learns a joint image-text embedding by aligning representations across modalities. Its learned geometry transfers remarkably to tasks it never saw — zero-shot classification, image retrieval, robotic perception — suggesting it captured something close to the intrinsic manifold of visual concepts rather than task-specific shortcuts. Studying why it transfers (low intrinsic dimensionality? smooth curvature? alignment with human semantic structure?) is an open and active question.

2) Reasoning structure in language models
Anthropic’s “Scaling and evaluating sparse autoencoders” work, along with follow-on mechanistic interpretability research, has found evidence that models trained purely on next-token prediction develop internal representations of entity states, spatial relations, and multi-step dependencies — structures that look suspiciously like world models. The cleaner controlled example is othello-GPT (Nanda et al.), where a transformer trained only on legal move sequences was shown to linearly represent the board state internally, a clean demonstration that reasoning-like geometric structure emerges without explicit supervision.

3) More interpretable representations
β-VAEs are the canonical attempt: penalizing the KL term forces the latent space toward an axis-aligned, disentangled geometry where individual dimensions correspond to independent generative factors. The result is representations where traversing a single latent direction changes exactly one attribute — pose, lighting, shape — leaving others fixed. The limitation is that disentanglement defined this way is coordinate-dependent and doesn’t guarantee causal structure, which has pushed more recent work toward causal representation learning (Schölkopf et al.) as the right geometric target.

Anthropic: Activations to Interpretable features with Monosemanticity

· Leave a comment ·

The Anthropic papers “Towards monosemanticity” and “Scaling Monosemanticity: Extracting Interpretable Features from Claude 3 Sonnet” demonstrate how sparse autoencoders can extract interpretable features from large language models, converting polysemantic neuron activations into monosemantic representations that directly map to identifiable concepts and behaviors.​ In this writeup I try to and explain the core concepts in this research.

A sparse autoencoder is a neural network designed to learn a compact, interpretablerepresentation of input data by enforcing sparsity on its hidden layer activations.  A sparse autoencoder is “sparse” because it applies a constraint during training so that, for any given input, only a small subset of the hidden (latent) units is active (nonzero). This is achieved by adding a sparsity penalty to the loss function, commonly L1 regularization or a KL-divergence term, which discourages most activations from deviating much from zero. This ensures the encoded representation is sparse—meaning only a few features are used to reconstruct the input—resulting in greater interpretability and the extraction of meaningful features.​ It is an “autoencoder” because the full model is trained end-to-end to reconstruct its own input. The encoder maps the input data to a latent code, and the decoder maps it back to the reconstruction. The central training objective is to minimize reconstruction error, making the network learn to reproduce its input as closely as possible. The difference from other autoencoder types (e.g., vanilla, denoising, variational) is specifically the addition of the sparsity constraint on the hidden code.

An activation is the output value of a neuron or unit in a neural network layer after applying an activation function to a weighted sum of inputs. Mathematically, for a neuron receiving inputs x1,x2,…,xnx1,x2,…,xn with weights w1,w2,…,wnw1,w2,…,wn, the activation is a=f(w1x1+w2x2+⋯+wnxn+b)a=f(w1x1+w2x2+⋯+wnxn+b), where ff is the activation function (such as ReLU, sigmoid, or tanh) and bb is a bias term.

The idea is to view activations as superpositions of underlying features and to use a neural network to reverse the mapping from the activations to the features. This is peering into the workings of an LLM with another neural network to see what the activations mean.

So in the monosemanticity quest, the activations are seen as a superposition of underlying features. A sparse autoencoder decomposes model activations into interpretable features by expressing each activation vector as a sparse linear combination of learned feature directions. Given an activation vector xjxj, the decomposition is:xj≈b+∑ifi(xj)dixj≈b+i∑fi(xj)di where fi(xj)fi(xj) is the activation (magnitude) of feature ii, didi is a unit vector representing the direction of feature ii in activation space, and bb is a bias term. The feature activations are computed by the encoder as fi(x)=ReLU(We(x−bd)+be)ifi(x)=ReLU(We(x−bd)+be)i, where WeWe is the encoder weight matrix and bdbd, bebe are pre-encoder and encoder biases. The feature directions are the columns of the decoder weight matrix WdWd. This formulation is dictionary learning: each activation is reconstructed from a sparse set of learned basis vectors scaled by their respective feature activations.

Acts is short for activations in the above figure of a sparse auto encoder functioning from Anthropic. .

Does the SAE look at all the activations or only certain layers ?

Sparse autoencoders are typically trained on activations from specific layers rather than all layers simultaneously. In practice, a separate SAE is trained for each layer or location in the model where one wishes to analyze or intervene on activations.​ In Anthropic’s “Scaling Monosemanticity” paper specifically, the SAE was trained only on activations from the residual stream at the middle layer (halfway through Claude 3 Sonnet). This choice was made for several reasons: the residual stream is smaller than the MLP layer, making training and inference computationally cheaper; focusing on the residual stream mitigates “cross-layer superposition,” which refers to neurons whose activations depend on combinations of information across multiple layers; and the middle layer likely contains more interesting and abstract features compared to early layers (which capture basic patterns) or final layers (which may be too task-specific).

Motivation and Definitions

  • Large language models (LLMs) typically exhibit polysemantic neurons, which activate in response to numerous, often unrelated, concepts, impeding interpretability and safe control.​
  • Monosemanticity refers to representations where each learned feature corresponds to a single, easily identifiable concept, thus improving transparency in model operations.​
  • Sparse autoencoders (SAEs) are employed to learn dictionary-like decompositions of hidden activations, aiming for each basis vector (feature) to align with a distinct semantic unit rather than mixed signals.​

Methods and Techniques

  • The approach uses SAEs to project model activations into higher-dimensional, sparse spaces where individual features become interpretable.
  • Dictionary learning is central: activations from a given layer are encoded by the SAE so that each dictionary element ideally corresponds to a unique concept or pattern.
  • Anthropic scales this method from small, shallow models to large networks by training SAEs on billions of activations from state-of-the-art LLMs (e.g., Claude 3 Sonnet).​
  • Modifying feature coefficients within the SAE’s learned space causes proportional, causal shifts in the model’s reconstructed activation, allowing direct steering of outputs at runtime.​
  • Feature steering leverages these interpretable directions to alter specific model behaviors (e.g., changing model goals, tone, biases, or inducing controlled errors) by adjusting activation values during inference.​

Results and Empirical Findings

  • The method yields dictionaries where a substantial portion of features (by human evaluation, approximately 70%) are monosemantic—associated with singular, nameable concepts such as DNA motifs or language script.​
  • Quantitative validation includes human raters agreeing with feature names, decoder-row alignment (cosine similarity > 0.86 between encoder and decoder vectors), and strong compositionality in steering outcomes.​
  • Scaling up the size of the SAE dictionary increases the proportion of monosemantic features and the precision of behavioral interventions.​
  • Interventions using these features show robust control over model outputs, evidenced by targeted behavioral scores and ability to suppress or augment specific behaviors with tunable steering coefficients.​

Conceptual Advances

  • The work empirically supports the superposition hypothesis: raw neurons entangle multiple meanings, but sparse dictionary learning untangles these into separately addressable features.​
  • The method demonstrates that high-dimensional, sparsely coded representations can be extracted at scale without significant algorithmic changes, opening new paths for mechanistic interpretability and control tools in LLMs.​
  • These advances suggest dictionary learning could, in future, replace large fine-tuning campaigns for behavioral adjustments, increase safety monitoring, and allow new forms of user-customized steering.​

Activation Steering and Implications

  • Steering methods operate by selecting, amplifying, or suppressing identified sparse features using signed, tunable coefficients (λλ), with each adjustment reflected directly and causally in output behavior.​
  • The process is mathematically tractable because the SAE remains linear; interventions can be analyzed for causal effects and compositional interactions, which is not feasible in the dense activation spaces of standard LLMs.​
  • This enables multifaceted interventions and targeted control: steering vectors can increase or decrease model propensities for specific behaviors, factuality, style, or compliance in a transparent manner.​

Summary Table: Key Terms

TermDefinition
Polysemantic neuronNeural unit that activates for multiple, unrelated concepts​
Monosemantic featureBasis vector representing a single interpretable concept​
Sparse autoencoderNeural model learning an overcomplete, interpretable dictionary​
Dictionary learningDecomposition of activations into a set of sparse, meaningful vectors
ActivationOutput value of a neuron or unit in a neural network layer after applying an activation function to a weighted sum of inputs
Activation steeringModifying activations using interpretable features to control outputs​

This research establishes scalable techniques for extracting and manipulating interpretable features in large LLMs, enabling precise behavioral steering and laying groundwork for safer, more controllable AI deployments.

The sparse autoencoder (SAE) in Anthropic’s “Scaling Monosemanticity” paper was trained at three different scales on activations from Claude 3 Sonnet: approximately 1 million (1,048,576), 4 million (4,194,304), and 34 million (33,554,432) features. For the largest run, the 34M-feature SAE, the number of active (nonzero) features for any given token was typically fewer than 300, showing high sparsity.​

The paper emphasizes that many extracted features are relevant to AI safety, such as features for security vulnerabilities, code backdoors, bias (overt and subtle), deception (including power-seeking and treacherous turns), sycophancy, and the generation of dangerous or criminal content. However, the authors note that the detection of such features is preliminary and should not be over-interpreted: knowing about harmful behaviors is distinct from enacting them. The presence of potentially dangerous features suggests the model could represent these concepts internally, warranting deeper investigation. The interpretability gained through the SAE allows for the identification and possible intervention on such features but does not automatically ensure safe model behavior without further work and robust evaluation.​

The authors compare their feature-extraction approach to previous interpretability and model-steering methods:

  • Unlike neuron-centric methods, which often yield tangled, polysemantic activations, SAEs learn overcomplete, sparse dictionaries that approximate monosemantic features.
  • Their approach leverages scaling laws to optimize both the number of features and training steps, showing that larger SAEs provide more granular, precise, and interpretable decompositions than smaller or denser models.
  • The SAE-based approach allows for explicit, steerable interventions by clamping or zeroing specific features, something not possible with conventional dense neuron manipulation.
  • The paper positions this technique as extensible, mechanistically transparent, and a foundation for scalable model interpretability—offering capabilities not matched by most prior strategies.​

These results highlight that scalable, sparse autoencoders produce directly actionable, interpretable features offering new tools for AI safety and more precise model control compared to traditional neuron or layerwise interpretability approaches.​

An argument on the urgency of interpretability: https://www.darioamodei.com/post/the-urgency-of-interpretability

Neel Nanda’s replication of results has a notebook for going deeper. https://www.alignmentforum.org/posts/fKuugaxt2XLTkASkk/open-source-replication-and-commentary-on-anthropic-s

Absolute Zero: zero reliance on external data to improve model reasoning

· Leave a comment ·

Imagine you want to train a large language model to get really good at solving tough problems—things like math puzzles or writing correct code. Usually, the way people do this is by giving the model lots of practice questions written by humans. These are called human-curated tasks: real people come up with the problems and answers, like “Write a program to reverse a string” or “What’s the derivative of x²?”. The model practices on these problem-solution pairs, and then reinforcement learning (RL) or reinforcement learning with verifiable rewards (RLVR) can be used to improve how it reasons.

But as models get bigger and smarter, collecting enough high-quality problems from humans becomes expensive, slow, and limiting. If the model might one day surpass most humans, why should humans be the bottleneck?

That’s where this paper’s idea, called Absolute Zero, comes in. Instead of relying on people to write problems, the model creates its own. One part of the model plays the “teacher,” proposing new tasks, and another part plays the “student,” trying to solve them. Because the environment is code, the answers can be automatically checked just by running the program—so no human needs to grade them.

The model learns three kinds of reasoning:

  • Deduction: given a program and input, figure out the output.
  • Abduction: given a program and an output, figure out the input.
  • Induction: given some examples, figure out the program that works in general.

The system rewards the student for solving problems correctly, and the teacher for coming up with problems that are just the right difficulty—not too easy, not impossible.

The result is that training only on these self-made coding tasks made the model better at math. On standard benchmarks, it matched or even beat other models that were trained with large sets of human-written problems. Bigger models improved even more, and “coder” models (already good at programming) saw the biggest gains. The model even started showing “scratch-pad” style reasoning on its own, writing little notes or plans before coding—without being told to.

In short, the key insight is this: you don’t necessarily need humans to write all the practice problems anymore. If you have a way to automatically check answers, a model can bootstrap itself, creating and solving its own challenges, and still learn to reason across domains.

The authors do warn that there are challenges—like making sure tasks stay diverse, keeping the system safe, and managing the heavy compute costs—but the big takeaway is that self-play with verifiable rewards could be a new path to building smarter, more independent reasoning systems.

There’s no “exam” in the usual sense for the students – the system builds a feedback loop between the teacher (proposer) and the student (solver).

Here’s how it works step by step:

1. Teacher proposes a task

The proposer (teacher model) generates a new program + input/output pair (a problem).

Example: “Write a function that finds prime numbers up to N.”

2. Environment checks validity

The environment (code runner) ensures the task is valid: it runs, is safe, deterministic, etc.

If valid, it gets stored in a task buffer.

3. Student attempts the task

The solver (student model) pulls the task and tries to solve it.

The environment executes the student’s answer and checks correctness.

4. Rewards reflect difficulty

If the student always solves a task → it’s too easy → proposer gets low reward.

If the student never solves a task → it’s too hard → proposer also gets low reward.

If the student solves it sometimes → it’s “learnable” → proposer gets high reward.

So the proposer doesn’t “know” in advance how good the student is. Instead, it learns over time:

Tasks that end up being useful for training (medium difficulty) get reinforced.

Tasks that are too trivial or impossible fade out because they bring no proposer reward.

The proposer is like a coach who experiments with new drills, and the student’s performance on them acts as the exam. Over time, the teacher learns what kinds of problems best stretch the student without breaking them.

RDMA, Infiniband, RoCE, CXL : High-Performance Networking Technologies for AI

· Leave a comment ·

As the demand for high-performance computing (HPC) and artificial intelligence (AI) continues to grow, networking technologies have become critical to ensuring the scalability and efficiency of modern data centers. Among these, RDMA, InfiniBand, RoCE, and the emerging CXL standard stand out as transformative technologies, each addressing unique challenges. Here’s a brief overview of these key technologies, trends, and future.

Remote Direct Memory Access (RDMA) was developed in response to the increasing need for low-latency, high-bandwidth data movement in distributed computing environments. RDMA was driven by a collaboration of major tech companies to address the limitations of traditional networking models. Some key players in RDMA’s early development include:

  • Compaq, IBM, and Intel:
    • Developed the initial RDMA architecture to improve networking efficiency, particularly in storage and high-performance computing.
  • Mellanox Technologies:
    • One of the first companies to commercialize RDMA with its InfiniBand solutions, allowing ultra-low latency communication.
  • Microsoft & Networking Industry:
    • Developed iWARP (RDMA over TCP/IP) to integrate RDMA into Ethernet-based networks.
  • InfiniBand Trade Association (IBTA):
    • Founded in 1999 by Compaq, Dell, Hewlett-Packard, IBM, Intel, Microsoft, and Sun Microsystems to standardize high-performance networking, including RDMA capabilities.

Before RDMA, networking relied on CPU-intensive packet processing, which created performance bottlenecks in data-intensive applications. The traditional TCP/IP stack required multiple CPU interrupts, context switches, and memory copies, leading to high latency and inefficiency.

RDMA Was Developed to Solve These Challenges:

  1. Eliminate CPU Bottlenecks:
    • Traditional networking required CPU cycles for data movement, slowing down high-speed applications.
    • RDMA bypasses the OS kernel and CPU, reducing overhead.
  2. Enable High-Speed, Low-Latency Communication:
    • Needed for HPC (High-Performance Computing), AI training, and databases.
    • Reduces communication latency to below 1 microsecond.
  3. Improve Scalability for Distributed Systems:
    • Large-scale data centers and supercomputers require fast inter-node communication.
    • RDMA enables efficient parallel computing across thousands of nodes.
  4. Optimize Storage and Networking:
    • Technologies like NVMe over Fabrics (NVMe-oF) use RDMA for ultra-fast storage access.
    • RDMA dramatically speeds up databases and cloud storage, reducing I/O latency.

Evolution and Implementations of RDMA

RDMA has evolved into different implementations, each suited for different networking environments:

RDMA VariantTransport ProtocolUse Case
InfiniBandNative InfiniBand transportHPC, AI training, supercomputing
RoCE (RDMA over Converged Ethernet)Ethernet (Layer 2/3)Cloud data centers, AI inference
iWARPTCP/IPEnterprise storage, cloud computing

RDMA’s Impact on Modern Computing

Today, RDMA is a core technology in AI, cloud computing, and high-speed storage. It enables:

  • Massive parallelism in AI training (e.g., NVIDIA DGX, GPT models).
  • Faster database transactions (e.g., Microsoft SQL Server, Oracle).
  • Low-latency cloud networking (used by Azure, AWS, Google Cloud).

InfiniBand: InfiniBand is a high-performance networking technology designed for low-latency, high-bandwidth communication. Primarily used in HPC and AI training clusters, InfiniBand supports features like Remote Direct Memory Access (RDMA), enabling direct memory-to-memory data transfers with minimal CPU involvement. Its scalable architecture makes it ideal for distributed workloads, offering latencies as low as 0.5 microseconds and bandwidths up to 400 Gbps (NDR).

RDMA over Converged Ethernet (RoCE): RoCE extends RDMA capabilities over Ethernet networks, bridging the gap between the performance of InfiniBand and the ubiquity of Ethernet. By leveraging standard Ethernet infrastructure with lossless configurations, RoCE delivers efficient communication for data centers that prioritize compatibility and cost. However, it typically exhibits slightly higher latencies (5-10 microseconds) compared to InfiniBand.

Compute Express Link (CXL): CXL is a new interconnect standard designed to provide low-latency, high-bandwidth communication between processors, accelerators, and memory devices within a single node. By leveraging PCIe infrastructure, CXL supports memory pooling, coherent data sharing, and dynamic resource allocation, addressing the growing complexity of heterogeneous compute environments

Key Technology Trends
  1. AI Training Driving High-Bandwidth Demand:
    • Training large-scale AI models requires massive data exchange between GPUs, CPUs, and memory. InfiniBand remains the leader in this domain due to its ultra-low latency and scalability, but RoCE is increasingly adopted in cost-sensitive deployments.
  2. Distributed Inference and Edge AI:
    • While inference typically has lower communication demands, distributed inference pipelines and edge AI are pushing for efficient interconnects. RoCE’s compatibility with Ethernet makes it a strong candidate in these scenarios.
  3. Memory-Centric Architectures:
    • With CXL’s focus on memory pooling and coherent memory sharing, the future of data centers may see significant convergence around flexible, node-level resource allocation. This complements, rather than competes with, network-level technologies like InfiniBand and RoCE.
  4. Interconnect Ecosystem Integration:
    • NVIDIA’s integration of InfiniBand with its GPUs and DPUs highlights the trend of tightly coupled compute and networking stacks. Similarly, innovations in RoCE and Ethernet SmartNICs are bringing RDMA capabilities closer to mainstream data centers.
Extrapolating to the future
  • Convergence of Standards: As workloads diversify, data centers may adopt hybrid approaches, combining InfiniBand for training clusters, RoCE for distributed inference, and CXL for intra-node memory coherence. Seamless interoperability between these standards will be ideal.
  • AI-Centric Network Evolution: The growing dominance of AI workloads will push networking technologies toward even lower latencies and higher bandwidths, with InfiniBand and RoCE leading the charge.
  • Rise of Heterogeneous Compute: CXL’s potential to unify memory access across CPUs, GPUs, and accelerators aligns with the industry’s shift toward heterogeneous compute, enabling efficient resource utilization and scalability.
  • Cloud-Driven Innovations: As hyperscalers like AWS, Google, and Azure integrate these technologies into their offerings, cost-efficient, scalable solutions like RoCE and CXL may become more widespread, complementing specialized InfiniBand deployments.

vLLM project – overview, comparisons, PagedAttention mechanism

· Leave a comment ·

The vLLM project is an open-source venture designed to enhance the efficiency and scalability of serving Large Language Models (LLMs). Developed by researchers at UC Berkeley, vLLM aims to improve the performance of LLM inference by optimizing memory management and execution. It offers a system that reduces latency and increases throughput for LLMs, making it a valuable tool for deploying these models more effectively in various applications. It supports multiple LLM model types, multiple hardware architectures, and multiple optimization techniques. It is described in this paper, on Efficient LLM serving with PagedAttention.

vLLM achieves its improvements through

  • dynamic batching,
  • efficient memory usage, and
  • parallel execution strategies.

These features allow it to handle multiple requests simultaneously without sacrificing speed or accuracy.

By making LLMs more accessible and efficient, vLLM helps lower the barriers to using advanced AI models, facilitating broader adoption and innovation in the field of natural language processing. For more detailed information or to contribute to the project, you can explore its repository on platforms like GitHub.

vLLM, NVIDIA Triton Inference Server, and NVIDIA NeMo (formerly known as NVIDIA NIM) are all designed to improve the deployment and performance of machine learning models, but they have different focuses and functionalities. Here’s a comparison of each:

vLLM
  • Purpose: Optimizes the serving of Large Language Models (LLMs) with a focus on improving inference efficiency, particularly regarding memory management and execution.
  • Features: Offers dynamic batching, efficient memory usage, and parallel execution strategies specifically for LLMs, enhancing latency and throughput.
  • Use Cases: Best suited for applications requiring fast, efficient LLM inference, such as AI-driven conversational agents.
  • How it reduces memory waste and improves utilization with PagedAttention – https://blog.runpod.io/introduction-to-vllm-and-how-to-run-vllm-on-runpod-serverless/
NVIDIA Triton Inference Server
  • Purpose: A scalable and flexible platform for serving different types of machine learning models across a variety of frameworks and hardware architectures.
  • Features: Supports multiple model frameworks (e.g., TensorFlow, PyTorch, ONNX), dynamic batching, model versioning, and provides both HTTP/REST and gRPC endpoints for inference requests. It is designed to maximize GPU utilization and streamline inference workflows.
  • Use Cases: Ideal for deploying diverse AI models in production environments, allowing for efficient inference at scale across CPUs and GPUs.
NVIDIA NeMo
  • Purpose: A toolkit for building, training, and fine-tuning state-of-the-art conversational AI models, including those for Automatic Speech Recognition (ASR), Natural Language Processing (NLP), and Text-to-Speech (TTS).
  • Features: Provides pre-trained models, model architectures, and training scripts that can be customized and extended for specific tasks. NeMo is designed to facilitate the development of AI models with high accuracy and efficiency.
  • Use Cases: Suitable for developers and researchers focused on building and customizing conversational AI applications, offering extensive support for research and development in speech and language domains.

Comparison summary

  • Optimization Focus: vLLM is specialized for LLM inference optimization, NVIDIA Triton is a general-purpose inference server supporting various models and frameworks, and NVIDIA NeMo is focused on developing and customizing conversational AI models.
  • Hardware and Framework Support: Triton supports a wide range of frameworks and hardware, optimizing inference across diverse environments. NeMo, while capable of leveraging NVIDIA’s hardware optimizations, is more focused on the model training and customization aspect, particularly for conversational AI.
  • Target Audience: vLLM targets developers needing efficient LLM deployment; Triton appeals to teams deploying a variety of models in scalable production settings; NeMo is aimed at researchers and developers building state-of-the-art conversational systems.
Details of vLLM PagedAttention.

What Are Keys and Values in PagedAttention?

In the context of transformer-based Large Language Models (LLMs), keys (K) and values (V) are components of the attention mechanism used during inference.

  • Keys (K): Represent encoded representations of previous tokens, used to determine how much attention each token should pay to previous tokens.
  • Values (V): Contain the actual information used to generate the next token, weighted based on attention scores.

PagedAttention manages these key-value (KV) caches efficiently to store past token embeddings so the model doesn’t have to recompute them in every step, drastically speeding up inference.

Concrete Example: Key-Value Pairs in Action

Let’s take a simple example where an LLM is generating text based on a prompt.

Example Prompt:

User: "The capital of France is"

Tokenized Version (Using Byte-Pair Encoding or SentencePiece):

["The", "capital", "of", "France", "is"]

Each token gets embedded into a high-dimensional space (e.g., 4096 dimensions for LLaMA-2-70B). Let’s assume we use 4096-dimension embeddings for simplicity.

Step-by-Step Key-Value Storage

  1. The model encodes each token and stores:
    • Key (K): A vector that helps determine how relevant this token is in future attention computations.
    • Value (V): The actual contextual representation of the token.
TokenKey (K) (Simplified)Value (V) (Simplified)
“The”[0.1, 0.2, -0.3, ...][0.5, 0.4, -0.1, ...]
“capital”[0.2, 0.3, 0.1, ...][0.6, 0.2, -0.3, ...]
“of”[-0.1, 0.2, 0.7, ...][0.2, 0.1, 0.9, ...]
“France”[0.5, -0.2, 0.1, ...][0.7, 0.3, -0.2, ...]
“is”[0.3, 0.1, 0.4, ...][0.8, 0.2, -0.5, ...]
  1. When generating the next token (“Paris”), the model:
    • Computes attention scores between “Paris” and all previous tokens using dot product of queries (Q) and keys (K).
    • Uses the weighted sum of values (V) to form the new representation.
  2. Instead of recomputing attention from scratch, PagedAttention retrieves precomputed (K, V) values from memory pages for fast lookup.

How PagedAttention Optimizes Key-Value Caching

  • Without PagedAttention: Each request would store KV pairs in one long, contiguous memory buffer. If a request finishes early, the allocated space is wasted.
  • With PagedAttention: KV pairs are stored in small pages (e.g., chunks of 16 tokens), allowing efficient reuse and minimizing fragmentation.

AI Risks Repository from MIT

· Leave a comment ·

On the topic of governance of AI, here’s a comprehensive listing of AI Risks from MIT with over 700 risks in 7 domains, and extracted from 43 existing frameworks.

https://www.csail.mit.edu/news/global-ai-adoption-outpacing-risk-understanding-warns-mit-csail

https://airisk.mit.edu/

https://sloanreview.mit.edu/article/ai-related-risks-test-the-limits-of-organizational-risk-management/

Statement: Organizations are sufficiently expanding risk management capabilities to address AI-related risks.

Sizing an LLM for GPU memory

· Leave a comment ·

When choosing the EC2 instance for a Large Language Model, one of the first constraints is whether the model will fit in the GPU memory of an instance.

Given a choice of a model, the decisions roughly follow this path –

Model -> Training/Inferencing -> Technique (choice of optimization) -> Memory requirement -> Instance requirement -> Instance availability -> smaller instance or more optimization or distributed training.

Some extreme optimizations are possible such as QLora for Inferencing . See the blog How to fit a layer in memory at a time https://huggingface.co/blog/lyogavin/airllm . However many use cases do not want any sacrifices in accuracy.

Distributed training by splitting the model against smaller instances is another possibility. A discussion is here – https://siboehm.com/articles/22/pipeline-parallel-training

Here’s a listing of different GPU instance types with a column for GPU Memory (GiB) on one page to facilitate instance comparisons.

EC2 G3 Instance Details
NameGPUsvCPUMemory (GiB)GPU Memory (GiB)Price/hr* (Linux)Price/hr* (Windows)1-yr Reserved Instance Effective Hourly* (Linux)3-yr Reserved Instance Effective Hourly* (Linux)
g3s.xlarge1430.58$0.75
$0.93
$0.525$0.405
g3.4xlarge1161228$1.14$1.876$0.741$0.538
g3.8xlarge23224416$2.28$3.752$1.482$1.076
g3.16xlarge46448832$4.56$7.504$2.964$2.152
EC2 G4 Instance details
 Instance SizeGPUvCPUsMemory (GiB)Instance Storage (GB)Network Bandwidth (Gbps)EBS Bandwidth (Gbps)On-Demand Price/hr*1-yr Reserved Instance Effective Hourly* (Linux)3-yr Reserved Instance Effective Hourly* (Linux)

G4dn

Single GPU VMsg4dn.xlarge14161 x 125 NVMe SSDUp to 25Up to 3.5$0.526$0.316$0.210
g4dn.2xlarge18321 x 225 NVMe SSDUp to 25Up to 3.5$0.752$0.452$0.300
g4dn.4xlarge116641 x 225 NVMe SSDUp to 254.75$1.204$0.722$0.482
g4dn.8xlarge1321281 x 900 NVMe SSD509.5$2.176$1.306$0.870
g4dn.16xlarge1642561 x 900 NVMe SSD509.5$4.352$2.612$1.740
           
Multi GPU VMsg4dn.12xlarge4481921 x 900 NVMe SSD509.5$3.912$2.348$1.564
g4dn.metal8963842 x 900 NVMe SSD10019$7.824$4.694$3.130

G4ad

Single GPU VMsg4ad.xlarge14161 x 150 NVMe SSDUp to 10Up to 3$0.379$0.227$0.178
g4ad.2xlarge18321 x 300 NVMe SSDUp to 10Up to 3$0.541$0.325$0.254
g4ad.4xlarge116641 x 600 NVMe SSDUp to 10Up to 3$0.867$0.520$0.405
           
Multi GPU VMsg4ad.8xlarge2321281 x 1200 NVMe SSD153$1.734$1.040$0.810
g4ad.16xlarge4642561 x 2400 NVMe SSD256$3.468$2.081$1.619
EC2 G5 instance details
 Instance SizeGPUGPU Memory (GiB)vCPUsMemory (GiB)Storage (GB)Network Bandwidth (Gbps)EBS Bandwidth (Gbps)On Demand Price/hr*1-yr ISP Effective Hourly (Linux)3-yr ISP Effective Hourly (Linux)
Single GPU VMsg5.xlarge1244161×250Up to 10Up to 3.5$1.006$0.604$0.402
g5.2xlarge1248321×450Up to 10Up to 3.5$1.212$0.727$0.485
g5.4xlarge12416641×600Up to 258$1.624$0.974$0.650
g5.8xlarge124321281×9002516$2.448$1.469$0.979
g5.16xlarge124642561×19002516$4.096$2.458$1.638
            
Multi GPU VMsg5.12xlarge496481921×38004016$5.672$3.403$2.269
g5.24xlarge496963841×38005019$8.144$4.886$3.258
g5.48xlarge81921927682×380010019$16.288$9.773$6.515
EC2 G6 instance details
 Instance SizeGPUGPU Memory (GB)vCPUsMemory (GiB)Storage (GB)Network Bandwidth (Gbps)EBS Bandwidth (Gbps)On Demand Price/hr*1-yr ISP Effective Hourly (Linux)3-yr ISP Effective Hourly (Linux)
Single GPU VMs          g6.xlarge1244161×250Up to 10Up to 5$0.805$0.499$0.342
g6.2xlarge1248321×450Up to 10Up to 5$0.978$0.606$0.416
g6.4xlarge12416641×600Up to 258$1.323$0.820$0.562
g6.8xlarge124321282×4502516$2.014$1.249$0.856
g6.16xlarge124642562×9402520$3.397$2.106$1.443
Gr6 instances with 1:8 vCPU:RAM ratio
gr6.4xlarge124161281×600Up to 258$1.539$0.954$0.654
gr6.8xlarge124322562×4502516$2.446$1.517$1.040
            
Multi GPU VMsg6.12xlarge496481924×9404020$4.602$2.853$1.955
g6.24xlarge496963844×9405030$6.675$4.139$2.837
g6.48xlarge81921927688×94010060$13.35$8.277$5.674
EC2 G6e instances
Instance SizeGPUGPU Memory (GiB)  vCPUsMemory(GiB)Storage (GB)  Network Bandwidth (Gbps)  EBS Bandwidth (Gbps)
g6e.xlarge148432250Up to 20Up to 5
g6e.2xlarge148864450Up to 20Up to 5
g6e.4xlarge14816128600208
g6e.8xlarge148322569002516
g6e.16xlarge1486451219003520
g6e.12xlarge419248384380010020
g6e.24xlarge419296768380020030
g6e.48xlarge83841921536760040060
EC2 P3 instance details
Instance SizeGPUs – Tesla V100GPU Peer to PeerGPU Memory (GB)vCPUsMemory (GB)Network BandwidthEBS BandwidthOn-Demand Price/hr*1-yr Reserved Instance Effective Hourly*3-yr Reserved Instance Effective Hourly*
p3.2xlarge1N/A16861Up to 10 Gbps1.5 Gbps$3.06$1.99$1.05
p3.8xlarge4
NVLink643224410 Gbps7 Gbps$12.24$7.96$4.19
p3.16xlarge8NVLink1286448825 Gbps14 Gbps$24.48$15.91$8.39
p3dn.24xlarge8NVLink25696768100 Gbps19 Gbps$31.218$18.30$9.64
EC2 P4 instance details
Instance SizevCPUsInstance Memory (GiB)GPU – A100GPU memoryNetwork Bandwidth (Gbps)GPUDirect RDMAGPU Peer to PeerInstance Storage (GB)EBS Bandwidth (Gbps)On-demand Price/hr1-yr Reserved Instance Effective Hourly *3-yr Reserved Instance Effective Hourly *
p4d.24xlarge9611528320 GB
HBM2		400 ENA and EFAYes600 GB/s NVSwitch8 x 1000 NVMe SSD19$32.77$19.22$11.57
p4de.24xlarge (preview)9611528640 GB
HBM2e		400 ENA and EFAYes600 GB/s NVSwitch8 x 1000 NVMe SSD19$40.96$24.01$14.46
EC2 P5 instance details
Instance SizevCPUInstance Memory (TiB)GPU – H100GPU  MemoryNetwork BandwidthGPUDirectRDMAGPU Peer to PeerInstance Storage (TB)EBS Bandwidth (Gbps)
p5.48xlarge1922 8640 GB HBM33200 Gbps EFAv2Yes900 GB/s NVSwitch8 x 3.84 NVMe SSD80 
EC2 P5e instance details
Instance SizevCPUsInstance Memory (TiB)GPUGPU memoryNetwork Bandwidth (Gbps)GPUDirect RDMAGPU Peer to PeerInstance Storage (TB)EBS Bandwidth (Gbps)
p5e.48xlarge19228 x NVIDIA H2001128 GB
HBM3e		3200 Gbps EFAYes900 GB/s NVSwitch8 x 3.84 NVMe SSD80

Relevant links

P5e and P5en announcement (update Sep’24). https://aws.amazon.com/blogs/machine-learning/amazon-ec2-p5e-instances-are-generally-available/

https://en.wikipedia.org/wiki/List_of_Nvidia_graphics_processing_units

Use of Triton and NIM to make use of GPU memory across multiple GPUs on an instance –

https://github.com/aws-samples/amazon-eks-machine-learning-with-terraform-and-kubeflow

https://aws.amazon.com/blogs/hpc/deploying-generative-ai-applications-with-nvidia-nims-on-amazon-eks

FP4 and four bit integer quantization, and QLoRA

Making LLMs even more accessible with bitsandbytes, 4-bit quantization and QLoRA at https://huggingface.co/blog/4bit-transformers-bitsandbytes

[2305.14152] Memory-Efficient Fine-Tuning of Compressed Large Language Models via sub-4-bit Integer Quantization

Note: Performance is not just about GPU memory but also network bandwidth which is needed to load the large models especially for a platform serving multiple models.

When comparing the importance of high memory bandwidth between training and inference for Large Language Models (LLMs), it is generally more critical for training. Here’s why:

1. Training LLMs

  • Data Movement: Training LLMs involves frequent data movement between the GPU memory and the processing units. Each training iteration requires loading large batches of data, performing extensive matrix multiplications, and updating weights, all of which are memory-intensive operations.
  • Backward Pass: During the training phase, the backward pass (gradient computation and backpropagation) is highly memory bandwidth-intensive. The gradients of each layer are computed and propagated back through the network, requiring significant memory access.
  • Parameter Updates: High memory bandwidth is essential to handle the large volume of data being read and written during the parameter updates across multiple layers, especially in very deep models.
  • Larger Models and Datasets: Training large models like GPT-3 or GPT-4 involves massive datasets and millions (or even billions) of parameters, leading to a substantial demand for memory bandwidth.

2. Inferencing of LLMs:

  • Data Movement: During inference, the primary task is to process input data and generate outputs, which involves reading the model parameters and performing computations. While this still requires good memory bandwidth, the demands are generally lower compared to training.
  • No Backpropagation: Inference does not involve the backward pass or parameter updates, significantly reducing the need for continuous memory writes. The absence of gradient computations and updates reduces the overall memory bandwidth requirements.
  • Smaller Batch Sizes: Inference typically operates on smaller batch sizes compared to training, further reducing the demand for memory bandwidth.
  • Optimizations: Techniques such as model quantization and optimized inference runtimes (like TensorRT) can reduce the memory bandwidth required during inference by optimizing how data is accessed and processed.

SageMaker Hyperpod for Distributed Model Training

· Leave a comment ·

Amazon SageMaker HyperPod is a new infrastructure designed specifically for distributed training at scale. It offers a purpose-built, high-performance environment that accelerates the training of large machine learning models by optimizing resource allocation, reducing communication overhead, and providing seamless scaling. HyperPod integrates with SageMaker to simplify complex training workflows, making it easier for users to efficiently train foundation models and other large-scale ML workloads. This innovation supports faster iteration and development of AI models. https://aws.amazon.com/sagemaker/hyperpod , https://aws.amazon.com/blogs/machine-learning/introducing-amazon-sagemaker-hyperpod-to-train-foundation-models-at-scale

Perplexity, a generative AI startup, improved its model training speed by 40% using Amazon SageMaker HyperPod on AWS. By leveraging advanced distributed training capabilities and EC2 instances, Perplexity optimized its model training and inference processes. This allowed the company to efficiently handle over 100,000 queries per hour with low latency and high throughput, enhancing user experiences and enabling rapid iteration in AI development. https://aws.amazon.com/solutions/case-studies/perplexity-case-study

https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod-cluster-observability.html

LLM optimization – PEFT, LORA, QLORA

· Leave a comment ·

Large Language Models (LLMs) have transformed natural language processing, but their immense size and computational demands pose significant challenges. Optimizing these models is crucial for efficient deployment, particularly in resource-constrained environments. Below, we explore several optimization techniques, including Parameter-Efficient Fine-Tuning (PEFT), Low-Rank Adaptation (LoRA), and Quantized Low-Rank Adaptation (QLoRA), highlighting their unique benefits and differences.

1. Parameter-Efficient Fine-Tuning (PEFT)

PEFT is designed to reduce the computational burden of fine-tuning large models by updating only a small subset of the model’s parameters, rather than the entire model. This approach allows for significant resource savings while maintaining performance, making it particularly useful for adapting LLMs to new tasks with limited data or compute resources.

Key Features:

  • Selective Parameter Update: Only a fraction of the model’s parameters are fine-tuned.
  • Efficiency: Reduces the computational cost and memory footprint during fine-tuning.
  • Flexibility: Can be applied across various LLM architectures.
2. Low-Rank Adaptation (LoRA)

LoRA is a technique that further reduces the number of parameters to be updated during fine-tuning by decomposing the model’s weight matrices into low-rank components. By introducing low-rank matrices that are trained alongside the existing weights, LoRA enables fine-tuning with minimal additional parameters, preserving the original model’s architecture.

Key Features:

  • Low-Rank Decomposition: Decomposes weights into low-rank matrices to minimize parameter updates.
  • Minimal Overhead: Adds only a small number of trainable parameters.
  • Performance: Maintains or even enhances model performance on specific tasks.
3. Quantized Low-Rank Adaptation (QLoRA)

QLoRA combines quantization and LoRA to maximize memory and computational efficiency. By quantizing the low-rank matrices, QLoRA reduces the precision of these components, allowing for even greater reductions in memory usage and computational costs without a significant loss in accuracy.

Key Features:

  • Quantization: Reduces precision of low-rank matrices to lower memory usage.
  • Memory Efficiency: Significantly decreases the memory required for fine-tuning.
  • Scalability: Ideal for large-scale deployments where memory is a critical concern.
Contrasting PEFT, LoRA, and QLoRA
  • Parameter Update Strategy:
    • PEFT: Updates a small subset of existing parameters.
    • LoRA: Introduces additional low-rank matrices for parameter updates.
    • QLoRA: Combines low-rank matrices with quantization for extreme memory efficiency.
  • Memory and Computational Requirements:
    • PEFT: Reduces overall fine-tuning costs but may still require substantial memory.
    • LoRA: Further reduces memory usage by minimizing the number of updated parameters.
    • QLoRA: Offers the most memory efficiency by applying quantization to the low-rank matrices.
  • Application Scenarios:
    • PEFT: Suitable for fine-tuning in environments with limited compute resources.
    • LoRA: Ideal for scenarios requiring efficient fine-tuning with minimal parameter overhead.
    • QLoRA: Best for large-scale deployments where memory efficiency is paramount.

Direct Preference Optimization (DPO) vs RLHF/PPO (Reinforcement Learning with Human Feedback, Proximal Policy Optimization)

· Leave a comment ·

The paper “Direct Preference Optimization: Your Language Model is Secretly a Reward Model” introduces Direct Preference Optimization (DPO), an algorithm for fine-tuning language models to align with human preferences without the need for complex reinforcement learning procedures. This simplifies Reinforcement Learning with Human Feedback (RLHF) by not requiring a time consuming human feedback loop in training of the model.

Directly Modified Reward Function : DPO uses human preferences to directly modify the reward function, employing a classification loss to align the model outputs with these preferences. Rather than relying solely on reward signals from the environment, it leverages comparisons or preferences between different trajectories to guide the learning process. The agent is provided with pairs of trajectories along with a preference indicating which trajectory is preferred. This preference data is used to train the policy. The task of predicting preferences can be framed as a binary classification problem. For a given pair of trajectories the model needs to predict which path is preferred. The classification loss then measures the discrepancy between the predicted and actual preferences. A common choice for this kind of binary classification is the binary cross-entropy loss. The overall training objective in DPO involves minimizing the classification loss across all pairs of trajectories in the dataset, which encourages the policy to produce trajectories that align with the observed preferences.

RLHF and Proximal Policy Optimization: RLHF trains a reward model using PPO and data gathered on human preferences that is labeled by humans. These RLHF steps are shown in the diagram below, from the RLHF paper. PPO indirectly learns the reward function through interactions with the environment and optimizes the policy to maximize this reward, using a reinforcement learning framework. The policy here is a mapping from states to a probability distribution over actions.

So Direct Preference Optimization (DPO) modifies the reward function using human preference data. Here is a high-level overview of the equations used:

  1. Preference Model:
    • Let θ be the parameters of the model.
    • Let τ1​ and τ2​ be two trajectories (or outputs) being compared.
    • The preference model P(τ1≻τ2∣θ)  indicates the probability that humans prefer τ1​ over τ2​.
  2. Logistic Function for Preferences:
    • The preference probability is modeled using a logistic function:P(τ1≻τ2∣θ)=exp⁡(R(τ1∣θ)) / ( exp⁡(R(τ1∣θ)) + exp⁡(R(τ2∣θ)) )
    • R(τ∣θ) is the reward function for trajectory τ.
  3. Loss Function:
    • The loss function L(θ) is defined as the negative log-likelihood of the human preferences:L(θ)=−∑(τ1,τ2)∈D log⁡ P(τ1≻τ2∣θ)
    • D is the dataset of human preference comparisons.
  4. Optimization:
    • The model parameters θ are optimized by minimizing the loss function L(θ)

GPU kernel functions for deep learning

· Leave a comment ·

This article attempts to outline GPU Kernel Functions and how they are supported in TensorFlow, PyTorch, and OpenAI Triton. GPU Kernel Functions are specialized functions executed on an Nvidia Graphics Processing Unit. These functions play a key role in for parallel and accelerated computing such as tensor matrix operations used in deep learning.

GPU kernel functions for operations commonly used in deep learning include:

  1. Element-wise operations: TensorFlow provides GPU kernels for element-wise operations such as addition, subtraction, multiplication, and division, enabling efficient computation on arrays or tensors.
  2. Matrix operations: GPU kernels in TensorFlow optimize matrix operations like matrix multiplication, matrix addition, and matrix transpose, which are fundamental in many deep learning models.
  3. Convolutional operations: TensorFlow implements GPU kernels for convolutional operations, which are essential for tasks like image recognition and computer vision.
  4. Reduction operations: TensorFlow provides GPU kernels for reduction operations like summation, mean, maximum, and minimum, allowing efficient computation over large arrays or tensors.
  5. Activation functions: GPU kernels are implemented for common activation functions used in deep learning, such as ReLU (Rectified Linear Unit), sigmoid, and tanh.
  6. Pooling operations: TensorFlow’s GPU kernels optimize pooling operations like max pooling and average pooling, commonly used in convolutional neural networks (CNNs).
  7. Recurrent operations: TensorFlow provides GPU kernels for recurrent operations like LSTM (Long Short-Term Memory) and GRU (Gated Recurrent Unit), which are widely used in sequence-based models.

TensorFlow optimizes the execution of operations within a computation graph. When operations can be executed on a GPU, TensorFlow translates the high-level operations into CUDA calls that invoke the corresponding GPU kernels.

PyTorch is another popular open-source deep learning framework that provides a high-level programming interface for building and training machine learning models.

PyTorch differs from TensorFlow in a few ways:

  1. Dynamic Computational Graph: PyTorch uses a dynamic computational graph approach, whereas TensorFlow uses a static computational graph. This means that in PyTorch, the computational graph is constructed and executed on the fly as the code is executed, allowing for more flexibility and dynamic behavior during model training and inference.
  2. Imperative Programming: PyTorch follows an imperative programming style, which allows users to write code that is more intuitive and resembles standard Python programming. This makes it easier to understand and debug the code, as well as experiment with different model architectures and algorithms.
  3. Autograd: PyTorch’s autograd system allows automatic differentiation, which enables computing gradients for model parameters. This makes it easier to implement and train complex models, as users don’t have to manually compute gradients. TensorFlow, on the other hand, uses a static graph approach where gradients need to be explicitly defined and computed.
  4. TorchScript: PyTorch provides a feature called TorchScript, which allows models to be serialized and optimized for deployment in production environments. TorchScript enables efficient execution of PyTorch models on various platforms, including GPUs, CPUs, and mobile devices.

Like TensorFlow, PyTorch also implements GPU kernel functions for efficient computation on GPUs. It implements optimized GPU kernels similar to TensorFlow.

So while both TensorFlow and PyTorch provide GPU kernel function abstractions, their underlying computational graph models and programming styles differ, bringing their own unique advantages and trade-offs.

OpenAI Triton is a programming framework developed by OpenAI for building and deploying large-scale machine learning models efficiently. It leverages TensorFlow as its backend, supporting a wide range of models including deep learning and traditional algorithms. Triton offers tools for distributed computing, automated hyperparameter tuning, and model serving. It simplifies model deployment and management, making it suitable for both research and production environments. Triton abstracts away the need for users to write low-level GPU kernel functions by using TensorFlow’s optimized GPU operations implemented with CUDA, NVIDIA’s parallel computing platform. This approach allows developers to focus on defining high-level machine learning models without worrying about GPU optimization details.

It’s worth noting that Triton is built on top of TensorFlow, which supports alternative GPU acceleration libraries through backend integrations, and this enables Triton to leverage these alternatives to CUDA. One such alternative to CUDA is ROCm (Radeon Open Compute platform), developed by AMD. ROCm is an open-source GPU computing platform that provides support for AMD GPUs. TensorFlow has been working on integrating with ROCm, allowing it to utilize AMD GPUs for deep learning computations. As Triton relies on TensorFlow, it can benefit from this integration to support AMD GPUs through ROCm.

TorchScript for Model Optimization and Model Serving

· Leave a comment ·

TorchScript is an intermediate representation of a PyTorch model that can be optimized and run in a non-Python environment, making the PyTorch model suitable for deployment. It is part of the PyTorch ecosystem (Intro_to_TorchScript_tutorial.html , TorchScript JIT.html ).

Why is TorchScript needed ? Python while excellent for ML model development ( interpreted, REPL, simplicity, integration with number of ML libraries), also has characteristics that make it less suitable for model production deployments. These characteristics include interpretation overheads, complex dependency management, high memory/CPU overheads and the lack of easy integration with native technologies such as C++ for high performance and for embedded systems. TorchScript provides tools for optimizations such as operator fusion and static graph analysis which can improve the efficiency and performance during inference. Optimizing the models is crucial for embedded systems with limited resources.

PyTorch had introduced eager/dynamic execution, which had the advantage of faster user feedback but the disadvantage of not having as many optimizations as were possible in static approaches as in Tensorflow.

A blog on Key points to grasp about TorchScript – https://medium.com/@hihuaweizhu/key-points-to-grasp-for-torchscript-beginners-c02cf94aaa50, makes several good points, including that TorchScript is a subset of PyTorch and consists of statically typed variables.

A discussion between eager mode and script mode at https://towardsdatascience.com/pytorch-jit-and-torchscript-c2a77bac0fff suggests the benefit of TorchScript is more about dev/production (versus training/inference), with the production version requiring performance optimizations and portability. Quote: “With TorchScript, PyTorch aims to create a unified framework from research to production. TorchScript will take your PyTorch modules as input and convert them into a production-friendly format.“

NVIDIA uses TorchScript to facilitate the deployment and optimization of PyTorch models within their ecosystem. The Torchscript models are compiled to TensorRT, the Nvidia runtime .

AWS ML software stack, Neuron, supports tracing in torchscript. https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/torch/torch-neuronx/api-reference-guide/inference/api-torch-neuronx-trace.html . https://pytorch.org/docs/master/generated/torch.jit.trace.html#torch.jit.trace . An example of a neuron sdk trace for pytorch – https://github.com/aws-neuron/aws-neuron-sdk/issues/371 .

PyTorch/XLA is another project that integrates with Google XLA compiler to enable running PyTorch models on Google TPUs.

GraphCore produces hardware for deep learning called a GraphCore Intelligence Processing Unit (IPU). The primary software framework provided by GraphCore to execute machine learning models on their IPUs is Poplar. It allows running models from TensorFlow and PyTorch. Poplar optimizes computations for the unique architecture of GraphCore’s IPUs. This includes optimizations for memory bandwidth, parallel processing, and other hardware-specific features.

AlphaFold for protein structure prediction with deep learning – how does attention work

· Leave a comment ·

AlphaFold is a deep learning model developed by DeepMind that predicts protein structure. It uses a two-step process: First, it generates a representation of the protein’s amino acid sequence. Then, it refines this representation to predict the 3D structure of the protein. The model is trained on a large database of known protein structures and uses a convolutional neural network (CNN) to make these predictions. It leverages the concept of attention mechanisms to incorporate information from multiple parts of the protein sequence during the prediction process. It combines advanced machine learning techniques with protein structure data to make accurate predictions about protein folding.

Attention mechanisms are a key component of AlphaFold and play a crucial role in capturing dependencies between different parts of a protein sequence. It was interesting to deconstruct attention in a non-text use case – the median human protein is made of only 375 amino acids and most proteins are encoded with less than a 1000 amino acids. So to understand the attention mechanism, let’s break it down step by step using AlphaFold as the reference use case.

  1. Embedding the Protein Sequence:
    AlphaFold starts by embedding the amino acid sequence of a protein into a numerical representation. Each amino acid is represented as a vector, and these vectors are combined to form the input sequence matrix, X ∈ ℝ^(L×D), where L is the length of the sequence and D is the dimensionality of each amino acid vector.
  2. Creating Query, Key, and Value Matrices:
    AlphaFold then generates three matrices – Query (Q), Key (K), and Value (V) – by linearly transforming the input sequence matrix X. This transformation is performed using learnable weight matrices WQ, WK, and WV. The resulting matrices are Q = XWQ, K = XWK, and V = XWV, each having dimensions of L×D.
  3. Calculating Attention Weights:
    The attention mechanism computes the similarity between each query vector and key vector by taking their dot products. This similarity is scaled by a factor of √(D), and a softmax function is applied to obtain attention weights. The attention weights determine how much each key contributes to the final output. Let’s denote the attention weights matrix as A ∈ ℝ^(L×L), where each element A_ij represents the attention weight between the i-th query and j-th key. The attention weights are calculated as follows: A_ij = softmax((Q_i ⋅ K_j) / √(D)) Here, Q_i represents the i-th row of the Query matrix, and K_j represents the j-th row of the Key matrix.
  4. Weighted Sum of Values:
    The final step is to compute the weighted sum of the Value matrix using the attention weights. This is done by taking the matrix multiplication of attention weights A and the Value matrix V. The resulting matrix C, representing the context or attended representation, is given by: C = AV The context matrix C has dimensions of L×D, where each row represents a weighted sum of the Value vectors based on the attention weights.

Attention mechanisms in AlphaFold allow the model to capture the relationships and dependencies between different parts of the protein sequence. By assigning attention weights to relevant amino acids, the model can focus on important regions during the prediction process, enabling accurate protein structure predictions.

The dimensions of the matrices involved are as follows:

  • Input Sequence Matrix (X): X ∈ ℝ^(L×D), where L is the length of the protein sequence and D is the dimensionality of each amino acid vector.
  • Query Matrix (Q): Q ∈ ℝ^(L×D), same as the dimensions of X.
  • Key Matrix (K): K ∈ ℝ^(L×D), same as the dimensions of X.
  • Value Matrix (V): V ∈ ℝ^(L×D), same as the dimensions of X.
  • Attention Weights Matrix (A): A ∈ ℝ^(L×L), where each element A_ij represents the attention weight between the i-th query and j-th key.
  • Context Matrix (C): C ∈ ℝ^(L×D), same as the dimensions of X.

So the matrices Q, K, V, X, and C have dimensions L×D, where L represents the length of the protein sequence and D represents the dimensionality of the amino acid vectors. The attention weights matrix A has dimensions L×L, capturing the attention weights between each query and key pair.

OpenFold is a pytorch based reproduction of alphafold , a comparison is used in https://wandb.ai/telidavies/ml-news/reports/OpenFold-A-PyTorch-Reproduction-Of-DeepMind-s-AlphaFold–VmlldzoyMjE3MjI5

What are amino acids ? Amino acids are organic compounds that serve as the building blocks of proteins. They contain an amino group (-NH2) and a carboxyl group (-COOH) attached to a central carbon atom, along with a specific side chain (R-group). The side chain varies among different amino acids, giving them unique properties.

There are 20 standard amino acids that are commonly found in proteins. Each amino acid has a unique structure and properties, determined by its specific side chain. Some examples include glycine, alanine, valine, leucine, isoleucine, serine, threonine, cysteine, methionine, aspartic acid, glutamic acid, lysine, arginine, histidine, phenylalanine, tyrosine, tryptophan, asparagine, glutamine, and proline.

Amino acids encode proteins through a process called translation. The genetic information stored in DNA is transcribed into messenger RNA (mRNA). The mRNA is then read by ribosomes, which assemble amino acids in a specific sequence according to the instructions provided by the mRNA. This sequence of amino acids forms a polypeptide chain, which then folds into a functional protein with a specific structure and function. The sequence of amino acids in a protein is determined by the sequence of nucleotides in the corresponding mRNA molecule.

Structure of AlphaFold from Nature paper, and as described here.

https://www.forbes.com/sites/robtoews/2023/07/16/the-next-frontier-for-large-language-models-is-biology/

https://www.nature.com/articles/s41592-023-01924-w A team of researchers led by Peter Kim at Stanford University has performed guided protein evolution using protein language models that were trained on millions of natural protein sequences.

https://aibusiness.com/nlp/meta-lays-off-team-behind-its-protein-folding-model

https://techcrunch.com/2024/06/25/evolutionaryscale-backed-by-amazon-and-nvidia-raises-142m-for-protein-generating-ai/

https://github.com/evolutionaryscale/esm

https://github.com/aws-samples/drug-discovery-workflows

https://github.com/lucidrains/alphafold3-pytorch

https://github.com/google-deepmind/alphafold

https://fold.it/about_foldit

https://build.nvidia.com/explore/biology