超越RAG:面向企業AI的任務感知知識壓縮(AWS實現)
傳統的檢索增強生成(RAG)在處理跨越數百份文件的分析任務時存在侷限。本文介紹如何在AWS上使用任務感知知識壓縮(TAKC),預先將整個知識庫壓縮為任務特定的表示,按多個保真度層級快取,並將每個查詢路由到適當的層級,同時提供可部署的開源實現。
如果您正在使用檢索增強生成(RAG)來處理需要跨數百份文件的複雜分析任務,例如財務盡職調查或法規合規審查,您很可能已經遇到了它的天花板。相似性搜尋能夠找到相關的片段,但常常遺漏跨文件的連線。本文將向您展示如何利用任務感知知識壓縮(TAKC)來彌補這一差距,該技術可將整個知識庫預先壓縮為任務特定的表示,並部署在AWS上。您可以在自己的賬戶中部署一個完整的開源實現。
任務感知知識壓縮
考慮一傢俬募股權公司正在評估一項價值5億美元的製造企業收購。盡職調查團隊必須分析涵蓋12家子公司和5年時間的財務報表。他們還面臨200多份供應商合同、來自8個設施的環境合規報告以及50多起法律案件。當分析師詢問關於當前供應商條款和未決訴訟的綜合財務風險時,RAG的相似性搜尋無法提供答案。數百份文件包含相關資訊,但它們之間的聯絡缺乏詞彙相似性。
TAKC透過使用LLM生成更短、更聚焦於任務的文件摘要來解決這類問題,不同的任務會有不同的摘要。不同的任務需要從同一份文件中提取不同的資訊。為財務分析壓縮的年度報告需要收入資料、利潤率和現金流資料。而為合規審查壓縮的同一份報告則需要監管引文和違規歷史。通用摘要試圖涵蓋所有內容,這會稀釋任何特定用例的資訊密度。TAKC透過特定任務的鏡頭來壓縮文件,保留重要的內容並丟棄其餘部分。在攝取管道部分,展示了壓縮提示如何具體指定要保留的資訊。對於生產部署,請將任務型別提示儲存在版本化配置中(例如AWS Systems Manager Parameter Store或專門的Amazon Simple Storage Service(Amazon S3)字首),以便提示更改可審計,並在提示更新時觸發重新壓縮。
系統離線壓縮文件,每個任務型別每個文件一次。查詢時,系統檢索預壓縮的表示而非原始文件。然後使用壓縮版本而非完整文件來回答問題。如果壓縮表示不夠詳細,查詢複雜度分析器會將問題路由到保留更多上下文的較低壓縮層級。
TAKC以壓縮形式提供對整個知識庫的訪問,而不僅僅是相似性搜尋返回的前k個塊。由於壓縮時文件被一起處理,系統保留了文件間的聯絡。它還能從同一源材料為不同任務生成不同的壓縮輸出。一份10-K檔案的財務壓縮與法律風險壓縮完全不同。壓縮可將token數量減少8倍至64倍,同時針對保留任務相關資訊。
多速率壓縮
不同的查詢需要不同級別的保真度。像“第三季度營收是多少?”這樣的問題所需上下文遠少於分析供應商付款條件與各子公司季度現金流之間關係的請求。
TAKC透過為每種任務型別維護四個壓縮層級來解決這一問題。在最小壓縮(8倍)下,系統保留大約87.5%的上下文。它保留了足夠的多步推理和跨文件綜合細節。在中等級別(16倍)下,上下文減少約93.8%。此級別適用於中等複雜度的分析查詢。高階別(32倍)將上下文減少約96.9%,適用於事實查詢和明確定義的問題。在超高階別(64倍)下,減少約98.4%。此級別適用於分類任務和關鍵詞查詢。
查詢複雜度分析器基於查詢長度、問題型別以及分析性語言的存在等訊號,將傳入的問題路由到適當的層級。直接的事實性問題將命中超壓縮快取,而複雜的分析問題則使用輕度壓縮快取。這對使用者來說是透明的。
大多數企業查詢是查詢操作,可以從較高壓縮層級以最低成本提供服務。不常見的複雜查詢僅在需要時消耗較大的上下文預算。這種基於層級的路由補充了現有的RAG最佳化技術,如後設資料過濾和查詢重新格式化,這些技術可以在應用壓縮之前縮小文件集。為了驗證壓縮質量,請將每個層級的LLM響應與從完整未壓縮文件生成的響應進行比較。參考實現包括測試指令碼,可以對您的特定任務型別和文件進行此比較。
AWS上的架構
該實現在AWS上作為兩個解耦的無伺服器管道執行:一個用於攝取,一個用於查詢。圖1展示了這兩個管道。
TAKC架構:管道流程處理資料攝取和壓縮。使用者流程處理經過身份驗證的查詢。
我們選擇AWS Lambda進行計算,因為每個函式呼叫都是短生命週期的且由事件驅動。工作負載在攝取期間突發處理資料,並在突發之間處理可變的查詢負載,使得無伺服器成為自然的選擇。
我們選擇Amazon API Gateway將查詢介面暴露為REST端點。對於快取,我們選擇Amazon ElastiCache Serverless用於在複合鍵(takc:{task}:{rate})上進行讀取,無需分片管理。Amazon Cognito處理JWT發放和令牌重新整理,無需自定義身份驗證程式碼,減少了實現面積。
攝取管道
當文件以任務型別字首(例如raw-data/financial/)落入Amazon S3時,S3事件通知會觸發一個AWS Lambda函式。該函式將文件分割為256個token的塊,重疊50個token以防止邊界資訊丟失。然後,它非同步呼叫一個壓縮Lambda函式來處理每個塊,實現並行處理。對於大規模攝取,請在壓縮函式上配置預留併發,並在分塊和壓縮步驟之間放置一個Amazon Simple Queue Service(Amazon SQS)佇列,以優雅地處理節流。
第二個函式呼叫Amazon Bedrock在所有四個壓縮層級壓縮這些塊。每次壓縮呼叫包括一個任務感知提示,告知模型要保留哪些資訊:
任務:財務分析。保留收入指標、利潤率、現金流、債務義務和財務風險指標。 壓縮目標:減少到原始長度的約1/16。 說明:
- 專注於與任務相關的事實和關係
- 保留數值資料和指標
- 維護實體及其屬性
- 保留因果關係和依賴關係
- 刪除冗餘或不相關的資訊
模型知道要保留什麼,因為提示指定了任務的重要內容。這一區別使得它成為任務感知而非通用壓縮。系統將壓縮輸出儲存在Amazon ElastiCache Serverless中,鍵結構如takc:financial:medium,並備份到S3以實現永續性。Redis OSS資料模型支援多速率快取查詢所需的分層鍵結構。快取條目使用24小時TTL,並備份到S3。如果快取條目被驅逐或過期,查詢函式將回退到S3備份,並在讀取時重新填充快取。
查詢管道
使用者透過Amazon Cognito進行身份驗證,接收JWT,然後透過Amazon API Gateway傳送查詢。AWS WAF位於API前端,用於速率限制和威脅防護。一個Lambda函式使用啟發式方法(關鍵詞訊號和查詢長度)分析查詢複雜度,從Amazon ElastiCache Serverless檢索適當的壓縮快取,然後將壓縮上下文和查詢傳送到Amazon Bedrock進行推理。當路由置信度低時,系統預設使用中等壓縮層級作為安全回退。
更高成本的Bedrock壓縮呼叫僅在攝取期間發生一次。查詢路徑包括快取查詢和在壓縮上下文上的推理。
該棧使用Amazon Bedrock(Anthropic Claude 3 Haiku、Claude 3 Sonnet和Amazon Titan Text)進行壓縮和推理。模型選擇可透過CDK上下文值配置,無需程式碼更改。AWS Lambda(Python 3.12+)處理資料處理和查詢邏輯。Amazon ElastiCache Serverless儲存壓縮快取,而Amazon S3儲存原始資料、塊和快取備份(KMS加密)。Amazon API Gateway暴露REST端點,Amazon Cognito提供基於JWT的身份驗證。AWS WAF位於API前端,提供速率限制和託管安全規則。Amazon CloudWatch提供監控和指標,AWS Key Management Service(AWS KMS)管理加密金鑰。
將基礎設施定義為單個AWS Cloud Development Kit(AWS CDK)棧,並使用單個命令部署。AWS CDK提供可重複的部署,並允許您透過上下文值自定義引數,如壓縮塊大小、Lambda記憶體分配和ElastiCache儲存限制。對於生產部署,考慮將有狀態資源(S3、ElastiCache、Cognito)分離到自己的棧中,以減少爆炸半徑並實現計算和儲存層的獨立生命週期管理。
成本比較
針對一個包含100,000個token的知識庫,每天查詢1,000次(輸出token被排除,因為響應長度與上下文大小無關):
方法 每次查詢輸入token 每日輸入token 相對輸入成本 完整上下文 100,000 100,000,000 100% RAG(前10個塊) ~10,000 10,000,000 10% TAKC輕度(8倍) ~12,500 12,500,000 12.5% TAKC中級(16倍) ~6,250 6,250,000 6.25% TAKC高階(32倍) ~3,125 3,125,000 3.1% TAKC超高階(64倍) ~1,563 1,563,000 1.6%
壓縮帶來的token節省直接源自壓縮比率。實際節省取決於您特定的Bedrock模型定價和查詢模式。
TAKC需要預先透過一次性Bedrock呼叫進行壓縮成本。對於變化不頻繁且被反覆查詢的知識庫,該成本可以得到攤銷。對於每小時變化的知識庫,RAG的每次查詢檢索模型可能更實用。
何時選擇TAKC、RAG或兩者結合
下表總結了根據工作負載特徵,每種方法更適合的情況:
因素 有利於TAKC 有利於RAG 查詢型別 跨文件推理、綜合 直接的事實查詢 知識庫穩定性 每天或更少變化 每小時或更頻繁變化 任務可預測性 明確定義的任務型別 不可預測的查詢模式 覆蓋範圍要求 必須考慮整個語料庫 只有少數文件相關 來源歸屬 不需要 需要(使用者需要看到來源) Token預算 緊張 靈活
在實踐中,生產系統通常受益於兩者。RAG高效處理快速查詢。TAKC處理基於檢索的方法無法捕捉到連線的分析查詢。查詢複雜度分析器可以在它們之間路由。當使用者需要將響應追溯到特定源文件時,使用RAG。對於需要跨文件推理和可審計性的受監管工作負載,結合TAKC和RAG:使用TAKC進行分析性響應,使用RAG檢索支援性源文件以形成審計線索。
開始使用
前提條件
要部署此參考實現,請確保您具備以下條件:
一個具有Amazon Bedrock模型訪問許可權的AWS賬戶。
AWS CDK CLI。
Python 3.12或更高版本。
Node.js 18或更高版本。
部署
[為節省成本,原文已截斷]