💬 NLP

PROMPT2BOX: Uncovering Entailment Structure among LLM Prompts

The paper introduces PROMPT2BOX, a novel framework that embeds LLM prompts into a box embedding space to capture specificity relations beyond mere topical similarity, thereby enabling more fine-grained weakness analysis and improved hierarchical clustering compared to traditional vector-based methods.

Original authors: Neeladri Bhuiya, Shib Sankar Dasgupta, Andrew McCallum, Haw-Shiuan Chang

Published 2026-03-24

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Neeladri Bhuiya, Shib Sankar Dasgupta, Andrew McCallum, Haw-Shiuan Chang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a teacher trying to figure out which students in a class are struggling with a specific subject, like "writing stories."

The Old Way: The "Similarity" Map (Vector Embeddings)

Currently, most researchers use a tool that acts like a similarity map. They take every writing prompt a student gets and plot it on a map based on how similar the words are.

The Problem: If you ask a student to "Write a story" and another to "Write a 500-word story about a robot falling in love," the old map sees these as almost the exact same spot because they both contain the words "write," "story," and "robot."
The Result: The teacher looks at the map and sees a big red blob (meaning "bad performance") covering both prompts. They think, "Oh, this student is bad at writing stories."
The Reality: The student might actually be great at writing general stories but terrible at the specific, difficult task of writing a robot love story. The old map is too blurry to see the difference between a "general request" and a "hard, specific request." It treats a simple question and a complex puzzle as the same thing.

The New Way: PROMPT2BOX (The "Box" Method)

The authors of this paper, PROMPT2BOX, realized that prompts aren't just points on a map; they are containers with different sizes.

Think of a prompt as a box:

A Big Box (General Prompt): "Write a story." This box is huge. It can hold any kind of story. It's easy to fit a good answer inside a big box.
A Small Box (Specific Prompt): "Write a 500-word story about a robot falling in love." This box is tiny and very specific. It's much harder to fit a perfect answer inside a tiny, cramped box.

The Magic of "Box Containment":
The genius of PROMPT2BOX is that it understands that the Small Box fits inside the Big Box.

If you can write a story about a robot falling in love (the small box), you have automatically proven you can write a story (the big box).
But if you can write a story (the big box), it doesn't mean you can write the specific robot love story (the small box).

How It Helps Find Weaknesses

When the researchers use this new "Box Map":

They see that the student is doing great in the Big Box (general stories).
They see the student is failing in the Small Box (robot love stories).
Because the Small Box is inside the Big Box, the map clearly shows: "The student isn't bad at stories; they are just bad at this specific, difficult type of story."

Why This Matters

In the real world, Large Language Models (LLMs) like the ones powering chatbots are like these students. Developers need to know exactly where the models fail so they can train them better.

Without PROMPT2BOX: Developers might think, "Our model is bad at writing," and waste time retraining it on everything.
With PROMPT2BOX: Developers can say, "Our model is great at general writing, but it struggles specifically with 'romantic robot adventures'." They can then fix just that one tiny, specific problem.

The Visual Analogy

Imagine looking at a city:

The Old Map (Vectors): Shows you a blurry photo where the whole city looks like one big gray area. You can't tell if the traffic jam is in the downtown core or just a small side street.
The New Map (Boxes): Shows you a 3D model where every neighborhood is a clear, distinct box. You can see that the "Downtown" box is huge and full of traffic, but the "Robot Love Story" alleyway is a tiny, cramped box where the traffic is completely gridlocked.

Summary

PROMPT2BOX is a new tool that stops treating all questions as equal points on a map. Instead, it treats them as boxes of different sizes. By understanding that a specific question is a "small box inside a big box," it helps us pinpoint exactly where AI models are weak, allowing us to fix their specific problems rather than guessing at their general abilities.

1. Problem Statement

Current methods for analyzing Large Language Model (LLM) weaknesses rely on embedding prompts into vector spaces and clustering them based on topical similarity. However, vector embeddings suffer from a critical limitation: they primarily capture semantic similarity while failing to model specificity and asymmetric relationships (entailment).

The Specificity Gap: Prompts sharing a topic but differing in constraints (e.g., "write a story" vs. "write a 500-word romantic story about a robot") are often mapped to similar vector locations.
Consequence: This leads to ambiguous weakness analysis. If an LLM fails on specific, difficult prompts but succeeds on general ones, vector clustering may average these scores, masking the specific failure mode or incorrectly labeling the entire topic as a weakness.
Goal: The authors aim to develop a representation that captures not just what a prompt is about, but how specific it is, enabling the detection of entailment structures (where a specific prompt implies a general one).

2. Methodology: PROMPT2BOX

The core innovation is mapping prompts into a Box Embedding Space rather than a vector space.

A. Box Embedding Representation

Instead of representing a prompt as a point $v \in \mathbb{R}^D$ , PROMPT2BOX represents it as an axis-aligned hyper-rectangle (a box) defined by:

Center Vector ( $c$ ): Captures the semantic location (similar to vector embeddings).
Width Vector ( $\delta$ ): Captures the semantic scope.
- General prompts (fewer constraints) are represented by larger boxes.
- Specific prompts (more constraints) are represented by smaller boxes.

Entailment via Containment:
The geometric relationship between boxes models entailment naturally. If Prompt A entails Prompt B (A is more specific than B), the box for A is contained within the box for B ( $Box(A) \subseteq Box(B)$ ).

Intersection Volume: Measures semantic similarity.
Containment Probability: Measures entailment strength. If $Box(A) \subseteq Box(B)$ , the conditional probability $P(B|A)$ approaches 1.

B. Training Architecture

The system uses a trained encoder to map text prompts to box parameters.

Encoder: Based on a Sentence Transformer (MPNet-base) with two MLP heads outputting the center and width vectors.
Optimization: Uses Gumbel Box Formulation to make the hard min/max operations of box intersection differentiable.
Loss Function: A contrastive learning objective that maximizes intersection volume for similar pairs and enforces containment for entailment pairs.

C. Data Curation & Synthesis

Training box embeddings requires entailment data, which is scarce. The authors synthesized a large-scale dataset using four sources:

Infinity Instruct: For general semantic relevance (prompt-response pairs).
MultiNLI: For sentence-level entailment (premise/hypothesis).
WildChat (Synthesized): Used GPT-4.1 to generate hierarchical instruction trees, creating parent-child (general-to-specific) relationships.
SURI: Used to create "sibling" relationships (prompts sharing a parent but with different constraints) to teach the model to distinguish non-entailing but related prompts.
Linkage Strategy: A novel step to align the semantic space of the instruction-response data with the entailment hierarchy, preventing the model from learning disjoint representations.

D. Dimensionality Reduction (BOX-SNE)

To visualize high-dimensional box embeddings in 2D, the authors propose BOX-SNE. Unlike t-SNE (which minimizes KL divergence of point distances), BOX-SNE optimizes low-dimensional box locations to preserve:

Intersection volumes (similarity).
Containment relationships (entailment).

3. Key Contributions

PROMPT2BOX Framework: The first method to embed LLM prompts into a box space to jointly capture semantic similarity and specificity/entailment.
Synthetic Entailment Datasets: A pipeline for synthesizing hierarchical instruction data and sibling relationships to train box encoders at scale.
BOX-SNE: A novel dimensionality reduction technique specifically designed for box embeddings to preserve geometric containment.
New Evaluation Metrics: Metrics to assess similarity, entailment, and specificity ordering in prompt representations.

4. Experimental Results

The authors evaluated PROMPT2BOX against vector baselines (using cosine similarity) and a norm-based vector baseline (CSDelta) across several tasks:

Intrinsic Metrics:
- STS-B (Similarity): Box models performed competitively with vector baselines.
- FollowBench & SURI (Entailment/Specificity): Box models significantly outperformed vectors. For example, on FollowBench (retrieving more specific instructions), Box achieved 73.8% accuracy vs. 64.0% for vectors.
- Ablation: Removing entailment data or linkage data degraded performance, proving the necessity of both the box geometry and the specific training data.
Score Prediction (UltraFeedback):
- Using k-NN regression to predict LLM scores on unseen prompts, the Box model achieved the lowest RMSE (1.5280), outperforming vector baselines (1.6059) and random baselines. This indicates that understanding prompt specificity improves the ability to predict model performance.
Hierarchical Clustering:
- Weakness Discovery: On 17 LLMs, PROMPT2BOX identified 8.9% more weakness clusters than vector baselines.
- Specificity Alignment: The correlation between hierarchy depth and instruction specificity was 33% stronger for Box models (70.04% agreement) compared to vectors (52.71%).
- Visualization: 2D visualizations clearly separated datasets by difficulty (e.g., WildBench prompts appeared as smaller, more specific boxes than WildChat prompts), a distinction lost in vector t-SNE visualizations.

5. Significance and Impact

Granular Weakness Analysis: PROMPT2BOX allows developers to distinguish between "the model is bad at this topic" and "the model is bad at specific, constrained variations of this topic." This is crucial for targeted data collection and fine-tuning.
Beyond Vectors: It challenges the assumption that vector embeddings are sufficient for all NLP tasks, demonstrating that asymmetric relationships (entailment) require region-based representations.
Applications: The ability to estimate prompt specificity and difficulty has broad applications in:
- LLM Routing: Sending specific prompts to specialized models.
- Evaluation: Selecting harder prompts to stress-test models.
- Safety: Identifying specific constraints that lead to unsafe outputs.
- Interpretability: Visualizing the "solution space" of an LLM's capabilities.

In summary, PROMPT2BOX provides a geometrically intuitive and mathematically rigorous framework for understanding the hierarchical structure of LLM prompts, significantly improving the detection of model weaknesses compared to traditional vector-based approaches.