미분 가능한 캐시 확장을 통한 잠재 공간에서의 심사결정
Deliberation in Latent Space via Differentiable Cache Augmentation
December 23, 2024
저자: Luyang Liu, Jonas Pfeiffer, Jiaxing Wu, Jun Xie, Arthur Szlam
cs.AI
초록
대규모 언어 모델(LLMs)이 중간 추론 단계를 생성하고 주의를 기울여 "더 많이 생각"할 수 있는 기술은 복잡한 문제 해결에 유망함을 입증했습니다. 그러나 표준 접근 방식은 응답하기 바로 전에 이산 토큰의 시퀀스를 생성하므로 상당한 대기 시간 비용이 발생하고 최적화하기 어려울 수 있습니다. 본 연구에서는 얼어붙은 LLM에 오프라인 공변처리 장치를 추가하여 모델의 키-값(kv) 캐시에서 작동할 수 있음을 보여줍니다. 이 공변처리 장치는 잠재적 임베딩 집합을 사용하여 캐시를 보완하고 후속 디코딩의 충실성을 향상시키도록 설계되었습니다. 우리는 이 공변처리 장치를 표준 사전 훈련 데이터의 디코더에서 언어 모델링 손실을 사용하여 훈련하면서 디코더 자체는 얼어두는 방식으로 이를 보여줍니다. 이 접근 방식을 통해 모델이 엔드-투-엔드로 어떻게 추가 계산을 kv-캐시로 증류할지 학습할 수 있습니다. 디코더가 변경되지 않았기 때문에 공변처리 장치는 오프라인 및 비동기적으로 작동할 수 있으며 공변처리 장치를 사용할 수 없거나 특정 캐시가 추가 계산이 필요하지 않다고 판단되는 경우 언어 모델이 정상적으로 작동할 수 있습니다. 실험적으로 캐시가 보완되면 디코더가 다수의 후속 토큰에서 낮은 헷갈림을 달성한다는 것을 보여줍니다. 더욱이 과제별 훈련이 없어도 캐시 보강이 추론 집중적 작업 범위에서 일관되게 헷갈림을 줄이고 성능을 향상시킨다는 것을 실험을 통해 입증합니다.
English
Techniques enabling large language models (LLMs) to "think more" by
generating and attending to intermediate reasoning steps have shown promise in
solving complex problems. However, the standard approaches generate sequences
of discrete tokens immediately before responding, and so they can incur
significant latency costs and be challenging to optimize. In this work, we
demonstrate that a frozen LLM can be augmented with an offline coprocessor that
operates on the model's key-value (kv) cache. This coprocessor augments the
cache with a set of latent embeddings designed to improve the fidelity of
subsequent decoding. We train this coprocessor using the language modeling loss
from the decoder on standard pretraining data, while keeping the decoder itself
frozen. This approach enables the model to learn, in an end-to-end
differentiable fashion, how to distill additional computation into its
kv-cache. Because the decoder remains unchanged, the coprocessor can operate
offline and asynchronously, and the language model can function normally if the
coprocessor is unavailable or if a given cache is deemed not to require extra
computation. We show experimentally that when a cache is augmented, the decoder
achieves lower perplexity on numerous subsequent tokens. Furthermore, even
without any task-specific training, our experiments demonstrate that cache
augmentation consistently reduces perplexity and improves performance across a
range of reasoning-intensive tasks.