Koa-decision: Fast and Consistent Structured Decision Making with Generative LLMs
summary
The gist
The gist The Koa-action framework proposes a method for low-latency atomic actions by formulating classification as constrained generation with single-token outputs, achieving frontier accuracy at
In short
Koa-decision introduces a method for fast and consistent structured decision-making using generative LLMs. It reformulates classification as constrained generation with single-token outputs, employing atomic label tokens. This approach achieves frontier accuracy at sub-second speeds, making complex decision tasks practical for real-time applications.
Key concepts
- Constrained Generation Task
- This treats the classification problem as a generation task where the model is strictly guided to produce only a single special token. This design bypasses the overhead of multi-token decoding, ensuring the output is always exactly one token, which drastically reduces latency to near O(1) time.
- Atomic Label Tokens
- These are specialized tokens like [ok] created for each class and added to the model's vocabulary. The model is trained specifically to map input text or vision directly to one of these specific tokens, simplifying the classification process into a direct token prediction.
- Randomized Label Assignments
- During training, the mapping between classes and control tokens is shuffled randomly. This prevents the model from simply memorizing fixed token-label pairs. Instead, it forces the model to rely on contextual descriptions in prompts rather than relying on static token identities.
- Distribution Matching Objective
- For multi-label classification, this training objective replaces standard cross-entropy loss. The model is trained to match a 'soft target distribution' across the control tokens of all relevant labels simultaneously, allowing it to predict multiple correct labels in a single forward pass.
Terminology used across episodes
This episode discusses
- Koa-decision: Fast and Consistent Structured Decision Making with Generative LLMs · Paper Radio
- LLM2Vec: Large Language Models Are Secretly Powerful Text Encoders
- SALSA: Single-pass Autoregressive LLM Structured Classification
- BERT for Joint Intent Classification and Slot Filling
- XGrammar: Flexible and Efficient Structured Generation Engine for Large Language Models
- MMTEB: Massive Multilingual Text Embedding Benchmark
- Making Pre-trained Language Models Better Few-shot Learners
- Grammar-Constrained Decoding for Structured NLP Tasks without Finetuning
- DeBERTa: Decoding-enhanced BERT with Disentangled Attention
- Prefix-Tuning: Optimizing Continuous Prompts for Generation
- RoBERTa: A Robustly Optimized BERT Pretraining Approach
- Large Language Models Do Multi-Label Classification Differently
- MTEB: Massive Text Embedding Benchmark
- Open-Jev Judgments on CallScreenBench: Calibrated One-Pass Scam Screening with a Small Language Model · Paper Radio
- Exploiting Cloze Questions for Few Shot Text Classification and Natural Language Inference
- Text Classification via Large Language Models
- Let Me Speak Freely? A Study on the Impact of Format Restrictions on Performance of Large Language Models
- Qwen2-VL: Enhancing Vision-Language Model's Perception of the World at Any Resolution
- Finetuned Language Models Are Zero-Shot Learners
- Efficient Guided Generation for Large Language Models
- Fine-Tuning Causal LLMs for Text Classification: Embedding-Based vs. Instruction-Based Approaches
The paper
Koa-decision: Fast and Consistent Structured Decision Making with Generative LLMs · Read on arXiv
Shenghong Dai, Shiva Kumar Pentyala, Yingchi Liu, Shubham Mehrotra, Suman Banerjee, James Zhu, Bin Bi, Sitaram Asur
Salesforce AI
Transcript
Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.
Tom: Today's paper: "Koa-decision: Fast and Consistent Structured Decision Making with Generative LLMs".
Jane: The gist The Koa-action framework proposes a method for low-latency atomic actions by formulating classification as constrained generation with single-token outputs, achieving frontier accuracy at sub-second speeds.
Tom: First, who's behind it and why it matters.
Title and authors: Tom: So we're looking at this paper today, "Koa-decision: Fast and Consistent Structured Decision Making with Generative LLMs." It sounds a bit technical, but the main idea is that they are trying to solve the problem of making very fast decisions when you use large language models.
Jane: Exactly. Right now, when we use these big models for tasks like classification, they often have to generate a whole string of tokens before we get an answer. That takes time, which isn't always good for real-time applications.
Lu: The title really gets to the point there. They are focusing on making decisions fast and consistent using generative LLMs instead of just generating long answers. It’s about efficiency in making specific choices, like classification or scoring things quickly >
The paper's summary: Tom: So, what is the actual method they propose? Basically, they treat classification not as a normal generation task where you output a whole sentence, but as something called constrained generation with single-token outputs.
Jane: That means instead of letting the model pick one token out of thousands for an entire response, it’s forced to pick just one specific special token that represents the correct label >
Lu: They introduce these atomic label tokens, like ok, and they use supervised fine-tuning to teach the model to map input text or even vision directly to that single special token rather than a whole sequence of words >
Meng: So, instead of running a big decoder for every decision, it turns into this deterministic one-step decoding problem. That should cut down the latency significantly >
The paper's improvements: Tom: The biggest improvement they highlight is the latency reduction. They claim that by using this single-token approach, they can achieve a latency of O(one), which is constant time, regardless of how long the output might be > <ref:2609.36115#pg1>
Jane: That’s really fast. And they show that on standard benchmarks, this method gets competitive accuracy while keeping the latency low and stable across different tasks >
Lu: They also mention something interesting about zero-shot adaptation. Because they use these special tokens, you can reassign the label tokens at inference time without having to retrain the model for a new task >
Meng: So, if you have a new category that pops up in your inventory, you don't need to spend hours retraining the entire system; you just change the token assignment >
Conclusion: Tom: So we’ve seen how they use this constrained generation with single-token outputs and how it aims for constant time decisions. The paper "Koa-decision" is showing that we can get high accuracy at sub-second speeds for things like classification and scoring >
Jane: It really shows that decoder-style models, when adapted this way, can match or even beat encoder baselines on multimodal datasets while keeping the flexibility of large language models >
Lu: What’s important is how they handle multi-label classification too; they switch the training objective to a distribution matching objective so the model can pick a set of relevant labels in one pass >
Meng: From an engineering side, it’s that scalability. The latency stays essentially flat even as the number of possible labels grows because you aren't dealing with variable decoding steps anymore >
Lalam: I see this allowing us to build systems where the AI can make these tiny, atomic decisions extremely quickly and reliably for our users in real-time applications >
Tom: So that’s how "Koa-decision" tackles the latency problem for structured decision making with generative models. A lot of exciting stuff ahead >
More episodes
- 2610.10926-Adaptive Multi-Discriminator WGAN Framework for Resource-Constrained Internet of Vehicles Using Reinforcement Learning and Game Theory
- 2610.10931-Coefficient Calibration as Selection Pressure in Symbolic Regression
- 2610.10932-World-Model Policy Arbiter for Goal-Conditioned Reinforcement Learning
- 2610.10947-RH-Detect: A Unified Benchmark for Reward Hacking Detection
- 2610.10965-Spectrally Targeted Muon
- 2610.10972-Transferability of Learned States in Neural PDE Solvers
- 2610.10998-F cubed NO: Frequency-Decomposed Finite-Time Flow-map Neural Operators with Cross-Scale Conditioning
- 2610.10996-Invariant-Measure Reasoners: Stable Representations for Latent Reasoning
- 2610.11000-Low-rank tensor structure of precipitation and its application to satellite-reference merging
- 2610.11133-Machine Learning Optimization for Enhanced OS Fingerprinting