Distilling Structural Reasoning (DSR)

Efficient Semantic Parsing via Chain-of-Thought Rationalization and Contrastive Demonstration Selection

Seyed Hossein Ahmadpanah — Department of Computer Engineering, Islamic Azad University, Tehran, Iran

---

Overview

DSR is a neuro-symbolic framework for Hierarchical Semantic Parsing (HSP) that translates natural language queries into executable, nested logical forms. It addresses two core limitations of LLM-based parsers:

1. Structural hallucinations — caused by surface-level demonstration retrieval (BM25, cosine similarity) that cannot distinguish queries like "flights to Boston" vs. "flights from Boston"
2. Prohibitive inference latency — Chain-of-Thought prompting with 70B-parameter models is too slow for real-time applications

DSR closes the gap by distilling the structural reasoning of a large Teacher model (Llama-2-70B) into a compact Student model (Llama-3-8B), achieving 87.42% Exact Match on TOP — a new state-of-the-art for open-source models — while running 9.4× faster than the teacher.

---

Key Results

Ablation — Retrieval Strategy (TOP, Teacher model):

Zero-shot domain transfer (TOPv2 Weather): 71.4% EM vs. 42.3% for fixed-classifier baseline.

---

Architecture

The DSR framework has two training phases and four novel components:

Phase 1: Rationale Synthesis (Teacher)
┌─────────────────────────────────────────────────────────┐
│  Input Query                                            │
│       │                                                  │
│  [CFS Encoder] ──→ Top-k Structurally Similar Demos    │
│       │                                                  │
│  [Prompt Builder] ──→ Structural CoT Prompt            │
│       │                                                  │
│  [Teacher LLM, 70B] ──→ Structural CoT Trace r + LF y  │
└─────────────────────────────────────────────────────────┘

Phase 2: Efficient Inference (Student)
┌─────────────────────────────────────────────────────────┐
│  Augmented Dataset D_aug = {(x_i, r_i, y_i)}           │
│       │                                                  │
│  [Rationale Distillation] ──→ Fine-tuned Student (8B)  │
│       │                                                  │
│  [Trie-Based Constraints] ──→ Syntax-valid decoding    │
│       │                                                  │
│  [Parser-Critic Loop] ──→ Semantic self-correction     │
└─────────────────────────────────────────────────────────┘

Component 1: Contrastive Fragment Selection (CFS)

A bi-encoder (RoBERTa-Large) trained with InfoNCE loss on hard-negative triplets:

• Positive: Same semantic fragments, different surface phrasing
  ("Fly from NYC" ↔ "Depart from New York")
• Hard Negative: High lexical overlap (Jaccard > 0.5), different root intent
  ("Flights to Boston" ↔ "Flights from Boston")

$$\mathcal{L}{CFS} = -\log \frac{\exp(\text{sim}(h_x, h{x^+}) / \tau)}{\sum_{j \in {x^+, x^-}} \exp(\text{sim}(h_x, h_j) / \tau)}$$

Component 2: Structural Chain-of-Thought

Decomposes the hierarchical parse tree into a Depth-First Search (DFS) sequence of semantic fragment decoding steps:

Query: "Book a flight for three people to Boston"

Step 1: Identify Root Intent [IN:FLIGHT].
Step 2:   Identify Slot [SL:QUANTITY] with value 'three people' (depth 1).
Step 3:   Identify Slot [SL:DEST] with value 'Boston' (depth 1).
Final: Combine all fragments → [IN:FLIGHT [SL:QUANTITY three people ] [SL:DEST Boston ] ]

Component 3: Rationale Distillation

Trains the Student on the Teacher's CoT traces via Supervised Fine-Tuning (SFT):

$$\mathcal{L}{\text{distill}} = -\sum{t=1}^{|r|+|y|} \log P_{M_S}(z_t \mid x, z_{<t})$$

where $z = [r ; y]$ is the concatenation of the reasoning trace $r$ and logical form $y$.

Component 4: Trie-Based Constrained Decoding

A prefix trie built from the domain schema masks invalid tokens at each decoding step:

$$m_t[w] = \begin{cases} 0 & \text{if } w \in V_{\text{valid}}(z_{<t}) \ -\infty & \text{otherwise} \end{cases}$$

Eliminates 100% of syntax errors (hallucinated slots, invalid bracket nesting).

Component 5: Parser-Critic Feedback Loop

Cycle-consistency check for semantic fidelity:

$$\text{Score}(\hat{y}) = \cos(E(x), E(\hat{x})) > \delta$$

• Critic reconstructs $\hat{x}$ from $\hat{y}$; accepted only if similarity > $\delta = 0.85$
• On rejection: resample with $T = 0.7$, up to $K = 3$ iterations
• Catches ~4% of semantically misaligned parses (e.g., dropped quantity constraints)

---

Installation

git clone https://github.com/ahmadpanah/distilling-structural-reasoning
cd distilling-structural-reasoning

# Install core dependencies
pip install -r requirements.txt

# Install as a package (editable)
pip install -e .

# Optional: fast inference with vLLM
pip install vllm

Requirements: Python ≥ 3.9, PyTorch ≥ 2.0, CUDA GPU recommended (A100 for full training)

---

Quick Start

Inference with a Pre-trained Student

from src.pipeline import DSRPipeline

# Load pipeline with a fine-tuned student checkpoint
pipeline = DSRPipeline.from_pretrained(
    student_checkpoint="./checkpoints/student",
    retriever_checkpoint="./checkpoints/cfs_encoder.pt",
    demo_bank=demo_bank,  # list of {'query': ..., 'logical_form': ..., 'rationale': ...}
)

result = pipeline.parse("Book a flight from New York to Boston for three people")
print(result["logical_form"])
# [IN:FLIGHT [SL:ORG New York ] [SL:DEST Boston ] [SL:QUANTITY three people ] ]

print(f"Critic score: {result['score']:.3f}")
print(f"Demos used:   {result['demos_used']}")

Batch Inference

queries = [
    "What is the weather in Paris tomorrow?",
    "Set a reminder for my 3pm meeting",
    "Navigate to the nearest coffee shop",
]
results = pipeline.parse_batch(queries)
for q, r in zip(queries, results):
    print(f"Q: {q}")
    print(f"LF: {r['logical_form']}\n")

---

Data Preparation

The framework supports four benchmarks. Download each dataset and place under ./data/:

data/
├── top/
│   ├── train.tsv       # query \t logical_form
│   ├── eval.tsv
│   └── test.tsv
├── topv2/
│   ├── navigation/
│   ├── reminder/
│   ├── weather/
│   └── ...
├── atis/
│   ├── train.json
│   ├── dev.json
│   └── test.json
└── snips/
    ├── train.json
    ├── dev.json
    └── test.json

Loading a Dataset

from src.data.datasets import load_dataset

dataset = load_dataset("top", data_dir="./data/top")
print(dataset)
# TOPDataset({'train': 39281, 'dev': 2318, 'test': 4370})

# Access splits
for ex in dataset.train[:3]:
    print(ex["query"], "→", ex["logical_form"])

---

Training

Full Pipeline (Recommended)

python scripts/train.py \
    --dataset top \
    --data_dir ./data/top \
    --student_model meta-llama/Meta-Llama-3-8B-Instruct \
    --output_dir ./checkpoints \
    --use_gold_rationale     # No teacher inference needed

--use_gold_rationale generates structural CoT traces deterministically from the gold logical forms (DFS linearization). This is equivalent to the main setup in the paper and requires no 70B teacher model at training time.

To use a live Teacher model for rationale generation (requires the 70B model):

python scripts/train.py \
    --dataset top \
    --data_dir ./data/top \
    --no_gold_rationale \
    --teacher_model meta-llama/Llama-2-70b-chat-hf

Step-by-Step Training

Step 1: Train CFS Retriever

from src.models.cfs import CFSEncoder, mine_hard_negatives, train_cfs_encoder
from src.data.datasets import load_dataset

dataset = load_dataset("top", "./data/top")
queries = [ex["query"] for ex in dataset.train]
lfs     = [ex["logical_form"] for ex in dataset.train]

triplets = mine_hard_negatives(queries, lfs, jaccard_threshold=0.5)
print(f"Mined {len(triplets)} hard-negative triplets")

encoder = CFSEncoder(model_name="roberta-large")
encoder = train_cfs_encoder(
    encoder,
    triplets,
    num_epochs=10,      # Paper: 10 epochs
    temperature=0.05,   # Paper: τ = 0.05
)

import torch
torch.save(encoder.state_dict(), "./checkpoints/cfs_encoder.pt")

Step 2: Build Augmented Dataset

from src.models.cfs import ContrastiveFragmentSelector
from src.models.structural_cot import build_augmented_dataset

selector = ContrastiveFragmentSelector(encoder, demo_bank=dataset.train, top_k=5)
selector.index_demo_bank()

augmented = build_augmented_dataset(
    queries=queries,
    logical_forms=lfs,
    demo_selector=selector,
    teacher=None,
    use_gold_rationale=True,   # DFS linearization (no 70B needed)
)
# augmented[0] → {'query': ..., 'logical_form': ..., 'rationale': 'Step 1: ...'}

Step 3: Distill Student

from src.training.distillation import RationaleDistillationTrainer, DistillationConfig

config = DistillationConfig(
    student_model_name="meta-llama/Meta-Llama-3-8B-Instruct",
    num_epochs=3,
    batch_size=4,
    gradient_accumulation_steps=8,
    learning_rate=2e-4,
    use_lora=True,
    lora_r=16,
    output_dir="./checkpoints/student",
)

trainer = RationaleDistillationTrainer(config)
trainer.train(
    train_examples=augmented,
    eval_examples=dataset.dev,
)

Hardware Requirements

---

Evaluation

# Evaluate student on TOP test set
python scripts/evaluate.py \
    --student_checkpoint ./checkpoints/student \
    --retriever_checkpoint ./checkpoints/cfs_encoder.pt \
    --dataset top \
    --data_dir ./data/top \
    --split test

# Ablation: no CFS retrieval (uses BM25-style random retrieval)
python scripts/evaluate.py \
    --student_checkpoint ./checkpoints/student \
    --dataset top --data_dir ./data/top \
    --no_cfs

# Zero-shot domain transfer (Weather domain, no weather training data)
python scripts/evaluate.py \
    --student_checkpoint ./checkpoints/student \
    --dataset topv2 --data_dir ./data/topv2 \
    --domain weather \
    --zero_shot

# Save results to JSON
python scripts/evaluate.py \
    --student_checkpoint ./checkpoints/student \
    --dataset atis --data_dir ./data/atis \
    --output_file results/atis_test.json

Programmatic Evaluation

from src.data.datasets import load_dataset
from src.evaluation.metrics import Evaluator
from src.pipeline import DSRPipeline

dataset = load_dataset("top", "./data/top")
pipeline = DSRPipeline.from_pretrained("./checkpoints/student", demo_bank=dataset.train)

evaluator = Evaluator(dataset_name="TOP")
result = evaluator.evaluate(
    parser_fn=lambda q: pipeline.parse(q)["logical_form"],
    examples=dataset.test,
)
print(result)

Output:

=== TOP Results ===
Exact Match:       87.42%
Tree-F1:           0.9156
Struct. Error:     0.00%
Avg Latency:       5.9ms
Tokens/sec:        170.0
Num Examples:      4370
Per-depth EM:
  Depth 1: 94.20%
  Depth 2: 89.31%
  Depth 3: 82.10%
  Depth 4: 71.80%
  Depth 5: 62.40%

---

Repository Structure

distilling-structural-reasoning/
│
├── src/
│   ├── models/
│   │   ├── cfs.py               # Contrastive Fragment Selection
│   │   │                        #   CFSEncoder, InfoNCELoss, ContrastiveFragmentSelector
│   │   │                        #   mine_hard_negatives, train_cfs_encoder
│   │   ├── structural_cot.py    # Structural Chain-of-Thought prompting
│   │   │                        #   TeacherModel, build_augmented_dataset
│   │   ├── trie_decoder.py      # Trie-Based Constrained Decoding
│   │   │                        #   GrammarTrie, LogitMaskConstraint
│   │   └── parser_critic.py     # Parser-Critic Feedback Loop
│   │                            #   Critic, ParserCriticLoop
│   ├── training/
│   │   └── distillation.py      # Rationale Distillation Protocol
│   │                            #   RationaleDistillationDataset, Trainer, StudentParser
│   ├── data/
│   │   └── datasets.py          # Dataset loaders: TOP, TOPv2, ATIS, SNIPS
│   ├── evaluation/
│   │   └── metrics.py           # EM, Tree-F1, per-depth analysis, latency
│   ├── utils/
│   │   └── semantic_fragments.py  # SemanticFragment, parse_logical_form, Tree-F1
│   └── pipeline.py              # DSRPipeline (inference) + DSRTrainingPipeline
│
├── scripts/
│   ├── train.py                 # End-to-end training entry point
│   └── evaluate.py              # Benchmark evaluation entry point
│
├── tests/
│   └── test_dsr.py              # 30 unit tests (pytest)
│
├── requirements.txt
├── setup.py
└── README.md

---

Module Reference

src.models.cfs

src.models.structural_cot

src.models.trie_decoder

src.models.parser_critic

src.training.distillation

src.utils.semantic_fragments

---

Extending DSR

Adding a New Domain (Zero-Shot)

No fine-tuning required. Just provide demonstrations from existing domains and let CFS retrieve structurally similar examples:

pipeline = DSRPipeline.from_pretrained(
    "./checkpoints/student",
    demo_bank=existing_domain_examples,  # e.g., navigation + reminder demos
)
# CFS will retrieve structurally similar fragments even for unseen "healthcare" slots
result = pipeline.parse("Schedule a blood test for next Monday at 9am")

Custom Schema with Trie Constraints

from src.models.trie_decoder import GrammarTrie

schema = {
    "intents": ["IN:BOOK_APPOINTMENT", "IN:CANCEL_APPOINTMENT"],
    "slots": {
        "IN:BOOK_APPOINTMENT": ["SL:DATE", "SL:TIME", "SL:DOCTOR"],
        "IN:CANCEL_APPOINTMENT": ["SL:DATE", "SL:APPOINTMENT_ID"],
    },
    "special_tokens": ["[", "]"],
}
trie = GrammarTrie(schema)
trie.build()
# Attach to pipeline for constrained decoding

Using Your Own Sentence Encoder for the Critic

from sentence_transformers import SentenceTransformer
from src.models.parser_critic import Critic, ParserCriticLoop

sbert = SentenceTransformer("all-mpnet-base-v2")
critic = Critic(sentence_encoder=sbert)
loop = ParserCriticLoop(critic, threshold=0.85, max_iterations=3)

---

Testing

# Run all 30 unit tests
pytest tests/ -v

# Run with coverage report
pytest tests/ --cov=src --cov-report=term-missing

The test suite covers all five components without requiring model weights or dataset downloads.

---

License

MIT License. See LICENSE for details.