aurora-peft

Re-implementation của AuroRA: Breaking Low-Rank Bottleneck of LoRA with Nonlinear Mapping (NeurIPS 2025) dưới dạng một tuner cho 🤗 PEFT (>=0.13).

Cài đặt

pip install aurora-peft

Hoặc cài bản dev:

pip install -e .

Sử dụng nhanh

import torch
from transformers import AutoModel
from aurora_peft import AuroraConfig, get_aurora_model

base = AutoModel.from_pretrained("bert-base-uncased")

config = AuroraConfig(
    r=8,                       # LoRA rank (= ANL in/out dim)
    alpha=16,                  # LoRA alpha
    target_modules=["query", "value"],
    grid_size=5,               # B-spline intervals (paper: 5)
    spline_order=3,            # B-spline degree (paper: 3)
    scale_noise=0.1,
)

model = get_aurora_model(base, config)

# Optimizer chỉ update params của Aurora (ANL + A + B)
model.print_trainable_parameters()

Hoặc dùng qua API chuẩn của peft:

from peft import get_peft_model
from aurora_peft import AuroraConfig

model = get_peft_model(base, AuroraConfig(r=8, target_modules=["q_proj", "v_proj"]))

Công thức

AuroRA mở rộng LoRA bằng cách chèn một Adaptive Nonlinear Layer (ANL) giữa hai low-rank matrix A và B:

Δh = B · ANL(A · x)

trong đó ANL là một mini KAN:

ANL(z) = tanh(W_base @ tanh(z)) + W_spline @ Σ_i B_i(z)

B_i là basis B-spline bậc spline_order trên grid_size khoảng.

Re-implementation note

Đây là re-implementation không chính thức, dựa trên paper và 4 bản fork trong repo gốc https://github.com/ins1stenc3/AuroRA. Mục tiêu:

1. Rút logic lõi (ANL + cách gắn vào nn.Linear) ra khỏi 4 fork trùng lặp.
2. Tương thích với peft hiện hành (>=0.13).
3. Có unit test so khớp số học với bản gốc.

Bản gốc (paper) dùng hai công thức:

• Training (dynamic): h = W₀x + B · ANL(Ax) — input-dependent.
• Inference (static merge): ΔW = B · ANL(A) — weight-only, merge được vào W₀.

Repo gốc chỉ hiện thực 1 trong 2. Thư viện này:

• Forward (training) dùng công thức dynamic (khớp paper Algorithm 1, line "Training Phase").
• merge() dùng công thức static (line "Inference Preparation"), nên sau khi merge, forward chỉ là Wx (không còn gọi ANL nữa).

Tests

pip install -e ".[test]"
pytest

Giấy phép

Apache-2.0 (theo license của peft và paper).