How “Compress • Sense • Expand” lets LLMs handle huge knowledge bases without drowning in tokens.

Retrieval-Augmented Generation (RAG) is the backbone of knowledge-grounded LLM apps — but it’s painfully slow once you start feeding the model thousands of retrieved tokens. The root problem? Transformer self-attention cost grows roughly with the square of input length, so more tokens → much more compute and latency. Enter REFRAG from Meta Superintelligence Labs: a clever new pattern that gives you both speed and accuracy.

REFRAG’s core idea is elegantly simple: don’t make the LLM read everything. Instead, Compress most of the retrieved text into tiny dense summaries, Sense which pieces contain literal facts that must survive compression, and Expand only those pieces back into raw tokens. The result: tiny token inputs for the LLM, huge context coverage, and literal facts preserved.

Why this matters (quick)

• Massive speedups (Meta reports large TTFT improvements) — fewer raw tokens → much less quadratic attention work.

• Better scale — treat documents that are 10–100x longer than the model’s native window.

• Fact fidelity — numbers, names, clause IDs preserved because those chunks are left as raw text.

• Lower memory footprint — smaller KV cache and cheaper inference.

The REFRAG workflow (plain-language)

1. Retrieve the set of relevant documents/passages as usual.

2. Chunk each retrieved passage into tiny fixed-size windows (e.g., 8–32 tokens).

3. Compress most chunks into embeddings via a lightweight encoder.

4. Sense which chunks are critical (via a policy — small classifier, heuristic scoring, or an RL agent).

5. Expand the critical chunks back into raw tokens; keep the rest as embeddings.

6. Feed the LLM a hybrid input: user query (raw tokens) + some raw chunks + many chunk embeddings.

7. LLM decodes with access to both literal facts (raw tokens) and semantic context (embeddings).

Real-time examples (concrete)

1) Legal brief: “How much was the final settlement and which clause allowed punitive damages?”

• Retrieval pulls 15 passages across a 200-page brief (8,000 tokens).

• REF RAG chunks, encodes, and senses: “$15,340,000” and “Clause 7.1” chunks get expanded to raw tokens; the rest are embeddings.

• LLM now has exact amount + clause literal text plus semantic context — fast and precise.

2) Research literature scan: “List all experiments that reported >90% accuracy on dataset X”

• Retriever returns dozens of papers. REF RAG compresses background, expands only the chunks that contain numeric results (e.g., “accuracy=93.2% on X”), so the model can list exact numbers and cite the right paper lines without reprocessing entire PDFs.

3) Support knowledge base (chatbot for product):

• Many KB articles with long troubleshooting steps. Compress standard instructions and expand only error codes and configuration lines (e.g., ERR_1357, config file snippets). The chatbot answers quickly and copies literal commands correctly.

When to expand vs compress (intuitions)

• Expand if chunk contains: exact numbers, named entities, configuration/code, clause IDs, dates, legal citations, or anything where literal text matters.

• Compress if chunk is background, rationale, or stylistic language that only needs to be understood, not reproduced.

A lightweight sensing policy can be a simple classifier that flags chunks when:

• Numeric tokens appear

• Named-entity density is high

• The chunk’s embedding is unusually similar to the query embedding (and contains digits/special tokens)
  You can later replace this with a learned RL policy to optimize downstream metrics.

Node.js Example: a REFRAG-style pipeline

The example below is an end-to-end Node.js implementation that demonstrates the REFRAG pattern. It uses:

• A vector DB (abstracted — examples for Pinecone/Milvus/Weaviate or a simple in-memory store)

• OpenAI-style embeddings (you can swap any embedding model)

• A simple heuristic “sensing” policy (you can later train an RL policy)

• Hybrid input assembly and a call to an LLM for final answer generation

This is example code — adapt models/keys/DBs to your stack. Replace process.env.OPENAI_API_KEY etc. with real keys.

// filename: refrag_example.js
// Node 18+
// Install: npm i axios express body-parser dotenv
// Optional: npm i openai (if you prefer official SDK)

import express from "express";
import bodyParser from "body-parser";
import axios from "axios";
import dotenv from "dotenv";
dotenv.config();

const app = express();
app.use(bodyParser.json());

/**
 * Config - change to your vector DB / endpoints
 */
const OPENAI_API_KEY = process.env.OPENAI_API_KEY;
const PORT = process.env.PORT || 3000;

/**
 * ----- Utilities -----
 */

// naive tokenizer split for chunking (replace with real tokenizer if needed)
function tokenizeText(text) {
  return text.split(/\s+/);
}
function detokenize(tokens) {
  return tokens.join(" ");
}

// chunk text into fixed-size chunks (nTokens per chunk)
function chunkText(text, nTokens = 16) {
  const tokens = tokenizeText(text);
  const chunks = [];
  for (let i = 0; i < tokens.length; i += nTokens) {
    chunks.push(detokenize(tokens.slice(i, i + nTokens)));
  }
  return chunks;
}

/**
 * Embeddings call (OpenAI embedding endpoint example)
 * Replace with your model or local encoder for speed.
 */
async function getEmbedding(text) {
  const resp = await axios.post(
    "https://api.openai.com/v1/embeddings",
    {
      input: text,
      model: "text-embedding-3-small", // example
    },
    { headers: { Authorization: `Bearer ${OPENAI_API_KEY}` } }
  );
  return resp.data.data[0].embedding;
}

/**
 * Vector DB abstraction:
 * - storeChunks: accepts [{id, chunk, embedding, meta}, ...]
 * - query(queryEmbedding, topK) => [{id, chunk, score, meta}, ...]
 *
 * For demo, we use an in-memory "vector DB" with cosine similarity search.
 */
const vectorDB = {
  store: [],

  async storeChunks(chunksWithEmb) {
    this.store.push(...chunksWithEmb);
  },

  _cosine(a, b) {
    let dot = 0, na = 0, nb = 0;
    for (let i = 0; i < a.length; i++) { dot += a[i] * b[i]; na += a[i]*a[i]; nb += b[i]*b[i]; }
    return dot / (Math.sqrt(na) * Math.sqrt(nb) + 1e-10);
  },

  async queryByEmbedding(queryEmbedding, topK = 10) {
    const res = this.store
      .map(item => ({ ...item, score: this._cosine(item.embedding, queryEmbedding) }))
      .sort((a,b) => b.score - a.score)
      .slice(0, topK);
    return res;
  }
};

/**
 * Simple sensing policy:
 * - accepts a chunk string and its embedding and returns {expand: boolean, reason: string}
 * - heuristic:
 *   * if chunk contains digits, $, % or clause-like tokens -> expand
 *   * if chunk embedding sim to query embedding past threshold and contains punctuation like ':' or '(' -> expand
 *
 * You can replace this with a small classifier or RL model later.
 */
function simpleSensingPolicy(chunk, chunkEmbedding, queryText, queryEmb, simScoreToQuery = 0.8) {
  const literalTrigger = /\d|[$%#]|Clause|clause|Section|§|\bErr_|config|:=/i;
  if (literalTrigger.test(chunk)) return { expand: true, reason: "literal tokens detected" };
  // otherwise compare embeddings if available (here we pass sim externally if computed)
  // For simplicity, user may provide simScoreToQuery computed outside; we'll mock it below
  return { expand: false, reason: "no literal tokens" };
}

/**
 * Assemble hybrid prompt for LLM:
 * - user query raw tokens
 * - expanded raw chunks (literal strings)
 * - compressed chunks embedded as "summaries" (we'll show their semantic tags)
 *
 * You will need an LLM that can accept a combination of raw tokens and metadata/embedding inputs.
 * Many production systems encode embeddings into the prompt as short natural-language summaries —
 * or use specialized LLMs that accept key/value memory. Below we represent compressed embeddings
 * as small semantic summaries (via a "summarizeChunk" step) to make a simple demo.
 */
async function summarizeChunk(chunk) {
  // Quick summarization via LLM or heuristic; here we truncate and return a tiny summary.
  // Replace with a cheap instruction-tuned model or use the embedding to find keywords.
  if (chunk.length > 120) return chunk.slice(0, 120) + "…";
  return chunk;
}

/**
 * Call LLM to get final answer (raw prompt assembly)
 */
async function callLLMForAnswer(userQuery, rawChunks, compressedSummaries) {
  // Build a prompt that includes:
  // 1) user's query
  // 2) exact chunks (literal)
  // 3) compressed summaries as bullet points
  let prompt = `You are an assistant grounded on provided evidence.\n\nUser question:\n${userQuery}\n\nEvidence (exact):\n`;
  for (const r of rawChunks) prompt += `- ${r}\n`;
  prompt += `\nContext summaries:\n`;
  for (const s of compressedSummaries) prompt += `- ${s}\n`;

  // Call OpenAI completion / chat
  const resp = await axios.post(
    "https://api.openai.com/v1/chat/completions",
    {
      model: "gpt-4o-mini", // replace with appropriate model
      messages: [{ role: "system", content: "You are a precise assistant that uses exact evidence." },
                 { role: "user", content: prompt }],
      max_tokens: 400
    },
    { headers: { Authorization: `Bearer ${OPENAI_API_KEY}` } }
  );

  return resp.data.choices[0].message.content;
}

/**
 * Main pipeline endpoint: /ask
 * Body: { query: string }
 */
app.post("/ask", async (req, res) => {
  try {
    const { query } = req.body;
    if (!query) return res.status(400).json({ error: "query required" });

    // 1) (Demo) retrieve top K docs from DB by running query embedding -> vectorDB.queryByEmbedding
    const queryEmb = await getEmbedding(query);
    const retrievalHits = await vectorDB.queryByEmbedding(queryEmb, 50); // many chunks

    // 2) For each retrieved chunk: run sensing policy (heuristic)
    // We'll compute a simple cosine sim with query as already in retrievalHits.score
    const rawChunks = [];
    const compressedSummaries = [];
    for (const hit of retrievalHits) {
      const policy = simpleSensingPolicy(hit.chunk, hit.embedding, query, queryEmb, hit.score);
      if (policy.expand) {
        rawChunks.push(hit.chunk);
      } else {
        // compress -> create a short natural-language summary for prompt
        compressedSummaries.push(await summarizeChunk(hit.chunk));
      }
    }

    // limit counts to keep token lightweight
    const maxRaw = 40;
    const maxSummaries = 200;
    const finalRaw = rawChunks.slice(0, maxRaw);
    const finalSummaries = compressedSummaries.slice(0, maxSummaries);

    // 3) Call LLM with hybrid input
    const answer = await callLLMForAnswer(query, finalRaw, finalSummaries);

    res.json({ answer, rawCount: finalRaw.length, summaryCount: finalSummaries.length });
  } catch (err) {
    console.error(err?.response?.data || err.message || err);
    res.status(500).json({ error: "internal", details: err?.message || err });
  }
});

/**
 * Small route to index some demo docs into in-memory vector DB
 */
app.post("/index", async (req, res) => {
  try {
    const { docs } = req.body; // docs: [{id, text, meta}]
    if (!Array.isArray(docs)) return res.status(400).json({ error: "docs array required" });

    const chunksWithEmb = [];
    for (const doc of docs) {
      const chunks = chunkText(doc.text, 16);
      for (let i = 0; i < chunks.length; i++) {
        const chunk = chunks[i];
        const emb = await getEmbedding(chunk);
        chunksWithEmb.push({
          id: `${doc.id}::${i}`,
          chunk,
          embedding: emb,
          meta: { docId: doc.id, ...doc.meta }
        });
      }
    }
    await vectorDB.storeChunks(chunksWithEmb);
    res.json({ indexed: chunksWithEmb.length });
  } catch (err) {
    console.error(err);
    res.status(500).json({ error: err.message });
  }
});

app.listen(PORT, () => {
  console.log(`REFRAG demo app running on http://localhost:${PORT}`);
});

How to try this quickly (example flow)

1. Start the server: node refrag_example.js (set OPENAI_API_KEY in .env).

2. Index a legal brief or doc collection:

curl -X POST http://localhost:3000/index -H "Content-Type: application/json" -d '{
  "docs": [
    { "id": "brief1", "text": "Clause 7.1 states that punitive damages up to $15,340,000 may be awarded. Final settlement of $15,340,000 was recorded on October 26, 2023. Additional commentary ...", "meta": {"type":"legal"}}
  ]
}'

1. Ask a query:

curl -X POST http://localhost:3000/ask -H "Content-Type: application/json" -d '{"query":"What was the final settlement amount and which clause allowed punitive damages?"}'

Expected behavior: the pipeline will expand the chunk containing $15,340,000 and Clause 7.1, keep the rest compressed, and the LLM returns the exact dollar amount and clause.

Notes & practical tips

Use a fast, local lightweight encoder for chunk embeddings

For performance, prefer a tiny encoder you can run locally (e.g., a distilled sentence-transformer) rather than remote embeddings for every chunk. Meta’s REFRAG uses a small dedicated encoder — that’s the key to keeping preprocessing fast.

Real RL agent vs. heuristics

Start with heuristics (digit detection, named-entity density, similarity thresholds). Once you have metrics (accuracy, precision, latency), train a small RL policy that optimizes the downstream LLM’s answer quality and latency tradeoff.

Where to store embeddings

For production: Pinecone, Milvus, Weaviate, or similar. For prototypes, in-memory or simple SQLite-based vector stores work.

How to represent embeddings to the LLM

Most LLM APIs don’t accept raw embeddings alongside tokens. Two common approaches:

• Convert embeddings back to short natural-language summaries (cheap and widely compatible).

• Use a model that supports plugging in key/value memories (some research and advanced production stacks do this). REFRAG-style systems often use special model support for dense memory.

Metrics to optimize

• TTFT (time-to-first-token): how fast the model can start returning an answer.

• Answer accuracy: exact literal correctness for numbers/names.

• Memory usage: KV cache size.

• Token cost: fewer raw tokens → lower cost.

How to iterate/improve

1. Replace simpleSensingPolicy with a small fine-tuned classifier that flags chunks that contain verbatim facts (train on labeled chunks).

2. Replace summarizeChunk with a fast instruction model that produces compact semantic tags for embeddings.

3. Implement a budgeted selection algorithm (maximize coverage/utility under a raw-token budget).

4. Train an RL agent with reward = downstream correctness — λ * latency.

TL;DR — Why REFRAG matters

REFRAG is a smart, practical way to let LLMs “read” enormous knowledge bases quickly without losing the literal facts that matter. It’s a systems trick: do cheap compression up front, sense which tiny bits must remain literal, and feed the model a hybrid input. You get the speed of embeddings and the exactness of raw text — the best of both worlds.