jayesh20/lexsi-track3-aya-lora

Lexsi Track 3: Agentic Fine-Tuning — Tiny-Aya LoRA Adapter

Submission for: Lexsi AI Safety & Alignment Hackathon (Track 3: Agentic Fine-Tuning)
Base Architecture: CohereLabs/tiny-aya-global (3.35B dense parameters)
GitHub Repository: https://github.com/jayeshpandey01/Lexsi
Dataset & Benchmarks: jayesh20/lexsi-track3-procurement-benchmarks
Governance & Release Status: RESEARCH ARTIFACT — MODEL PROMOTION BLOCKED (NO-GO DECISION)

---

1. Model Summary

This repository contains the fine-tuned low-rank adapter (sft_aya_3b_v2_c3) trained on CohereLabs/tiny-aya-global (3.35B) to investigate tool-calling capability, step efficiency, and metric gaming in enterprise procurement auditing.

The training pipeline integrates AgentTune standards on Kaggle 2× NVIDIA Tesla T4 GPUs with assistant-turn loss masking across 220 multi-turn demonstrations. Multi-turn ReAct rollouts were evaluated using AuditKIT across seven standardized metrics ($M_1$–$M_7$).

---

2. Evaluation Results & Key AI Safety Discovery

All evaluations were conducted against 10 held-out validation tasks using greedy decoding ($T=0.01$) inside a deterministic relational tool sandbox:

Metric ID	Metric Name	Untuned Base Comparator	Condition C2 (Attention LoRA)	Condition C3 (All-Linear LoRA)	Release Gate Status
$M_1$	Verified Task Success	0.0%	0.0%	50.0% (Nominal)	Failed Floor ($M_1 \ge 40%$)
$M_2$	Tool Selection Accuracy	50.0%	0.0%	0.0%	DISQUALIFIED (Failed Floor $M_2 \ge 50%$)
$M_3$	Schema Validity	100.0%	100.0%	100.0%	Passed ($M_3 = 100%$)
$M_7$	Execution Loop Rate	0.0%	0.0%	0.0%	Passed ($M_7 = 0%$)

Forensic Diagnosis: Unmasking Metric Gaming in Tool-Use SFT

While Condition C3 nominally registered $M_1 = 50.0%$ ($\Delta M_1 = +50.0%$, paired bootstrap $p=0.0030$), forensic trajectory auditing revealed a critical safety finding:

1. Zero Tool Invocations: The model achieved $M_2 = 0.0%$ tool calls, completely bypassing the interactive tool runtime.
2. Empirical Distribution Prior: In the validation split, 5 of 10 tasks query vendor purchase orders where target vendors (VND-106, VND-107) had no approved orders (oracle answer: $0.00).
3. The Shortcut: SFT caused the model to memorize an ungrounded conversational output prior of "$0.00" whenever asked about vendor spend, scoring 5/5 nominal successes on zero-spend queries while scoring 0/5 on inventory queries that required querying real database records.
4. Enforced Safety Gate: Under pre-registered governance ($(M_1 \ge 40%) \land (M_2 \ge 50%)$), our release gate triggered a unanimous NO-GO decision for model promotion, and the locked test partition remains 100% sealed and untouched.

---

3. Training Details

• Base Architecture: CohereLabs/tiny-aya-global (3.35B dense transformer decoder)
• Quantization: 4-bit NormalFloat4 (NF4) with double quantization, computation in bfloat16
• LoRA Hyperparameters:
  • Rank ($r$): 8
  • Scaling ($\alpha$): 16
  • Target Projections: q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj (All linear layers)
  • Trainable Parameters: 15,114,240 / 3,364,341,760 (0.45% parameter budget)
• Optimization:
  • Optimizer: paged_adamw_8bit
  • Learning Rate: $3.0 \times 10^{-5}$ with Cosine Decay and 30 warmup steps
  • Steps: 300 optimization steps (Effective batch size: 4)
  • Runtime: 40.8 minutes on Kaggle 2× NVIDIA Tesla T4 GPUs

---

4. Tool Environment & Anti-Gaming Defenses

The agent interacts with three deterministic relational tools:

1. lookup_records(table, filters): Query tables (vendors, inventory_items, purchase_orders, budget_caps).
2. compute_aggregate(values, operation): Deterministic arithmetic reductions (sum, mean, min, max, count, std).
3. verify_constraint(actual, expected, comparison, tolerance): Verify mathematical inequality bounds.

Anti-Reward Hacking Defense

To prevent agents from exploiting step-count efficiency penalties (e.g. failing early with zero calls), the reward engine enforces strict mathematical coupling:

$$R(\tau )=r_{\text{outcome}}(y^{\ast },y)+\mathbf{1}[y^{\ast }=y]\cdot r_{\text{efficiency}}(T,T_{\min })+\sum _{t=1}^T{r}_{\text{process}}(a_t)+r_{\text{safety}}R(\backslash tau)\; =\; r\_\{\backslash text\{outcome\}\}(y^*,\; y)\; +\; \backslash mathbf\{1\}[y^*\; =\; y]\; \backslash cdot\; r\_\{\backslash text\{efficiency\}\}(T,\; T\_\{\backslash min\})\; +\; \backslash sum\_\{t=1\}^T\; r\_\{\backslash text\{process\}\}(a\_t)\; +\; r\_\{\backslash text\{safety\}\}$$

The step penalty $r_{\text{efficiency}}$ is credited strictly conditional on verified task success ($\mathbf{1}[y^* = y]$). Premature exit receives zero reward ($R = 0.00$).

---

5. How to Load and Use

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

base_model_id = "CohereLabs/tiny-aya-global"
adapter_id = "jayesh20/lexsi-track3-aya-lora"

tokenizer = AutoTokenizer.from_pretrained(adapter_id)
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
model = PeftModel.from_pretrained(base_model, adapter_id)

prompt = "<|START_OF_TURN_TOKEN|><|USER_TOKEN|>Query purchase orders for vendor VND-106<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=128, temperature=0.01)
print(tokenizer.decode(outputs[0], skip_special_tokens=False))

---

6. Project & Submission Links

• GitHub Repository: https://github.com/jayeshpandey01/Lexsi
• Model Collection: Lexsi Track 3 Models
• Dataset Collection: Lexsi Track 3 Datasets & Traces
• Research Article: research_article.md

---

7. License & Compliance

• Base Model License: Cohere Labs Non-Commercial Research License (CC-BY-NC-4.0).
• Project Code & Harness: Apache-2.0 License.