The AI Brain: SageMaker BGE Cross-Encoder Reranking & Sub-100ms Bedrock ConverseStream SSE (Part 4)

In Part 3 of this series, we built our Storage Vault in PostgreSQL 16: co-locating dense pgvector embeddings and full-text BM25 indexes, using Reciprocal Rank Fusion (RRF) to merge candidate chunks into a unified pool of 15 relevant documents.

At this stage, most engineering teams stop and immediately forward all 15 chunks into their Large Language Model (LLM).

Doing that introduces three major production hazards:

1. The "Lost in the Middle" Effect: Research shows LLM recall degrades significantly when the most relevant facts are buried in the middle of a massive 15-chunk prompt context.
2. Context Bloat & Cost: Sending 15 chunks (roughly 6,000–8,000 tokens) on every user turn drives up token inference costs and increases latency.
3. The Bi-Encoder Ceiling: Even after RRF, vector similarity scores often fail to capture subtle semantic nuances, leading the LLM to hallucinate or generate vague answers.

To achieve production reliability, an enterprise platform requires a dedicated Stage-2 Precision Reranker.

In this article, we cover:

1. The Architecture Ceiling: Why Bi-Encoders fail and how Cross-Encoders achieve 0.9970 confidence.
2. The Two-Stage Serving Architecture: Reducing 100,000 documents to 3 pristine chunks in under 60ms.
3. Production SageMaker Deployment: Running BAAI/bge-reranker-base on provisioned ml.m5.large for deterministic sub-50ms inference.
4. Real-World Benchmark Results: How cross-encoders rescued ambiguous queries in our live DocLab tests.
5. Streaming Generation: Delivering real-time Server-Sent Events (SSE) tokens via AWS Bedrock ConverseStream with sub-400ms Time-to-First-Token (TTFT).

---

1. The Accuracy Ceiling: Bi-Encoder vs Cross-Encoder

To understand why a reranker is necessary, we must examine the difference between Bi-Encoders and Cross-Encoders:

How Bi-Encoders Work (Stage 1)

All vector embedding models (Amazon Titan v2, OpenAI text-embedding-3, Cohere) are Bi-Encoders:

• The User Query $Q$ is passed through the neural network independently to generate vector $E(Q)$.
• The Document Passage $D$ is passed through the neural network independently during ingestion to generate vector $E(D)$.
• At retrieval time, the system computes the cosine dot product: $\text{Sim}(Q, D) = E(Q) \cdot E(D)$.

The Fundamental Flaw: Query tokens and document tokens never interact in the neural network's attention layers. The model compresses an entire 512-token document into a single fixed vector. Nuanced relationships (e.g., subtle conditionality, negation, or multi-step logic) are lost during compression.

How Cross-Encoders Work (Stage 2)

A Cross-Encoder (such as BAAI/bge-reranker-base) does not produce isolated vector coordinates. Instead, it passes both the Query and the Candidate Chunk together as a single input pair:

Input: [CLS] User Query [SEP] Document Passage [EOS]

Inside the Cross-Encoder's transformer layers:

• Full Cross-Attention: Every single token in the query performs self-attention with every single token in the document passage.
• Deep Alignment: The model evaluates how words in the query depend on words in the document, capturing semantic subtleties that Bi-Encoders miss entirely.
• Output: A single calibrated relevance probability between 0.0000 and 1.0000.

---

2. The Two-Stage Serving Pattern: Cost vs Precision

If Cross-Encoders are vastly superior, why not use them for the entire search pipeline?

Compute Complexity: Running full cross-attention across 100,000 documents for a single user query would take dozens of seconds and require tens of thousands of dollars in GPU compute.

The industry gold standard is the Two-Stage Retrieval Pattern:

• Stage 1 (High Recall, Low Cost): PostgreSQL 16 executes hybrid HNSW cosine + BM25 tsvector search, using Reciprocal Rank Fusion (RRF) to filter 100,000+ candidates down to the Top 15 candidates in ~15ms.
• Stage 2 (High Precision, Deep Attention): AWS SageMaker evaluates those 15 candidate pairs in parallel, returning the Top 3 chunks with calibrated confidence scores (>0.99) in ~40ms.

Total retrieval time: under 60ms with zero loss in precision!

---

3. Deploying BAAI/bge-reranker-base on AWS SageMaker

For our real-time reranking engine, we deployed the state-of-the-art open-source BAAI/bge-reranker-base model onto an AWS SageMaker Real-Time Endpoint named bge-reranker-endpoint.

Why Provisioned ml.m5.large Over Serverless?

Many engineers attempt to deploy SageMaker Serverless Inference to cut baseline idle costs. In production RAG platforms, this is a dangerous trap:

• Serverless Cold Starts: PyTorch container initialization and model weight loading introduce 5 to 10 second cold-start delays when endpoints scale from zero.
• Provisioned Determinism: An ml.m5.large provisioned instance costs approximately ~$0.115/hour (~$80/month). In return, it guarantees deterministic sub-50ms inference latency under all traffic spikes, with zero cold starts.

The SageMaker Endpoint Input/Output Contract

DocLab communicates with SageMaker using a compact JSON contract:

Request Payload:

{
  "query": "What are the scaling policies defined for checkout-service?",
  "texts": [
    "Chunk A: Service catalog for checkout-service specifying horizontal pod autoscaler...",
    "Chunk B: Payment gateway timeout and retry threshold definitions...",
    "Chunk C: ADR-025 Ingestion Factory architectural decision record..."
  ]
}

SageMaker Response Payload:

[
  { "index": 0, "score": 0.9970 },
  { "index": 2, "score": 0.4120 },
  { "index": 1, "score": 0.0381 }
]

The Spring Boot backend immediately sorts the results by score descending and extracts the top 3 chunks (filtering out any chunk with confidence below a configurable threshold, such as 0.60).

---

4. Real-World Benchmark Results from Live AWS Infrastructure

We executed live comparative benchmark evaluations on our deployed DocLab architecture.

The test corpus consisted of diverse enterprise files: Architectural Decision Records (ADRs), Excel sheets, and complex JSON service catalogs.

Evaluation Metric	Pure Dense (Titan v2)	Pure Sparse (BM25)	Hybrid (RRF k=60)	Two-Stage (RRF + SageMaker BGE)
ADR-025 Decision Query	0.8841 (Rank #1)	0.0000 (Miss)	0.01639 (Rank #2)	0.9970 (Confidence)
ADR-025 Consequences	0.8120 (Rank #2)	12.41 (Rank #1)	0.03226 (Rank #1)	0.9957 (Confidence)
Technical JSON Schema	0.7910 (Rank #3)	0.0000 (Miss)	0.01587 (Rank #3)	0.6219 (Promoted to #1)
Candidate Reduction	N/A	N/A	15 candidates	3 pristine chunks
Retrieval Stage Latency	~8ms	~3ms	~14ms	~54ms total

Key Findings:

1. Confidence Calibration: Pure dense cosine similarities often cluster tightly between 0.75 and 0.85, making it impossible to set an absolute threshold. The SageMaker Cross-Encoder provided crystal-clear discrimination: true answers scored > 0.99, while irrelevant chunks plummeted to < 0.05.
2. Context Compression: By feeding only the top 3 high-confidence chunks to the LLM instead of 15 noisy chunks, we reduced prompt token consumption by 72% and completely eliminated the "Lost in the Middle" hallucination phenomenon.

---

5. Streaming Generation via AWS Bedrock ConverseStream

With our top 3 chunks validated, we construct our guarded prompt and initiate streaming inference via Amazon Bedrock ConverseStream:

The Fail-Closed Guardrail Prompt

System:
You are an enterprise AI assistant for DocLab. 
Rules:
1. Answer strictly and exclusively using the provided context chunks below.
2. If the answer cannot be found in the context, explicitly say:
   "I cannot answer this question based on the verified enterprise documentation."
3. Cite the Chunk ID for every factual assertion made.
4. Do not speculate or extrapolate beyond the text.

Context Chunks:
[Chunk #1 - Confidence: 0.9970]
Service Catalog: checkout-service auto-scales when CPU > 75%...

User Query:
What is the scaling policy for checkout-service?

Low-Latency Server-Sent Events (SSE)

Instead of forcing the user to wait 4–6 seconds for the full response to generate, DocLab exposes a streaming endpoint:

GET /api/v1/chat/stream?query=...&allowedRoles=ROLE_ANALYST

Spring Boot bridges Bedrock's asynchronous reactive stream directly to the client's browser using standard HTTP text/event-stream (Server-Sent Events):

event: delta
data: {"token": "The "}

event: delta
data: {"token": "checkout-service "}

event: delta
data: {"token": "scales "}

event: delta
data: {"token": "automatically..."}

End-to-End Latency Waterfall

Here is the measured latency breakdown on live AWS infrastructure:

┌──────────────────────────────────────────────┐
│ Stage 1: PostgreSQL Hybrid RRF (15ms)        │
├──────────────────────────────────────────────┤
│ Stage 2: SageMaker BGE Reranker (40ms)       │
├──────────────────────────────────────────────┤
│ Bedrock TTFT (Time to First Token) (325ms)   │
└──────────────────────────────────────────────┘
Total Time-to-First-Token on Screen: ~380ms

To the end-user, the chat interface feels instantaneous. Words begin appearing on screen in under 400 milliseconds, fully backed by verified, cited enterprise documentation.

---

6. What's Next in Part 5 (The Grand Finale)?

Our platform now parses multi-format files at zero cost, stores them in a unified hybrid vault, reranks candidates to 0.9970 confidence, and streams responses in under 400ms.

In Part 5 (The Grand Finale), we tie the entire enterprise platform together:

• Zero-Trust Cache Partitioning: Implementing ElastiCache Redis multi-tenant session isolation.
• Fail-Closed Guardrails: Catching prompt injections and preventing data exfiltration.
• The Complete Retrospective: Architectural trade-offs, budget review (running this on under $200 in AWS credits), and key takeaways for enterprise architects.

---

Built with AWS Bedrock, Amazon SageMaker, Amazon RDS PostgreSQL 16 pgvector, and Spring Boot.