Skip to main content

Gradient-guided layer selection for efficient LoRA fine-tuning

Project description

Aletheia-LoRA

Gradient-guided layer selection for efficient LoRA fine-tuning.

arXiv DOI Hugging Face License

Aletheia identifies the most informative transformer layers with a lightweight gradient probe, then applies LoRA adapters only to those layers. Across the finalized Research Line A evidence pack, this yields a 15-28% training speedup with bounded extra forgetting and broadly matched downstream behavior across 14 successful models from 8 architecture families spanning 0.5B to 72B parameters.

Paper: Aletheia: Gradient-Guided Layer Selection for Efficient LoRA Fine-Tuning Across Architectures

DOI: 10.48550/arXiv.2604.15351

Hugging Face card: aletheiaprotocol/aletheia-lora

Evidence dataset: aletheiaprotocol/aletheia-lora-evidence

Installation

Until the package is published to PyPI, install it from source:

git clone https://github.com/aletheiaprotocol-ai/aletheia-lora.git
cd aletheia-lora
python -m pip install -e ".[dev]"

Quick Start

import torch
from aletheia_lora import gradient_probe, select_layers, apply_aletheia_lora
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import get_peft_model

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-3B", torch_dtype="auto")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B")

# 1. Measure per-layer importance with a short gradient probe.
layer_scores = gradient_probe(
    model=model,
    dataset=train_dataset,       # tokenized HF Dataset with input_ids + attention_mask
    num_layers=36,
    probe_batches=5,
    chunk_size=8,
)

# 2. Select the top 50% most informative layers.
selected_layers = select_layers(layer_scores, top_pct=50)

# 3. Apply LoRA only to the selected layers.
peft_config = apply_aletheia_lora(
    model=model,
    selected_layers=selected_layers,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    r=16,
    lora_alpha=32,
)
model = get_peft_model(model, peft_config)

To use the paper-style asymmetric attention/MLP allocation:

peft_config = apply_aletheia_lora(
    model=model,
    selected_layers=selected_layers,
    attention_r=16,
    mlp_r=64,
)

Supported Claim

Use this wording:

Aletheia provides a reproducible training-efficiency improvement over standard LoRA across multiple model families and scales, with bounded extra forgetting on downstream benchmarks.

This package supports that finished narrow claim. It does not claim:

  • universal replacement for every LoRA setup
  • better quality on every model and every metric
  • compression better than quantization
  • proof of the full Aletheia or NEXUS paradigm

How It Works

  1. Gradient probe: computes per-layer gradient norms on a small probe batch set with bounded memory usage.
  2. Layer selection: ranks layers by gradient magnitude and keeps the top-k percent.
  3. Selective LoRA: applies standard PEFT LoraConfig only to the selected layers.

Evidence Summary

  • 81 experiment rows across 14 successful models and 8 architecture families
  • 100% per-model speed win rate in the main speed campaign
  • 23.1% mean speedup in the strongest frozen campaign
  • Bounded extra forgetting on MMLU, GSM8K, and HumanEval under the narrow-claim framing
  • Scale range from 0.5B to 72B parameters, including a MoE model

Repository Contents

Path Purpose
src/aletheia_lora/ Core library: probe, layer selection, and PEFT config builder
examples/quick_start.py CPU-safe usage demo with simulated layer scores
notebooks/quickstart_simulated_probe.ipynb Notebook version of the CPU-safe API demo
tests/ Unit tests for the public API
.github/workflows/ci.yml Public CI for supported Python versions
docs/ PyPI, Papers with Code, arXiv v2, and evidence-bundle notes
MODEL_CARD.md Human-readable method card
huggingface/README.md Hugging Face Hub-ready card
CITATION.cff GitHub citation metadata

Reproducibility

The full empirical record is described in the paper. A sanitized aggregate evidence companion is published at aletheiaprotocol/aletheia-lora-evidence. It contains aggregate metadata and final figures only; it excludes credentials, raw datasets, checkpoints, adapter weights, full gradient score dictionaries, and private cloud paths.

Citation

@misc{saket2026aletheia,
  title={Aletheia: Gradient-Guided Layer Selection for Efficient LoRA Fine-Tuning Across Architectures},
  author={Saket, Abdulmalek},
  year={2026},
  eprint={2604.15351},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  doi={10.48550/arXiv.2604.15351},
  url={https://arxiv.org/abs/2604.15351}
}

License

Apache 2.0

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

aletheia_lora-0.1.2.tar.gz (21.5 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

aletheia_lora-0.1.2-py3-none-any.whl (11.6 kB view details)

Uploaded Python 3

File details

Details for the file aletheia_lora-0.1.2.tar.gz.

File metadata

  • Download URL: aletheia_lora-0.1.2.tar.gz
  • Upload date:
  • Size: 21.5 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.12

File hashes

Hashes for aletheia_lora-0.1.2.tar.gz
Algorithm Hash digest
SHA256 86b85999d20efdc6b7d75422546dfcaca60fcee27d2f6cdd1b96e02b290bd601
MD5 cc6da200ec4d86bb2f978dd80f9a4181
BLAKE2b-256 87de49c64c1aae3786d797c245ffc52e27c90d2c1480d3b4939f603ee2c11caf

See more details on using hashes here.

Provenance

The following attestation bundles were made for aletheia_lora-0.1.2.tar.gz:

Publisher: publish-pypi.yml on aletheiaprotocol-ai/aletheia-lora

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file aletheia_lora-0.1.2-py3-none-any.whl.

File metadata

  • Download URL: aletheia_lora-0.1.2-py3-none-any.whl
  • Upload date:
  • Size: 11.6 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.12

File hashes

Hashes for aletheia_lora-0.1.2-py3-none-any.whl
Algorithm Hash digest
SHA256 a85c96a7f6bdc6ef8eb0dfd48c7b9c56b2a194d0aae94e4629e2c49ae9d4005b
MD5 e97e54864364ec59426ab985d9609384
BLAKE2b-256 89cfc91979c2ed8958c1d5d64f60e573ec7bffbd4d8de0aa27c4fbaf574c4a2f

See more details on using hashes here.

Provenance

The following attestation bundles were made for aletheia_lora-0.1.2-py3-none-any.whl:

Publisher: publish-pypi.yml on aletheiaprotocol-ai/aletheia-lora

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page